diff options
Diffstat (limited to 'crates')
313 files changed, 56803 insertions, 0 deletions
diff --git a/crates/atuin-ai/Cargo.toml b/crates/atuin-ai/Cargo.toml new file mode 100644 index 00000000..1b81646c --- /dev/null +++ b/crates/atuin-ai/Cargo.toml @@ -0,0 +1,44 @@ +[package] +name = "atuin-ai" +edition = "2024" +description = "AI integration for Atuin CLI" + +rust-version = { workspace = true } +version = { workspace = true } +authors = { workspace = true } +license = { workspace = true } +homepage = { workspace = true } +repository = { workspace = true } + +# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html + +[dependencies] +atuin-client = { workspace = true } +atuin-common = { workspace = true } +tokio = { workspace = true } +eyre = { workspace = true } +clap = { workspace = true, features = ["derive", "env"] } +tracing = { workspace = true } +tracing-subscriber = { workspace = true, features = [ + "ansi", + "fmt", + "registry", + "env-filter", +] } +directories = { workspace = true } +tracing-appender = "0.2.4" +reqwest = { workspace = true } +serde = { workspace = true } +serde_json = { workspace = true } +crossterm = { workspace = true, features = ["use-dev-tty", "event-stream"] } +ratatui = { workspace = true, features = ["unstable-rendered-line-info"] } +futures = "0.3" +eventsource-stream = "0.2" +pulldown-cmark = "0.13.0" +async-stream = "0.3" +uuid = { workspace = true } +tui-textarea-2 = "0.9.1" +unicode-width = "0.2" + +[dev-dependencies] +pretty_assertions = { workspace = true } diff --git a/crates/atuin-ai/render-tests.sh b/crates/atuin-ai/render-tests.sh new file mode 100755 index 00000000..8dedc76e --- /dev/null +++ b/crates/atuin-ai/render-tests.sh @@ -0,0 +1,34 @@ +#!/bin/bash +# Render all test cases from test-renders.json +# Usage: ./render-tests.sh [test_name] +# With no args: renders all tests +# With arg: renders only matching test (e.g., ./render-tests.sh 05) + +set -e +cd "$(dirname "$0")" + +JSON_FILE="test-renders.json" +FILTER="${1:-}" + +# Build once +cargo build -p atuin-ai --quiet + +# Count tests +TOTAL=$(jq length "$JSON_FILE") + +for i in $(seq 0 $((TOTAL - 1))); do + NAME=$(jq -r ".[$i].name" "$JSON_FILE") + DESC=$(jq -r ".[$i].description" "$JSON_FILE") + STATE=$(jq -c ".[$i].state" "$JSON_FILE") + + # Skip if filter provided and doesn't match + if [[ -n "$FILTER" && ! "$NAME" =~ $FILTER ]]; then + continue + fi + + echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━" + echo "[$NAME] $DESC" + echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━" + echo "$STATE" | cargo run -p atuin-ai --quiet -- debug-render -f plain + echo "" +done diff --git a/crates/atuin-ai/replay-states.sh b/crates/atuin-ai/replay-states.sh new file mode 100755 index 00000000..4f586709 --- /dev/null +++ b/crates/atuin-ai/replay-states.sh @@ -0,0 +1,101 @@ +#!/bin/bash +# Replay state snapshots from a debug state JSONL file +# Usage: ./replay-states.sh <state-file.jsonl> [entry-number] +# With no entry: renders all frames in sequence (press Enter to advance) +# With entry number: renders just that frame + +set -e +# cd "$(dirname "$0")" + +STATE_FILE="${1:-}" +ENTRY_FILTER="${2:-}" + +if [[ -z "$STATE_FILE" ]]; then + echo "Usage: $0 <state-file.jsonl> [entry-number]" + echo "" + echo "Examples:" + echo " $0 /tmp/state.jsonl # Interactive replay of all frames" + echo " $0 /tmp/state.jsonl 15 # Show just entry 15" + exit 1 +fi + +if [[ ! -f "$STATE_FILE" ]]; then + echo "Error: File not found: $STATE_FILE" + exit 1 +fi + +# Build once +cargo build -p atuin-ai --quiet + +# Count entries +TOTAL=$(wc -l < "$STATE_FILE" | tr -d ' ') + +if [[ -n "$ENTRY_FILTER" ]]; then + # Show single entry + LINE=$(sed -n "${ENTRY_FILTER}p" "$STATE_FILE") + if [[ -z "$LINE" ]]; then + echo "Error: Entry $ENTRY_FILTER not found (file has $TOTAL entries)" + exit 1 + fi + + ENTRY=$(echo "$LINE" | jq -r '.entry') + LABEL=$(echo "$LINE" | jq -r '.label') + STATE=$(echo "$LINE" | jq -c '.state') + + echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━" + echo "[$ENTRY/$TOTAL] $LABEL" + echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━" + echo "$STATE" | cargo run -p atuin-ai --quiet -- debug-render -f plain +else + # Interactive replay + echo "Replaying $TOTAL frames from $STATE_FILE" + echo "Press Enter to advance, 'q' to quit, or number+Enter to jump" + echo "" + + CURRENT=1 + while [[ $CURRENT -le $TOTAL ]]; do + LINE=$(sed -n "${CURRENT}p" "$STATE_FILE") + ENTRY=$(echo "$LINE" | jq -r '.entry') + LABEL=$(echo "$LINE" | jq -r '.label') + STATE=$(echo "$LINE" | jq -c '.state') + + clear + echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━" + echo "[$CURRENT/$TOTAL] $LABEL" + echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━" + echo "$STATE" | cargo run -p atuin-ai --quiet -- debug-render -f plain + echo "" + echo "[Enter: next] [p: prev] [number: jump] [s: show state JSON] [q: quit]" + + read -r INPUT + case "$INPUT" in + q|Q) + break + ;; + p|P) + if [[ $CURRENT -gt 1 ]]; then + CURRENT=$((CURRENT - 1)) + fi + ;; + s|S) + echo "" + echo "State JSON:" + echo "$STATE" | jq . + echo "" + echo "Press Enter to continue..." + read -r + ;; + ''|' ') + CURRENT=$((CURRENT + 1)) + ;; + *[0-9]*) + if [[ "$INPUT" =~ ^[0-9]+$ ]] && [[ "$INPUT" -ge 1 ]] && [[ "$INPUT" -le $TOTAL ]]; then + CURRENT=$INPUT + else + echo "Invalid entry number (1-$TOTAL)" + sleep 1 + fi + ;; + esac + done +fi diff --git a/crates/atuin-ai/src/commands.rs b/crates/atuin-ai/src/commands.rs new file mode 100644 index 00000000..d04875ea --- /dev/null +++ b/crates/atuin-ai/src/commands.rs @@ -0,0 +1,201 @@ +use std::{ + fs, + path::{Path, PathBuf}, +}; + +use atuin_common::shell::Shell; +use clap::{Args, Subcommand}; +use eyre::Result; +use tracing_appender::rolling::{RollingFileAppender, Rotation}; +use tracing_subscriber::{EnvFilter, Layer, fmt, layer::SubscriberExt, util::SubscriberInitExt}; +#[cfg(debug_assertions)] +pub mod debug_render; + +pub mod init; +pub mod inline; + +#[derive(Args, Debug)] +pub struct AiArgs { + /// Enable verbose logging + #[arg(short, long, global = true)] + verbose: bool, + + /// Custom API endpoint; defaults to reading from the `ai.endpoint` setting. + #[arg(long, global = true)] + api_endpoint: Option<String>, + + /// Custom API token; defaults to reading from the `ai.api_token` setting. + #[arg(long, global = true)] + api_token: Option<String>, +} + +#[derive(Subcommand, Debug)] +pub enum Commands { + /// Initialize shell integration + Init { + /// Shell to generate integration for; defaults to "auto" + #[arg(value_name = "SHELL", default_value = "auto")] + shell: String, + }, + + /// Inline completion mode with small TUI overlay + Inline { + #[command(flatten)] + args: AiArgs, + + /// Current command line to complete + #[arg(value_name = "COMMAND")] + command: Option<String>, + + /// Keep TUI output visible after exit (default: erase) + #[arg(long)] + keep: bool, + + /// Use the hook mode + #[arg(long, hide = true)] + hook: bool, + + /// Log state changes to file for debugging (dev tool) + #[arg(long, value_name = "FILE", hide = true)] + debug_state: Option<String>, + }, + + /// Debug render: output a single frame from JSON state (dev tool) + #[cfg(debug_assertions)] + DebugRender { + /// Input file (reads from stdin if not provided) + #[arg(short, long)] + input: Option<String>, + + /// Output format: ansi (default), plain, json + #[arg(short, long, default_value = "ansi")] + format: String, + }, +} + +pub async fn run( + command: Commands, + settings: &atuin_client::settings::Settings, +) -> eyre::Result<()> { + match command { + Commands::Init { shell } => init::run(shell).await, + Commands::Inline { + command, + keep, + debug_state, + hook, + args, + .. + } => { + if settings.logs.ai_enabled() { + init_logging(settings, args.verbose)?; + } + + inline::run( + command, + args.api_endpoint, + args.api_token, + keep, + debug_state, + settings, + hook, + ) + .await + } + #[cfg(debug_assertions)] + Commands::DebugRender { input, format } => { + let output_format = match format.as_str() { + "plain" => debug_render::OutputFormat::Plain, + "json" => debug_render::OutputFormat::Json, + _ => debug_render::OutputFormat::Ansi, + }; + debug_render::run(input, output_format).await + } + } +} + +pub fn detect_shell() -> Option<String> { + Some(Shell::current().to_string()) +} + +/// Initializes logging for the AI commands. +fn init_logging(settings: &atuin_client::settings::Settings, verbose: bool) -> Result<()> { + // ATUIN_LOG env var overrides config file level settings + let env_log_set = std::env::var("ATUIN_LOG").is_ok(); + + // Base filter from env var (or empty if not set) + let base_filter = + EnvFilter::from_env("ATUIN_LOG").add_directive("sqlx_sqlite::regexp=off".parse()?); + + // Use config level unless ATUIN_LOG is set + let filter = if env_log_set { + base_filter + } else { + EnvFilter::default() + .add_directive(settings.logs.ai_level().as_directive().parse()?) + .add_directive("sqlx_sqlite::regexp=off".parse()?) + }; + + let log_dir = PathBuf::from(&settings.logs.dir); + let ai_log_filename = settings.logs.ai.file.clone(); + + // Clean up old log files + cleanup_old_logs(&log_dir, &ai_log_filename, settings.logs.ai_retention()); + + let console_layer = if verbose { + Some( + fmt::layer() + .with_writer(std::io::stderr) + .with_ansi(true) + .with_target(false) + .with_filter(filter.clone()), + ) + } else { + None + }; + + let file_appender = RollingFileAppender::new(Rotation::DAILY, &log_dir, &ai_log_filename); + + let base = tracing_subscriber::registry().with( + fmt::layer() + .with_writer(file_appender) + .with_ansi(false) + .with_filter(filter), + ); + + if let Some(console_layer) = console_layer { + base.with(console_layer).init(); + } else { + base.init(); + }; + + Ok(()) +} + +fn cleanup_old_logs(log_dir: &Path, prefix: &str, retention_days: u64) { + let cutoff = std::time::SystemTime::now() + - std::time::Duration::from_secs(retention_days * 24 * 60 * 60); + + let Ok(entries) = fs::read_dir(log_dir) else { + return; + }; + + for entry in entries.flatten() { + let path = entry.path(); + let Some(name) = path.file_name().and_then(|n| n.to_str()) else { + continue; + }; + + // Match files like "search.log.2024-02-23" or "daemon.log.2024-02-23" + if !name.starts_with(prefix) || name == prefix { + continue; + } + + if let Ok(metadata) = entry.metadata() + && let Ok(modified) = metadata.modified() + && modified < cutoff + { + let _ = fs::remove_file(&path); + } + } +} diff --git a/crates/atuin-ai/src/commands/debug_render.rs b/crates/atuin-ai/src/commands/debug_render.rs new file mode 100644 index 00000000..b35d73c9 --- /dev/null +++ b/crates/atuin-ai/src/commands/debug_render.rs @@ -0,0 +1,466 @@ +//! Debug render command for TUI development +//! +//! Takes JSON state as input and outputs a single rendered frame as text. +//! Useful for debugging view model derivation and rendering without running the full TUI. + +use eyre::{Context, Result}; +use ratatui::{Terminal, backend::TestBackend}; +use serde::Deserialize; +use std::io::{self, Read}; +use std::time::Instant; + +use crate::tui::{ + render::{RenderContext, render}, + state::{AppMode, AppState, ConversationEvent, StreamingStatus}, + view_model::Blocks, +}; + +/// JSON input format for debug rendering +#[derive(Debug, Deserialize)] +pub struct DebugInput { + /// Conversation events in API format + pub events: Vec<EventInput>, + /// Current mode: "Input", "Generating", "Streaming", "Review", "Error" + #[serde(default = "default_mode")] + pub mode: String, + /// Text being streamed (for Streaming mode) + #[serde(default)] + pub streaming_text: String, + /// Current input buffer + #[serde(default)] + pub input: String, + /// Cursor position + #[serde(default)] + pub cursor_pos: usize, + /// Spinner frame (0-3) + #[serde(default)] + pub spinner_frame: usize, + /// Error message + #[serde(default)] + pub error: Option<String>, + /// Session ID from server + #[serde(default)] + pub session_id: Option<String>, + /// Streaming status + #[serde(default)] + pub streaming_status: Option<String>, + /// Whether current turn was interrupted + #[serde(default)] + pub was_interrupted: bool, + /// Terminal width for rendering + #[serde(default = "default_width")] + pub width: u16, + /// Terminal height for rendering + #[serde(default = "default_height")] + pub height: u16, +} + +fn default_mode() -> String { + "Review".to_string() +} + +fn default_width() -> u16 { + 80 +} + +fn default_height() -> u16 { + // Default to a reasonable height; state files include calculated height + 50 +} + +/// Event input matching the API protocol format +#[derive(Debug, Clone, Deserialize)] +#[serde(tag = "type", rename_all = "snake_case")] +pub enum EventInput { + UserMessage { + content: String, + }, + Text { + content: String, + }, + ToolCall { + id: String, + name: String, + input: serde_json::Value, + }, + ToolResult { + tool_use_id: String, + content: String, + #[serde(default)] + is_error: bool, + }, +} + +impl From<EventInput> for ConversationEvent { + fn from(input: EventInput) -> Self { + match input { + EventInput::UserMessage { content } => ConversationEvent::UserMessage { content }, + EventInput::Text { content } => ConversationEvent::Text { content }, + EventInput::ToolCall { id, name, input } => { + ConversationEvent::ToolCall { id, name, input } + } + EventInput::ToolResult { + tool_use_id, + content, + is_error, + } => ConversationEvent::ToolResult { + tool_use_id, + content, + is_error, + }, + } + } +} + +impl DebugInput { + /// Parse JSON from string + pub fn from_json(json: &str) -> Result<Self> { + serde_json::from_str(json).context("Failed to parse debug input JSON") + } + + /// Convert to AppState + pub fn to_state(&self) -> AppState { + let mode = match self.mode.as_str() { + "Input" => AppMode::Input, + "Generating" => AppMode::Generating, + "Streaming" => AppMode::Streaming, + "Review" => AppMode::Review, + "Error" => AppMode::Error, + _ => AppMode::Review, + }; + + let events: Vec<ConversationEvent> = self.events.iter().cloned().map(Into::into).collect(); + + let streaming_status = self + .streaming_status + .as_ref() + .map(|s| StreamingStatus::from_status_str(s)); + + // Create textarea from input and set cursor position + let mut textarea = tui_textarea::TextArea::from(self.input.lines()); + // Disable underline on cursor line + textarea.set_cursor_line_style(ratatui::style::Style::default()); + // Enable word wrapping + textarea.set_wrap_mode(tui_textarea::WrapMode::Word); + // Note: cursor_pos from old format is character-based; new format has row/col + // For compatibility, just move to end if we have text + if !self.input.is_empty() { + textarea.move_cursor(tui_textarea::CursorMove::End); + } + + AppState { + mode, + events, + streaming_text: self.streaming_text.clone(), + textarea, + error: self.error.clone(), + should_exit: false, + exit_action: None, + session_id: self.session_id.clone(), + streaming_status, + was_interrupted: self.was_interrupted, + spinner_frame: self.spinner_frame, + last_spinner_tick: Instant::now(), + streaming_started: None, + confirmation_pending: false, + } + } +} + +/// Output format options +#[derive(Debug, Clone, Copy, Default)] +pub enum OutputFormat { + /// Raw terminal output (ANSI) + #[default] + Ansi, + /// Plain text (strips ANSI codes) + Plain, + /// JSON with blocks structure + Json, +} + +/// Run the debug render command +pub async fn run(input_file: Option<String>, format: OutputFormat) -> Result<()> { + // Read input JSON + let json = if let Some(path) = input_file { + std::fs::read_to_string(&path).context(format!("Failed to read input file: {}", path))? + } else { + let mut buffer = String::new(); + io::stdin() + .read_to_string(&mut buffer) + .context("Failed to read from stdin")?; + buffer + }; + + let debug_input = DebugInput::from_json(&json)?; + let state = debug_input.to_state(); + + match format { + OutputFormat::Json => { + // Output the derived blocks as JSON + let blocks = Blocks::from_state(&state); + println!( + "{}", + serde_json::to_string_pretty(&blocks_to_json(&blocks))? + ); + } + OutputFormat::Plain | OutputFormat::Ansi => { + // Render to a test backend + let backend = TestBackend::new(debug_input.width, debug_input.height); + let mut terminal = Terminal::new(backend)?; + + // Load default theme + let settings = atuin_client::settings::Settings::new()?; + let mut theme_manager = atuin_client::theme::ThemeManager::new(None, None); + let theme = theme_manager.load_theme(&settings.theme.name, None); + + let ctx = RenderContext { + theme, + anchor_col: 0, + textarea: Some(&state.textarea), + max_height: debug_input.height, + popup_mode: false, + render_above: false, + }; + + terminal.draw(|frame| { + render(frame, &state, &ctx); + })?; + + // Get buffer content + let buffer = terminal.backend().buffer(); + let output = buffer_to_string(buffer, matches!(format, OutputFormat::Plain)); + print!("{}", output); + } + } + + Ok(()) +} + +/// Convert blocks to JSON for debugging +fn blocks_to_json(blocks: &Blocks) -> serde_json::Value { + serde_json::json!({ + "count": blocks.items.len(), + "blocks": blocks.items.iter().map(|block| { + serde_json::json!({ + "separator_above": block.separator_above, + "title": block.title, + "content": block.content.iter().map(content_to_json).collect::<Vec<_>>() + }) + }).collect::<Vec<_>>(), + "status_bar": blocks.status_bar.as_ref().map(|sb| serde_json::json!({ + "frame": sb.frame, + "text": sb.text + })) + }) +} + +fn content_to_json(content: &crate::tui::view_model::Content) -> serde_json::Value { + use crate::tui::view_model::Content; + match content { + Content::Input { + text, + active, + cursor_pos, + } => serde_json::json!({ + "type": "Input", + "text": text, + "active": active, + "cursor_pos": cursor_pos + }), + Content::Command { text, faded } => serde_json::json!({ + "type": "Command", + "text": text, + "faded": faded + }), + Content::Text { markdown } => serde_json::json!({ + "type": "Text", + "markdown": markdown + }), + Content::Error { message } => serde_json::json!({ + "type": "Error", + "message": message + }), + Content::Warning { + kind, + text, + pending_confirm, + } => serde_json::json!({ + "type": "Warning", + "kind": format!("{:?}", kind), + "text": text, + "pending_confirm": pending_confirm + }), + Content::Spinner { frame, status_text } => serde_json::json!({ + "type": "Spinner", + "frame": frame, + "status_text": status_text + }), + Content::ToolStatus { + completed_count, + current_label, + frame, + } => serde_json::json!({ + "type": "ToolStatus", + "completed_count": completed_count, + "current_label": current_label, + "frame": frame + }), + } +} + +/// Convert ratatui buffer to string +fn buffer_to_string(buffer: &ratatui::buffer::Buffer, strip_ansi: bool) -> String { + let area = buffer.area; + let mut output = String::new(); + + for y in 0..area.height { + for x in 0..area.width { + let cell = &buffer[(x, y)]; + if strip_ansi { + output.push_str(cell.symbol()); + } else { + // Include ANSI styling + let fg = cell.fg; + let bg = cell.bg; + let mods = cell.modifier; + + // Simple ANSI encoding + if fg != ratatui::style::Color::Reset + || bg != ratatui::style::Color::Reset + || !mods.is_empty() + { + output.push_str("\x1b["); + let mut first = true; + + if mods.contains(ratatui::style::Modifier::BOLD) { + output.push('1'); + first = false; + } + if mods.contains(ratatui::style::Modifier::DIM) { + if !first { + output.push(';'); + } + output.push('2'); + first = false; + } + if mods.contains(ratatui::style::Modifier::REVERSED) { + if !first { + output.push(';'); + } + output.push('7'); + first = false; + } + if mods.contains(ratatui::style::Modifier::UNDERLINED) { + if !first { + output.push(';'); + } + output.push('4'); + first = false; + } + + if let Some(code) = color_to_ansi(fg, true) { + if !first { + output.push(';'); + } + output.push_str(&code); + first = false; + } + + if let Some(code) = color_to_ansi(bg, false) { + if !first { + output.push(';'); + } + output.push_str(&code); + } + + output.push('m'); + } + + output.push_str(cell.symbol()); + + if fg != ratatui::style::Color::Reset + || bg != ratatui::style::Color::Reset + || !mods.is_empty() + { + output.push_str("\x1b[0m"); + } + } + } + output.push('\n'); + } + + output +} + +fn color_to_ansi(color: ratatui::style::Color, foreground: bool) -> Option<String> { + use ratatui::style::Color; + let base = if foreground { 30 } else { 40 }; + + match color { + Color::Reset => None, + Color::Black => Some((base).to_string()), + Color::Red => Some((base + 1).to_string()), + Color::Green => Some((base + 2).to_string()), + Color::Yellow => Some((base + 3).to_string()), + Color::Blue => Some((base + 4).to_string()), + Color::Magenta => Some((base + 5).to_string()), + Color::Cyan => Some((base + 6).to_string()), + Color::Gray | Color::White => Some((base + 7).to_string()), + Color::DarkGray => Some((base + 60).to_string()), + Color::LightRed => Some((base + 61).to_string()), + Color::LightGreen => Some((base + 62).to_string()), + Color::LightYellow => Some((base + 63).to_string()), + Color::LightBlue => Some((base + 64).to_string()), + Color::LightMagenta => Some((base + 65).to_string()), + Color::LightCyan => Some((base + 66).to_string()), + Color::Indexed(i) => Some(format!("{}8;5;{}", if foreground { 3 } else { 4 }, i)), + Color::Rgb(r, g, b) => Some(format!( + "{}8;2;{};{};{}", + if foreground { 3 } else { 4 }, + r, + g, + b + )), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_parse_simple_input() { + let json = r#"{ + "events": [ + {"type": "user_message", "content": "list files"}, + {"type": "tool_call", "id": "123", "name": "suggest_command", "input": {"command": "ls -la"}} + ], + "mode": "Review" + }"#; + + let input = DebugInput::from_json(json).unwrap(); + assert_eq!(input.events.len(), 2); + assert_eq!(input.mode, "Review"); + + let state = input.to_state(); + assert_eq!(state.events.len(), 2); + assert_eq!(state.mode, AppMode::Review); + } + + #[test] + fn test_parse_streaming_state() { + let json = r#"{ + "events": [ + {"type": "user_message", "content": "explain flags"} + ], + "mode": "Streaming", + "streaming_text": "The -l flag means..." + }"#; + + let input = DebugInput::from_json(json).unwrap(); + let state = input.to_state(); + assert_eq!(state.mode, AppMode::Streaming); + assert_eq!(state.streaming_text, "The -l flag means..."); + } +} diff --git a/crates/atuin-ai/src/commands/init.rs b/crates/atuin-ai/src/commands/init.rs new file mode 100644 index 00000000..6b23e936 --- /dev/null +++ b/crates/atuin-ai/src/commands/init.rs @@ -0,0 +1,230 @@ +use crate::commands::detect_shell; + +pub async fn run(shell: String) -> eyre::Result<()> { + let integration = match shell.as_str() { + "zsh" => generate_zsh_integration(), + "bash" => generate_bash_integration(), + "fish" => generate_fish_integration(), + "auto" => generate_auto_integration()?, + _ => eyre::bail!("Unsupported shell: {}", shell), + }; + + println!("{}", integration); + Ok(()) +} + +fn generate_auto_integration() -> eyre::Result<&'static str> { + let shell = detect_shell(); + match shell.as_deref() { + Some("zsh") => Ok(generate_zsh_integration()), + Some("bash") => Ok(generate_bash_integration()), + Some("fish") => Ok(generate_fish_integration()), + Some(s) => eyre::bail!("Unsupported shell: {}", s), + None => eyre::bail!("Could not detect shell"), + } +} + +/// Generate the zsh integration function - pure function for easy testing +pub fn generate_zsh_integration() -> &'static str { + r#" +# TUI uses an alternate screen, so no explicit cleanup is needed. +_atuin_ai_cleanup() { + true +} + +# Question mark at start of line - natural language mode +_atuin_ai_question_mark() { + # If buffer is empty or just contains '?', trigger natural language mode + if [[ -z "$BUFFER" || "$BUFFER" == "?" ]]; then + BUFFER="" + local output + output=$(atuin ai inline --hook 3>&1 1>&2 2>&3) + + # Clean up the inline viewport + _atuin_ai_cleanup + + if [[ $output == __atuin_ai_print__:* ]]; then + zle -I + echo "${output#__atuin_ai_print__:}" + elif [[ $output == __atuin_ai_cancel__ ]]; then + zle reset-prompt + elif [[ $output == __atuin_ai_execute__:* ]]; then + RBUFFER="" + LBUFFER=${output#__atuin_ai_execute__:} + zle reset-prompt + zle accept-line + elif [[ $output == __atuin_ai_insert__:* ]]; then + RBUFFER="" + LBUFFER=${output#__atuin_ai_insert__:} + zle reset-prompt + elif [[ -n $output ]]; then + RBUFFER="" + LBUFFER=$output + zle reset-prompt + else + zle reset-prompt + fi + else + LBUFFER="${LBUFFER}?" + fi +} + +# Set up keybindings +zle -N _atuin_ai_question_mark +bindkey '?' _atuin_ai_question_mark # Question mark +"# + .trim() +} + +/// Generate the bash integration function - pure function for easy testing +pub fn generate_bash_integration() -> &'static str { + r#" +# Question mark at start of line - natural language mode +_atuin_ai_question_mark() { + # If buffer is empty or just contains '?', trigger natural language mode + if [[ -z "$READLINE_LINE" || "$READLINE_LINE" == "?" ]]; then + READLINE_LINE="" + READLINE_POINT=0 + + local output + output=$(atuin ai inline --hook 3>&1 1>&2 2>&3) + + if [[ $output == __atuin_ai_print__:* ]]; then + echo "${output#__atuin_ai_print__:}" + READLINE_LINE="" + READLINE_POINT=0 + elif [[ $output == __atuin_ai_cancel__ ]]; then + READLINE_LINE="" + READLINE_POINT=0 + elif [[ $output == __atuin_ai_execute__:* ]]; then + # Execute the command immediately + READLINE_LINE=${output#__atuin_ai_execute__:} + READLINE_POINT=${#READLINE_LINE} + # Note: We can't directly execute in bash bind -x, but we can + # use a workaround by binding to a macro that accepts the line + bind '"\C-x\C-a": accept-line' + bind -x '"\C-x\C-e": _atuin_ai_question_mark' + elif [[ $output == __atuin_ai_insert__:* ]]; then + # Insert the command for editing + READLINE_LINE=${output#__atuin_ai_insert__:} + READLINE_POINT=${#READLINE_LINE} + elif [[ -n $output ]]; then + # Default: insert for editing + READLINE_LINE=$output + READLINE_POINT=${#READLINE_LINE} + fi + else + # Not at empty prompt, just insert the question mark + READLINE_LINE="${READLINE_LINE:0:READLINE_POINT}?${READLINE_LINE:READLINE_POINT}" + ((READLINE_POINT++)) + fi +} + +# Set up keybindings +# Bash requires special handling: we use bind -x for the function, +# but need a two-step approach for execute mode +__atuin_ai_accept_line="" + +_atuin_ai_question_mark_wrapper() { + _atuin_ai_question_mark + if [[ -n "$__atuin_ai_accept_line" ]]; then + __atuin_ai_accept_line="" + fi +} + +bind -x '"?": _atuin_ai_question_mark' +"# + .trim() +} + +/// Generate the fish integration function - pure function for easy testing +pub fn generate_fish_integration() -> &'static str { + r#" +# Question mark at start of line - natural language mode +function _atuin_ai_question_mark + set -l buf (commandline -b) + + # If buffer is empty or just contains '?', trigger natural language mode + if test -z "$buf" -o "$buf" = "?" + commandline -r "" + + # Run atuin ai inline, swapping stdout and stderr + set -l output (atuin ai inline --hook 3>&1 1>&2 2>&3 | string collect) + + if string match --quiet '__atuin_ai_print__:*' "$output" + echo (string replace "__atuin_ai_print__:" "" -- "$output" | string collect) + commandline -f repaint + else if test "$output" = "__atuin_ai_cancel__" + commandline -f repaint + else if string match --quiet '__atuin_ai_execute__:*' "$output" + # Execute the command immediately + set -l cmd (string replace "__atuin_ai_execute__:" "" -- "$output" | string collect) + commandline -r "$cmd" + commandline -f repaint + commandline -f execute + else if string match --quiet '__atuin_ai_insert__:*' "$output" + # Insert the command for editing + set -l cmd (string replace "__atuin_ai_insert__:" "" -- "$output" | string collect) + commandline -r "$cmd" + commandline -f repaint + else if test -n "$output" + # Default: insert for editing + commandline -r "$output" + commandline -f repaint + else + commandline -f repaint + end + else + # Not at empty prompt, just insert the question mark + commandline -i "?" + end +end + +# Set up keybindings +bind "?" _atuin_ai_question_mark +"# + .trim() +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_generate_zsh_integration() { + let result = generate_zsh_integration(); + assert!(result.contains("_atuin_ai_question_mark")); + assert!(result.contains("bindkey")); + assert!(result.contains("atuin ai inline --hook")); + assert!(result.contains("__atuin_ai_print__")); + assert!(result.contains("__atuin_ai_cancel__")); + assert!(result.contains("__atuin_ai_execute__")); + assert!(result.contains("__atuin_ai_insert__")); + } + + #[test] + fn test_generate_bash_integration() { + let result = generate_bash_integration(); + assert!(result.contains("_atuin_ai_question_mark")); + assert!(result.contains("bind")); + assert!(result.contains("READLINE_LINE")); + assert!(result.contains("atuin ai inline --hook")); + assert!(result.contains("__atuin_ai_print__")); + assert!(result.contains("__atuin_ai_cancel__")); + assert!(result.contains("__atuin_ai_execute__")); + assert!(result.contains("__atuin_ai_insert__")); + } + + #[test] + fn test_generate_fish_integration() { + let result = generate_fish_integration(); + assert!(result.contains("_atuin_ai_question_mark")); + assert!(result.contains("bind")); + assert!(result.contains("commandline")); + assert!(result.contains("atuin ai inline --hook")); + assert!(result.contains("__atuin_ai_print__")); + assert!(result.contains("__atuin_ai_cancel__")); + assert!(result.contains("__atuin_ai_execute__")); + assert!(result.contains("__atuin_ai_insert__")); + } +} diff --git a/crates/atuin-ai/src/commands/inline.rs b/crates/atuin-ai/src/commands/inline.rs new file mode 100644 index 00000000..ce566be1 --- /dev/null +++ b/crates/atuin-ai/src/commands/inline.rs @@ -0,0 +1,749 @@ +use crate::commands::detect_shell; +use crate::tui::render::render; +use crate::tui::{ + App, AppEvent, AppMode, ConversationEvent, EventLoop, ExitAction, RenderContext, TerminalGuard, + calculate_needed_height, install_panic_hook, +}; +use atuin_client::theme::ThemeManager; +use atuin_common::tls::ensure_crypto_provider; +use crossterm::{ + event::{self, Event, KeyCode}, + terminal::{disable_raw_mode, enable_raw_mode}, +}; +use eventsource_stream::Eventsource; +use eyre::{Context as _, Result, bail}; +use futures::StreamExt; +use reqwest::Url; +use std::io::Write; +use tracing::{debug, error, info, trace}; + +pub async fn run( + initial_command: Option<String>, + api_endpoint: Option<String>, + api_token: Option<String>, + keep_output: bool, + debug_state_file: Option<String>, + settings: &atuin_client::settings::Settings, + output_for_hook: bool, +) -> Result<()> { + if !settings.ai.enabled { + emit_shell_result( + Action::Print( + "Atuin AI is not enabled. Please enable it in your settings or run `atuin setup`." + .to_string(), + ), + output_for_hook, + ); + return Ok(()); + } + + // Install panic hook once at entry point to ensure terminal restoration + install_panic_hook(); + + // Token and endpoint priority: + // 1. Command line arguments/environment variables + // 2. Settings file + // 3. Default + let endpoint = api_endpoint.as_deref().unwrap_or( + settings + .ai + .endpoint + .as_deref() + .unwrap_or("https://hub.atuin.sh"), + ); + let api_token = api_token.as_deref().or(settings.ai.api_token.as_deref()); + + let token = if let Some(token) = &api_token { + token.to_string() + } else { + // If no token is provided, assume we're using Hub as the endpoint if we're using Hub sync + if settings.is_hub_sync() { + ensure_hub_session(settings).await? + } else { + bail!("No API token provided in ai.api_token settings or command line argument.") + } + }; + + let action = run_inline_tui( + endpoint.to_string(), + token, + initial_command, + keep_output, + debug_state_file, + settings, + ) + .await?; + emit_shell_result(action, output_for_hook); + + Ok(()) +} + +async fn ensure_hub_session(settings: &atuin_client::settings::Settings) -> Result<String> { + if let Some(token) = atuin_client::hub::get_session_token().await? { + debug!("Found Hub session, using existing token"); + return Ok(token); + } + + let hub_address = settings + .active_hub_endpoint() + .unwrap_or("https://hub.atuin.sh".to_string()); + + let will_sync = settings.is_hub_sync(); + + info!("No Hub session found, prompting for authentication"); + + println!("Atuin AI requires authenticating with Atuin Hub."); + if will_sync { + println!( + "Once logged in, your shell history will be synchronized via Atuin Hub if auto_sync is enabled or when manually syncing." + ) + } + println!( + "If you have an existing Atuin sync account, you can log in with your existing credentials." + ); + println!("Press enter to begin (or esc to cancel)."); + if !wait_for_login_confirmation()? { + bail!("authentication canceled"); + } + + debug!("Starting Atuin Hub authentication..."); + + println!("Authenticating with Atuin Hub..."); + let session = atuin_client::hub::HubAuthSession::start(&hub_address).await?; + println!("Open this URL to continue:"); + println!("{}", session.auth_url); + + let token = session + .wait_for_completion( + atuin_client::hub::DEFAULT_AUTH_TIMEOUT, + atuin_client::hub::DEFAULT_POLL_INTERVAL, + ) + .await?; + + info!("Authentication complete, saving session token"); + + atuin_client::hub::save_session(&token).await?; + + // Silently attempt to link CLI account to Hub if one exists + // This enables unified auth - users can use their Hub token for sync + if let Ok(meta) = atuin_client::settings::Settings::meta_store().await + && let Ok(Some(cli_token)) = meta.session_token().await + { + debug!("CLI session found, attempting to link accounts"); + if let Err(e) = atuin_client::hub::link_account(&hub_address, &cli_token).await { + // Don't fail AI flow if linking fails - it's not critical + debug!("Could not link CLI account to Hub: {}", e); + } else { + info!("Successfully linked CLI account to Hub"); + } + } + + Ok(token) +} + +/// SSE event received from chat endpoint +#[derive(Debug, Clone)] +enum ChatStreamEvent { + /// Text chunk to display + TextChunk(String), + /// Tool call event (need to echo back, may contain suggest_command) + ToolCall { + id: String, + name: String, + input: serde_json::Value, + }, + /// Tool result from server-side execution + ToolResult { + tool_use_id: String, + content: String, + is_error: bool, + }, + /// Status update from server + Status(String), + /// Stream complete + Done { session_id: String }, + /// Error from server + Error(String), +} + +fn create_chat_stream( + hub_address: String, + token: String, + session_id: Option<String>, + messages: Vec<serde_json::Value>, + settings: &atuin_client::settings::Settings, +) -> std::pin::Pin<Box<dyn futures::Stream<Item = Result<ChatStreamEvent>> + Send>> { + let send_cwd = settings.ai.send_cwd; + + Box::pin(async_stream::stream! { + ensure_crypto_provider(); + let endpoint = match hub_url(&hub_address, "/api/cli/chat") { + Ok(url) => url, + Err(e) => { + yield Err(e); + return; + } + }; + + debug!("Sending SSE request to {endpoint}"); + + // Build request body + let mut request_body = serde_json::json!({ + "messages": messages, + "context": { + "os": detect_os(), + "shell": detect_shell(), + "pwd": if send_cwd { std::env::current_dir() + .ok() + .map(|path| path.to_string_lossy().into_owned()) } else { None }, + } + }); + + // Include session_id only if present (not on first request) + if let Some(ref sid) = session_id { + trace!("Including session_id in request: {sid}"); + request_body["session_id"] = serde_json::json!(sid); + } + + + let client = reqwest::Client::new(); + let response = match client + .post(endpoint.clone()) + .header("Accept", "text/event-stream") + .bearer_auth(&token) + .json(&request_body) + .send() + .await + { + Ok(resp) => resp, + Err(e) => { + yield Err(eyre::eyre!("Failed to send SSE request: {}", e)); + return; + } + }; + + let status = response.status(); + if status == reqwest::StatusCode::UNAUTHORIZED { + // Clear saved session on auth error + error!("SSE request failed with status: {status}, clearing session"); + let _ = atuin_client::hub::delete_session().await; + yield Err(eyre::eyre!("Hub session expired. Re-run to authenticate again.")); + return; + } + if !status.is_success() { + let body = response.text().await.unwrap_or_default(); + error!("SSE request failed ({}): {}", status, body); + yield Err(eyre::eyre!("SSE request failed ({}): {}", status, body)); + return; + } + + let byte_stream = response.bytes_stream(); + let mut stream = byte_stream.eventsource(); + + while let Some(event) = stream.next().await { + match event { + Ok(sse_event) => { + let event_type = sse_event.event.as_str(); + let data = sse_event.data.clone(); + + debug!(event_type = %event_type, "SSE event received"); + + match event_type { + "text" => { + if let Ok(json) = serde_json::from_str::<serde_json::Value>(&data) + && let Some(content) = json.get("content").and_then(|v| v.as_str()) + { + yield Ok(ChatStreamEvent::TextChunk(content.to_string())); + } + } + "tool_call" => { + if let Ok(json) = serde_json::from_str::<serde_json::Value>(&data) { + let id = json.get("id").and_then(|v| v.as_str()).unwrap_or("").to_string(); + let name = json.get("name").and_then(|v| v.as_str()).unwrap_or("").to_string(); + let input = json.get("input").cloned().unwrap_or(serde_json::json!({})); + yield Ok(ChatStreamEvent::ToolCall { id, name, input }); + } + } + "tool_result" => { + if let Ok(json) = serde_json::from_str::<serde_json::Value>(&data) { + let tool_use_id = json.get("tool_use_id").and_then(|v| v.as_str()).unwrap_or("").to_string(); + let content = json.get("content").and_then(|v| v.as_str()).unwrap_or("").to_string(); + let is_error = json.get("is_error").and_then(|v| v.as_bool()).unwrap_or(false); + yield Ok(ChatStreamEvent::ToolResult { tool_use_id, content, is_error }); + } + } + "status" => { + if let Ok(json) = serde_json::from_str::<serde_json::Value>(&data) + && let Some(state) = json.get("state").and_then(|v| v.as_str()) + { + yield Ok(ChatStreamEvent::Status(state.to_string())); + } + } + "done" => { + if let Ok(json) = serde_json::from_str::<serde_json::Value>(&data) { + let session_id = json.get("session_id") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + yield Ok(ChatStreamEvent::Done { session_id }); + } else { + yield Ok(ChatStreamEvent::Done { session_id: String::new() }); + } + break; + } + "error" => { + if let Ok(json) = serde_json::from_str::<serde_json::Value>(&data) { + let message = json.get("message").and_then(|v| v.as_str()).unwrap_or("Unknown error").to_string(); + error!("SSE error: {}", message); + yield Ok(ChatStreamEvent::Error(message)); + } else { + error!("SSE error: {}", data); + yield Ok(ChatStreamEvent::Error(data)); + } + break; + } + _ => { + // Unknown event type, ignore + } + } + } + Err(e) => { + yield Err(eyre::eyre!("SSE error: {}", e)); + break; + } + } + } + }) +} + +fn hub_url(base: &str, path: &str) -> Result<Url> { + let base_with_slash = if base.ends_with('/') { + base.to_string() + } else { + format!("{base}/") + }; + let stripped = path.strip_prefix('/').unwrap_or(path); + Url::parse(&base_with_slash)? + .join(stripped) + .context("failed to build hub URL") +} + +fn detect_os() -> String { + match std::env::consts::OS { + "macos" => "macos".to_string(), + "linux" => "linux".to_string(), + "windows" => "windows".to_string(), + _ => "linux".to_string(), + } +} + +#[derive(Clone)] +enum Action { + Execute(String), + Insert(String), + Print(String), + Cancel, +} + +/// Serialize AppState to JSON for debug logging +fn state_to_json(state: &crate::tui::AppState) -> serde_json::Value { + let events: Vec<serde_json::Value> = state.events.iter().map(|e| e.to_json()).collect(); + + let mode = match state.mode { + AppMode::Input => "Input", + AppMode::Generating => "Generating", + AppMode::Streaming => "Streaming", + AppMode::Review => "Review", + AppMode::Error => "Error", + }; + + // Get input and cursor from textarea + let input = state.input(); + let cursor = state.textarea.cursor(); + + let mut json = serde_json::json!({ + "events": events, + "mode": mode, + "input": input, + "cursor_row": cursor.0, + "cursor_col": cursor.1, + "spinner_frame": state.spinner_frame, + "confirmation_pending": state.confirmation_pending, + }); + + // Add streaming fields if in streaming mode + if !state.streaming_text.is_empty() { + json["streaming_text"] = serde_json::json!(state.streaming_text); + } + if let Some(ref status) = state.streaming_status { + json["streaming_status"] = serde_json::json!(status.display_text()); + } + if let Some(ref err) = state.error { + json["error"] = serde_json::json!(err); + } + + json +} + +/// Debug logger that writes state changes to a file +struct DebugStateLogger { + file: std::fs::File, + entry_count: usize, + width: u16, +} + +impl DebugStateLogger { + fn new(path: &str) -> Result<Self> { + let file = std::fs::File::create(path) + .with_context(|| format!("Failed to create debug state file: {}", path))?; + // Get terminal width, default to 80 + let (width, _) = crossterm::terminal::size().unwrap_or((80, 24)); + Ok(Self { + file, + entry_count: 0, + width, + }) + } + + fn log(&mut self, label: &str, state: &crate::tui::AppState) { + use crate::tui::calculate_needed_height; + + self.entry_count += 1; + let timestamp_ms = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map(|d| d.as_millis()) + .unwrap_or(0); + + // Calculate the actual content height needed for this state + let content_height = calculate_needed_height(state, 0); + + let mut state_json = state_to_json(state); + // Add dimensions for accurate replay + state_json["width"] = serde_json::json!(self.width); + state_json["height"] = serde_json::json!(content_height); + + let entry = serde_json::json!({ + "entry": self.entry_count, + "label": label, + "timestamp_ms": timestamp_ms, + "state": state_json, + }); + + // Write as JSONL (one JSON object per line) + if let Err(e) = writeln!(self.file, "{}", entry) { + tracing::warn!("Failed to write debug state: {}", e); + } + let _ = self.file.flush(); + } +} + +async fn run_inline_tui( + endpoint: String, + token: String, + initial_prompt: Option<String>, + keep_output: bool, + debug_state_file: Option<String>, + settings: &atuin_client::settings::Settings, +) -> Result<Action> { + // Detect popup mode (only on Unix where atuin-hex socket is available) + #[cfg(unix)] + let mut popup_state = crate::tui::popup::try_setup_popup(); + #[cfg(not(unix))] + let mut popup_state: Option<()> = None; + + let popup_mode = popup_state.is_some(); + + // Initialize terminal guard: popup mode uses Fixed viewport, inline uses Inline + #[cfg(unix)] + let mut guard = if let Some(ref ps) = popup_state { + TerminalGuard::new_popup(ps.current_rect, ps.saved_screen.cursor_col)? + } else { + TerminalGuard::new(keep_output)? + }; + #[cfg(not(unix))] + let mut guard = TerminalGuard::new(keep_output)?; + let mut app = App::new(); + if let Some(prompt) = initial_prompt { + // Set initial text in textarea + let mut textarea = tui_textarea::TextArea::from(prompt.lines()); + // Disable underline on cursor line + textarea.set_cursor_line_style(ratatui::style::Style::default()); + // Enable word wrapping + textarea.set_wrap_mode(tui_textarea::WrapMode::Word); + // Move cursor to end + textarea.move_cursor(tui_textarea::CursorMove::End); + app.state.textarea = textarea; + } + + // Initialize debug state logger if requested + let mut debug_logger = debug_state_file + .map(|path| DebugStateLogger::new(&path)) + .transpose()?; + + // Helper macro to log state changes + macro_rules! log_state { + ($label:expr) => { + if let Some(ref mut logger) = debug_logger { + logger.log($label, &app.state); + } + }; + } + + // Log initial state + log_state!("init"); + + // Load theme + let mut theme_manager = ThemeManager::new(None, None); + let theme = theme_manager.load_theme(&settings.theme.name, None); + + // Initialize event loop + let mut event_loop = EventLoop::new(); + + // Track chat stream + let mut chat_stream: Option< + std::pin::Pin<Box<dyn futures::Stream<Item = Result<ChatStreamEvent>> + Send>>, + > = None; + + loop { + // Ensure viewport is large enough for current content (capped at terminal height) + // In popup mode, use the actual popup width for accurate height calculation + let card_width = if popup_mode { + #[cfg(unix)] + { + popup_state + .as_ref() + .map(|ps| { + ps.current_rect + .width + .saturating_sub(crate::tui::popup::POPUP_MARGIN * 2) + }) + .unwrap_or(0) + } + #[cfg(not(unix))] + { + 0 + } + } else { + 0 + }; + let needed_height = calculate_needed_height(&app.state, card_width); + + // Grow popup dynamically as content arrives + #[cfg(unix)] + if let Some(ref mut ps) = popup_state { + // Add vertical margin for visual separation from terminal content + let popup_height = needed_height.saturating_add(crate::tui::popup::POPUP_MARGIN * 2); + if let Some(new_rect) = ps.fit_to(popup_height) { + guard.resize_popup(new_rect)?; + } + } + + let actual_height = guard.ensure_height(needed_height)?; + + // Render current state + let anchor_col = guard.anchor_col(); + #[cfg(unix)] + let render_above = popup_state.as_ref().is_some_and(|ps| ps.render_above); + #[cfg(not(unix))] + let render_above = false; + + let ctx = RenderContext { + theme, + anchor_col, + textarea: Some(&app.state.textarea), + max_height: actual_height, + popup_mode, + render_above, + }; + // Handle draw errors gracefully - cursor position reads can fail during resize + if let Err(e) = guard.terminal().draw(|frame| { + render(frame, &app.state, &ctx); + }) { + let err_msg = e.to_string(); + if err_msg.contains("cursor position") { + // Cursor position read failed (common during terminal resize) + // Skip this frame and continue - next frame will likely succeed + tracing::debug!( + "Skipping frame due to cursor position read error: {}", + err_msg + ); + continue; + } + return Err(e.into()); + } + + // Get next event + let event = event_loop.run().await?; + + // Handle event based on app mode + match event { + AppEvent::Key(key) => { + app.handle_key(key); + log_state!("key"); + } + AppEvent::Tick => { + app.state.tick(); + + // Poll chat stream if active - keep polling until done regardless of mode + // (mode may change to Review before we receive the done event with session_id) + if let Some(stream) = &mut chat_stream { + let mut cx = std::task::Context::from_waker(futures::task::noop_waker_ref()); + match stream.as_mut().poll_next(&mut cx) { + std::task::Poll::Ready(Some(Ok(event))) => match event { + ChatStreamEvent::TextChunk(text) => { + trace!(text = %text, "Processing TextChunk"); + app.state.append_streaming_text(&text); + log_state!("text_chunk"); + } + ChatStreamEvent::ToolCall { id, name, input } => { + trace!(id = %id, name = %name, "Processing ToolCall"); + app.state.add_tool_call(id, name, input); + log_state!("tool_call"); + } + ChatStreamEvent::ToolResult { + tool_use_id, + content, + is_error, + } => { + trace!(tool_use_id = %tool_use_id, "Processing ToolResult"); + app.state.add_tool_result(tool_use_id, content, is_error); + log_state!("tool_result"); + } + ChatStreamEvent::Status(status) => { + trace!(status = %status, "Processing Status"); + app.state.update_streaming_status(&status); + log_state!("status"); + } + ChatStreamEvent::Done { session_id } => { + trace!(session_id = %session_id, "Processing Done"); + chat_stream = None; + if !session_id.is_empty() { + app.state.store_session_id(session_id); + } + app.state.finalize_streaming(); + log_state!("done"); + } + ChatStreamEvent::Error(msg) => { + trace!(error = %msg, "Processing Error"); + chat_stream = None; + app.state.streaming_error(msg); + log_state!("error"); + } + }, + std::task::Poll::Ready(Some(Err(e))) => { + chat_stream = None; + app.state.streaming_error(e.to_string()); + log_state!("stream_error"); + } + std::task::Poll::Ready(None) => { + chat_stream = None; + app.state.finalize_streaming(); + log_state!("stream_end"); + } + std::task::Poll::Pending => {} + } + } + } + _ => {} + } + + // Handle user cancellation (Esc during streaming) - drop the stream + if app.state.was_interrupted && chat_stream.is_some() { + debug!("User cancelled streaming, dropping chat stream"); + chat_stream = None; + app.state.was_interrupted = false; // Reset the flag + } + + // Check exit condition (includes Ctrl+C / SIGINT from event loop) + if app.state.should_exit || event_loop.is_shutdown() { + break; + } + + // Handle generation trigger - unified path for all turns + if app.state.mode == AppMode::Generating && chat_stream.is_none() { + // Get the last user message from events + let last_user_content = app.state.events.iter().rev().find_map(|e| { + if let ConversationEvent::UserMessage { content } = e { + Some(content.clone()) + } else { + None + } + }); + + if last_user_content.is_some() { + // Build messages in Claude API format + let messages = app.state.events_to_messages(); + + // Transition to streaming mode + app.state.start_streaming(); + log_state!("start_streaming"); + + // Start the chat stream + chat_stream = Some(create_chat_stream( + endpoint.clone(), + token.clone(), + app.state.session_id.clone(), + messages, + settings, + )); + } + } + } + + // Restore popup area before guard drops (guard skips cleanup in popup mode) + #[cfg(unix)] + if let Some(ref ps) = popup_state { + crate::tui::popup::restore(ps); + } + + // Map exit action to return value + let result = match app.state.exit_action { + Some(ExitAction::Execute(cmd)) => Action::Execute(cmd), + Some(ExitAction::Insert(cmd)) => Action::Insert(cmd), + _ => Action::Cancel, + }; + + Ok(result) +} + +struct RawModeGuard; + +impl Drop for RawModeGuard { + fn drop(&mut self) { + let _ = disable_raw_mode(); + } +} + +fn emit_shell_result(action: Action, output_for_hook: bool) { + if output_for_hook { + match action { + Action::Execute(output) => eprintln!("__atuin_ai_execute__:{output}"), + Action::Insert(output) => eprintln!("__atuin_ai_insert__:{output}"), + Action::Print(output) => eprintln!("__atuin_ai_print__:{output}"), + Action::Cancel => eprintln!("__atuin_ai_cancel__"), + } + } else { + match action { + Action::Execute(output) => eprintln!("{output}"), + Action::Insert(output) => eprintln!("{output}"), + Action::Print(output) => eprintln!("{output}"), + Action::Cancel => eprintln!(), + } + } +} + +fn wait_for_login_confirmation() -> Result<bool> { + enable_raw_mode().context("failed enabling raw mode for login prompt")?; + let _guard = RawModeGuard; + + loop { + let ev = event::read().context("failed to read login confirmation key")?; + if let Event::Key(key) = ev { + match key.code { + KeyCode::Enter => return Ok(true), + KeyCode::Esc => return Ok(false), + _ => {} + } + } + } +} diff --git a/crates/atuin-ai/src/lib.rs b/crates/atuin-ai/src/lib.rs new file mode 100644 index 00000000..2d86271d --- /dev/null +++ b/crates/atuin-ai/src/lib.rs @@ -0,0 +1,2 @@ +pub mod commands; +pub mod tui; diff --git a/crates/atuin-ai/src/tui/app.rs b/crates/atuin-ai/src/tui/app.rs new file mode 100644 index 00000000..ecb1eb81 --- /dev/null +++ b/crates/atuin-ai/src/tui/app.rs @@ -0,0 +1,157 @@ +use super::state::{AppMode, AppState, ExitAction}; +use crossterm::event::{KeyCode, KeyEvent, KeyModifiers}; +use tui_textarea::{Input, Key}; + +/// Thin wrapper around AppState for compatibility +/// All state lives in AppState, this just provides the handle_key interface +pub struct App { + pub state: AppState, +} + +impl App { + pub fn new() -> Self { + Self { + state: AppState::new(), + } + } + + /// Handle a key event. Returns true if render is needed. + pub fn handle_key(&mut self, key: KeyEvent) -> bool { + match self.state.mode { + AppMode::Input => self.handle_input_key(key), + AppMode::Generating => self.handle_generating_key(key), + AppMode::Streaming => self.handle_streaming_key(key), + AppMode::Review => self.handle_review_key(key), + AppMode::Error => self.handle_error_key(key), + } + } + + fn handle_input_key(&mut self, key: KeyEvent) -> bool { + // Handle special keys ourselves + match key.code { + KeyCode::Esc => { + self.state.exit(ExitAction::Cancel); + return true; + } + KeyCode::Enter => { + if self.state.input_is_empty() { + self.state.exit(ExitAction::Cancel); + } else { + self.state.start_generating(); + } + return true; + } + _ => {} + } + + // Delegate all other keys to textarea + // Manually convert crossterm KeyEvent to tui-textarea Input + // (needed due to crossterm version mismatch) + let tui_key = match key.code { + KeyCode::Char(c) => Key::Char(c), + KeyCode::Backspace => Key::Backspace, + KeyCode::Delete => Key::Delete, + KeyCode::Left => Key::Left, + KeyCode::Right => Key::Right, + KeyCode::Up => Key::Up, + KeyCode::Down => Key::Down, + KeyCode::Home => Key::Home, + KeyCode::End => Key::End, + KeyCode::PageUp => Key::PageUp, + KeyCode::PageDown => Key::PageDown, + KeyCode::Tab => Key::Tab, + _ => Key::Null, + }; + + if tui_key != Key::Null { + let input = Input { + key: tui_key, + ctrl: key.modifiers.contains(KeyModifiers::CONTROL), + alt: key.modifiers.contains(KeyModifiers::ALT), + shift: key.modifiers.contains(KeyModifiers::SHIFT), + }; + self.state.textarea.input(input); + } + true + } + + fn handle_generating_key(&mut self, key: KeyEvent) -> bool { + match key.code { + KeyCode::Esc => { + self.state.cancel_generation(); + true + } + _ => false, // Discard other keys during generation + } + } + + fn handle_streaming_key(&mut self, key: KeyEvent) -> bool { + match key.code { + KeyCode::Esc => { + self.state.cancel_streaming(); + true + } + _ => false, // Ignore other keys during streaming + } + } + + fn handle_review_key(&mut self, key: KeyEvent) -> bool { + match key.code { + KeyCode::Esc => { + self.state.confirmation_pending = false; // Clear confirmation state + self.state.exit(ExitAction::Cancel); + true + } + KeyCode::Enter => { + let cmd = self.state.current_command().map(|c| c.to_string()); + if let Some(cmd) = cmd { + if self.state.is_current_command_dangerous() && !self.state.confirmation_pending + { + // First Enter on dangerous command: enter confirmation mode + self.state.confirmation_pending = true; + } else { + // Second Enter (confirmation), or non-dangerous command: execute + self.state.confirmation_pending = false; + self.state.exit(ExitAction::Execute(cmd)); + } + } + true + } + KeyCode::Tab => { + let cmd = self.state.current_command().map(|c| c.to_string()); + if let Some(cmd) = cmd { + self.state.confirmation_pending = false; // Clear on Tab too + self.state.exit(ExitAction::Insert(cmd)); + } + true + } + KeyCode::Char('f') => { + // Changed from 'e' to 'f' for follow-up mode + self.state.confirmation_pending = false; // Clear on follow-up + self.state.start_edit_mode(); + true + } + _ => false, + } + } + + fn handle_error_key(&mut self, key: KeyEvent) -> bool { + match key.code { + KeyCode::Esc => { + self.state.exit(ExitAction::Cancel); + true + } + KeyCode::Enter | KeyCode::Char('r') => { + self.state.retry(); + true + } + _ => false, + } + } +} + +impl Default for App { + fn default() -> Self { + Self::new() + } +} diff --git a/crates/atuin-ai/src/tui/event.rs b/crates/atuin-ai/src/tui/event.rs new file mode 100644 index 00000000..8efbf522 --- /dev/null +++ b/crates/atuin-ai/src/tui/event.rs @@ -0,0 +1,303 @@ +use crate::tui::App; +use crossterm::event::{Event, EventStream, KeyEvent, KeyEventKind}; +use eyre::{Result, eyre}; +use futures::StreamExt; +use std::time::Duration; +use tokio::time; + +/// Base tick interval for the event loop (fast for responsive streaming) +const BASE_TICK_INTERVAL: Duration = Duration::from_millis(50); + +/// Application events that drive the TUI state machine. +/// +/// # Event Types +/// - `Key`: Keyboard input (filtered to KeyEventKind::Press only) +/// - `Tick`: Periodic event for updates (50ms base interval) +/// - `Resize`: Terminal window resize +/// - `StreamChunk/StreamDone/StreamError`: Placeholders for Phase 3 streaming +/// +/// # Design Decisions +/// - Fast 50ms base tick for responsive streaming; spinner timing handled in AppState +/// - Stream events are placeholders - will be wired to channels in Phase 3 +/// - Resize handling enables responsive layout adjustments +#[derive(Debug, Clone)] +pub enum AppEvent { + /// Keyboard input event (filtered to Press events only) + Key(KeyEvent), + + /// Periodic tick for updates (50ms base interval; spinner timing in AppState) + Tick, + + /// Terminal resize event (width, height) + Resize(u16, u16), + + /// Stream chunk received (Phase 3 placeholder) + StreamChunk(String), + + /// Stream completed successfully (Phase 3 placeholder) + StreamDone, + + /// Stream error occurred (Phase 3 placeholder) + StreamError(String), +} + +/// Async event loop that drives the TUI with prioritized event handling. +/// +/// # Priority Model (Biased Select) +/// 1. **Stream data** - Highest priority (future Phase 3 streaming) +/// 2. **Keyboard input** - Medium priority (user responsiveness) +/// 3. **Tick events** - Lowest priority (spinner animation) +/// +/// This ensures stream data is processed immediately when available, +/// keyboard input is responsive, and spinner updates don't block higher priority events. +/// +/// # Graceful Shutdown +/// - SIGINT (Ctrl+C) sets shutdown flag and breaks the loop +/// - EventStream close (stdin EOF) triggers shutdown +/// - Shutdown flag can be checked/set externally for controlled termination +/// +/// # Example +/// ```no_run +/// use atuin_ai::tui::EventLoop; +/// +/// # async fn example() -> eyre::Result<()> { +/// let mut event_loop = EventLoop::new(); +/// loop { +/// let event = event_loop.run().await?; +/// // Handle event... +/// # break; +/// } +/// # Ok(()) +/// # } +/// ``` +pub struct EventLoop { + /// Tick interval timer (created lazily on first run) + tick_timer: Option<time::Interval>, + + /// Flag indicating a render was requested (future use in Phase 2) + #[allow(dead_code)] + render_requested: bool, + + /// Shutdown flag - when true, event loop will terminate + shutdown: bool, +} + +impl EventLoop { + /// Create a new EventLoop with default settings. + /// + /// # Defaults + /// - Tick interval: 50ms base rate (spinner timing handled separately in AppState) + /// - Render requested: false + /// - Shutdown: false + pub fn new() -> Self { + Self { + tick_timer: None, + render_requested: false, + shutdown: false, + } + } + + /// Run the event loop, returning the next application event. + /// + /// # Priority Model + /// Uses `tokio::select!` with `biased;` mode to enforce priority: + /// 1. Stream data (placeholder for Phase 3) + /// 2. Keyboard input with rapid keypress batching + /// 3. Tick for spinner animation + /// + /// # Keyboard Handling + /// - Filters to KeyEventKind::Press on all platforms for safety + /// - Batching of rapid keypresses will be implemented in Phase 2 + /// - Currently returns individual key events + /// + /// # Graceful Shutdown + /// - SIGINT (Ctrl+C) triggers shutdown and returns last event + /// - EventStream close (stdin EOF) triggers shutdown + /// - Shutdown flag can be checked after this returns + /// + /// # Errors + /// - Returns error if terminal event stream encounters an error + /// - EventStream close is handled gracefully as shutdown signal + /// + /// # Example + /// ```no_run + /// # use atuin_ai::tui::EventLoop; + /// # async fn example() -> eyre::Result<()> { + /// let mut event_loop = EventLoop::new(); + /// while !event_loop.is_shutdown() { + /// match event_loop.run().await? { + /// // Handle events... + /// # _ => break, + /// } + /// } + /// # Ok(()) + /// # } + /// ``` + pub async fn run(&mut self) -> Result<AppEvent> { + // Create async event stream for keyboard/terminal events + let mut reader = EventStream::new(); + + // Get or create the tick timer (reused across calls to maintain timing) + // Uses fast base tick for responsive streaming; spinner timing handled in AppState + let tick_timer = self.tick_timer.get_or_insert_with(|| { + let mut interval = time::interval(BASE_TICK_INTERVAL); + // Skip the first immediate tick + interval.reset(); + interval + }); + + loop { + if self.shutdown { + break; + } + + // Biased select: prioritize stream > keyboard > tick + let event = tokio::select! { + biased; + + // Priority 1: Stream data (placeholder for Phase 3) + // In Phase 3, this will be: + // Some(chunk) = stream_rx.recv() => { ... } + + // Priority 2: Keyboard input + maybe_event = reader.next() => { + match maybe_event { + Some(Ok(Event::Key(key))) => { + // Filter to Press events only for cross-platform safety + if key.kind == KeyEventKind::Press { + // Note: Rapid keypress batching will be implemented in Phase 2 + // when we integrate with the state machine. + // For now, just return individual key events. + Some(AppEvent::Key(key)) + } else { + None + } + } + Some(Ok(Event::Resize(w, h))) => { + Some(AppEvent::Resize(w, h)) + } + Some(Err(e)) => { + return Err(eyre!("terminal event error: {}", e)); + } + None => { + // EventStream closed (stdin EOF) - trigger shutdown + self.shutdown = true; + None + } + _ => { + // Ignore other event types (mouse, focus, etc.) + None + } + } + } + + // Priority 3: Tick for spinner animation + _ = tick_timer.tick() => { + Some(AppEvent::Tick) + } + + // SIGINT handling (Ctrl+C) - cross-platform + _ = tokio::signal::ctrl_c() => { + self.shutdown = true; + // Return one more event to allow graceful shutdown handling + Some(AppEvent::Tick) + } + }; + + if let Some(app_event) = event { + return Ok(app_event); + } + } + + // Loop exited due to shutdown - return final tick to allow cleanup + Ok(AppEvent::Tick) + } + + /// Check if the event loop has been signaled to shut down. + /// + /// This can be used to cleanly exit the main TUI loop after receiving + /// a shutdown signal (Ctrl+C, stdin close, etc.) + pub fn is_shutdown(&self) -> bool { + self.shutdown + } + + /// Signal the event loop to shut down. + /// + /// The shutdown will take effect on the next iteration of `run()`. + pub fn shutdown(&mut self) { + self.shutdown = true; + } + + /// Poll for next event and apply to app state. + /// + /// This is a convenience method that combines `run()` with `App` state updates. + /// Returns true if app should continue, false if should exit. + /// + /// # Example + /// ```no_run + /// # use atuin_ai::tui::{EventLoop, App}; + /// # async fn example() -> eyre::Result<()> { + /// let mut event_loop = EventLoop::new(); + /// let mut app = App::new(); + /// + /// while event_loop.poll_and_apply(&mut app).await? { + /// // Render app state... + /// } + /// # Ok(()) + /// # } + /// ``` + pub async fn poll_and_apply(&mut self, app: &mut App) -> Result<bool> { + let event = self.run().await?; + + match event { + AppEvent::Key(key) => { + app.handle_key(key); + } + AppEvent::Tick => { + app.state.tick(); + } + AppEvent::Resize(_, _) => { + // Render will be triggered anyway + } + AppEvent::StreamChunk(_) | AppEvent::StreamDone | AppEvent::StreamError(_) => { + // Placeholder for Phase 3 + } + } + + Ok(!app.state.should_exit) + } +} + +impl Default for EventLoop { + fn default() -> Self { + Self::new() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_event_loop_creation() { + let event_loop = EventLoop::new(); + assert!(!event_loop.shutdown); + } + + #[test] + fn test_shutdown_flag() { + let mut event_loop = EventLoop::new(); + assert!(!event_loop.is_shutdown()); + + event_loop.shutdown(); + assert!(event_loop.is_shutdown()); + } + + // Note: Cannot easily test run() in unit tests since it requires a TTY. + // Integration tests should verify: + // 1. Tick events are generated at 150ms intervals + // 2. Keyboard events are properly filtered to Press only + // 3. Rapid keypresses are batched + // 4. SIGINT triggers graceful shutdown + // 5. Resize events are propagated correctly +} diff --git a/crates/atuin-ai/src/tui/mod.rs b/crates/atuin-ai/src/tui/mod.rs new file mode 100644 index 00000000..03a9c007 --- /dev/null +++ b/crates/atuin-ai/src/tui/mod.rs @@ -0,0 +1,16 @@ +pub mod app; +pub mod event; +#[cfg(unix)] +pub mod popup; +pub mod render; +pub mod spinner; +pub mod state; +pub mod terminal; +pub mod view_model; + +pub use app::App; +pub use event::{AppEvent, EventLoop}; +pub use render::{RenderContext, calculate_needed_height, markdown_to_spans}; +pub use state::{AppMode, AppState, ConversationEvent, ExitAction}; +pub use terminal::{TerminalGuard, install_panic_hook}; +pub use view_model::{Block, Blocks, Content}; diff --git a/crates/atuin-ai/src/tui/popup.rs b/crates/atuin-ai/src/tui/popup.rs new file mode 100644 index 00000000..c62b0e62 --- /dev/null +++ b/crates/atuin-ai/src/tui/popup.rs @@ -0,0 +1,363 @@ +use ratatui::layout::Rect; + +/// Maximum popup height (lines). Keeps context visible around the popup. +const MAX_POPUP_HEIGHT: u16 = 24; + +/// Minimum usable popup height. +const MIN_POPUP_HEIGHT: u16 = 5; + +/// Initial popup height — just enough for input + a small response. +const INITIAL_POPUP_HEIGHT: u16 = 5; + +/// Margin around the card in popup mode. +pub(crate) const POPUP_MARGIN: u16 = 0; + +/// Screen state captured from atuin-hex's screen server. +pub struct SavedScreen { + #[allow(dead_code)] + pub rows: u16, + #[allow(dead_code)] + pub cols: u16, + pub cursor_row: u16, + pub cursor_col: u16, + /// Pre-formatted ANSI bytes for each screen row, ready to write to stdout. + pub rows_data: Vec<Vec<u8>>, +} + +/// Popup mode state: saved screen + computed placement. +pub struct PopupState { + pub saved_screen: SavedScreen, + /// Maximum rect computed from placement (the ceiling for growth). + pub max_rect: Rect, + /// Current rect — starts small, grows as content arrives. + pub current_rect: Rect, + pub scroll_offset: u16, + /// True when the popup renders above the cursor (input at bottom of card). + pub render_above: bool, +} + +impl PopupState { + /// Resize the popup to fit `needed` lines of content. + /// + /// Grows or shrinks the popup as needed (clamped to max_rect / INITIAL_POPUP_HEIGHT). + /// When growing, clears the new rect area. When shrinking, restores freed rows + /// from the saved screen data. + /// + /// Returns `Some(new_rect)` if the size changed (caller must resize terminal), + /// or `None` if no change is needed. + pub fn fit_to(&mut self, needed: u16) -> Option<Rect> { + let new_height = needed.clamp(INITIAL_POPUP_HEIGHT, self.max_rect.height); + if new_height == self.current_rect.height { + return None; + } + + let old_rect = self.current_rect; + let growing = new_height > old_rect.height; + + if self.render_above { + let new_y = self.max_rect.y + self.max_rect.height - new_height; + self.current_rect = Rect::new(old_rect.x, new_y, old_rect.width, new_height); + } else { + self.current_rect = Rect::new(old_rect.x, old_rect.y, old_rect.width, new_height); + } + + if growing { + // Clear the entire new rect so the new Terminal doesn't leave + // ghost content from the old card. + self.clear_rows( + self.current_rect.y, + self.current_rect.y + self.current_rect.height, + ); + } else { + // Shrinking: restore freed rows from saved screen data, then + // clear the new (smaller) rect for the re-rendered card. + self.restore_rows(&old_rect); + self.clear_rows( + self.current_rect.y, + self.current_rect.y + self.current_rect.height, + ); + } + + Some(self.current_rect) + } + + /// Clear a range of terminal rows within the popup width. + fn clear_rows(&self, from_row: u16, to_row: u16) { + use crossterm::cursor::MoveTo; + use crossterm::execute; + use crossterm::style::{Attribute, SetAttribute}; + use std::io::{Write, stdout}; + + let mut out = stdout(); + for row in from_row..to_row { + let _ = execute!( + out, + MoveTo(self.current_rect.x, row), + SetAttribute(Attribute::Reset) + ); + let _ = write!( + out, + "{:width$}", + "", + width = self.current_rect.width as usize + ); + } + let _ = out.flush(); + } + + /// Restore rows that were freed by shrinking — the rows in old_rect + /// that are no longer covered by current_rect. + fn restore_rows(&self, old_rect: &Rect) { + use crossterm::cursor::MoveTo; + use crossterm::execute; + use crossterm::style::{Attribute, SetAttribute}; + use std::io::{Write, stdout}; + + let mut out = stdout(); + + // Determine which rows are freed + let (freed_start, freed_end) = if self.render_above { + // Shrinking from above: freed rows are at the old top + (old_rect.y, self.current_rect.y) + } else { + // Shrinking from below: freed rows are at the old bottom + ( + self.current_rect.y + self.current_rect.height, + old_rect.y + old_rect.height, + ) + }; + + for row in freed_start..freed_end { + let source_row = (row + self.scroll_offset) as usize; + + // Clear the popup region + let _ = execute!(out, MoveTo(old_rect.x, row), SetAttribute(Attribute::Reset),); + let _ = write!(out, "{:width$}", "", width = old_rect.width as usize); + + // Write back saved row data from column 0 + let _ = execute!(out, MoveTo(0, row)); + if let Some(row_bytes) = self.saved_screen.rows_data.get(source_row) { + let _ = out.write_all(row_bytes); + } + } + let _ = out.flush(); + } +} + +/// Try to set up popup overlay mode. +/// +/// Checks for `ATUIN_HEX_SOCKET`, fetches screen state, computes placement, +/// and scrolls the terminal if needed. Returns `None` if popup mode is not +/// available (no socket, fetch failed, etc.), in which case the caller should +/// fall back to inline mode. +pub fn try_setup_popup() -> Option<PopupState> { + use std::io::Write; + + let socket_path = std::env::var("ATUIN_HEX_SOCKET").ok()?; + let saved = fetch_screen_state(&socket_path)?; + + let (term_cols, term_rows) = crossterm::terminal::size().unwrap_or((saved.cols, saved.rows)); + // Full-width popup with margin for visual separation + let popup_width = term_cols; + let (rect, scroll, render_above) = compute_popup_placement( + saved.cursor_row, + saved.cursor_col, + term_rows, + term_cols, + popup_width, + ); + + // Scroll terminal up if needed to make room for the popup + if scroll > 0 { + let mut stdout = std::io::stdout(); + let _ = crossterm::execute!(stdout, crossterm::cursor::MoveTo(0, term_rows - 1)); + for _ in 0..scroll { + let _ = writeln!(stdout); + } + let _ = stdout.flush(); + } + + // Start with a small rect that grows as content arrives + let initial_height = INITIAL_POPUP_HEIGHT.min(rect.height); + let current_rect = if render_above { + // Anchor at the bottom of max_rect (near cursor), grow upward + Rect::new( + rect.x, + rect.y + rect.height - initial_height, + rect.width, + initial_height, + ) + } else { + // Anchor at the top of max_rect (near cursor), grow downward + Rect::new(rect.x, rect.y, rect.width, initial_height) + }; + + Some(PopupState { + saved_screen: saved, + max_rect: rect, + current_rect, + scroll_offset: scroll, + render_above, + }) +} + +/// Restore the screen area that was covered by the popup. +/// +/// Clears the popup region, then writes pre-formatted per-row ANSI bytes from +/// column 0 to correctly restore wide characters, colors, and all attributes. +pub fn restore(state: &PopupState) { + use crossterm::cursor::MoveTo; + use crossterm::execute; + use crossterm::style::{Attribute, SetAttribute}; + use std::io::{Write, stdout}; + + let saved = &state.saved_screen; + let popup_rect = state.current_rect; + let scroll_offset = state.scroll_offset; + + let mut stdout = stdout(); + + for dy in 0..popup_rect.height { + let target_row = popup_rect.y + dy; + let source_row = (target_row + scroll_offset) as usize; + + // Clear only the popup region with spaces + let _ = execute!( + stdout, + MoveTo(popup_rect.x, target_row), + SetAttribute(Attribute::Reset), + ); + let _ = write!(stdout, "{:width$}", "", width = popup_rect.width as usize); + + // Write back full row ANSI data from column 0 + let _ = execute!(stdout, MoveTo(0, target_row)); + if let Some(row_bytes) = saved.rows_data.get(source_row) { + let _ = stdout.write_all(row_bytes); + } + } + + // Restore cursor position (adjusted for any scrolling) + let _ = execute!( + stdout, + MoveTo( + saved.cursor_col, + saved.cursor_row.saturating_sub(scroll_offset) + ) + ); + let _ = stdout.flush(); +} + +/// Connect to atuin-hex's Unix socket and fetch the current screen state. +/// +/// The wire format is: +/// ```text +/// [rows: u16 BE][cols: u16 BE][cursor_row: u16 BE][cursor_col: u16 BE] +/// [row_0_len: u32 BE][row_0_bytes...] +/// [row_1_len: u32 BE][row_1_bytes...] +/// ... +/// ``` +fn fetch_screen_state(socket_path: &str) -> Option<SavedScreen> { + use std::io::Read; + use std::os::unix::net::UnixStream; + use std::time::Duration; + + let mut stream = UnixStream::connect(socket_path).ok()?; + stream.set_read_timeout(Some(Duration::from_secs(2))).ok()?; + + let mut data = Vec::new(); + stream.read_to_end(&mut data).ok()?; + + if data.len() < 8 { + return None; + } + + let rows = u16::from_be_bytes([data[0], data[1]]); + let cols = u16::from_be_bytes([data[2], data[3]]); + let cursor_row = u16::from_be_bytes([data[4], data[5]]); + let cursor_col = u16::from_be_bytes([data[6], data[7]]); + + let mut rows_data = Vec::with_capacity(rows as usize); + let mut offset = 8; + while offset + 4 <= data.len() { + let row_len = u32::from_be_bytes([ + data[offset], + data[offset + 1], + data[offset + 2], + data[offset + 3], + ]) as usize; + offset += 4; + if offset + row_len > data.len() { + break; + } + rows_data.push(data[offset..offset + row_len].to_vec()); + offset += row_len; + } + + Some(SavedScreen { + rows, + cols, + cursor_row, + cursor_col, + rows_data, + }) +} + +/// Compute popup placement for the AI card. +/// +/// Positions the popup near the cursor: below if there's room, above otherwise. +/// Uses a capped height (MAX_POPUP_HEIGHT) so the popup doesn't fill the screen. +/// +/// Returns `(popup_rect, scroll_offset, render_above)`: +/// - `render_above`: true when popup is above cursor (input should be at bottom) +/// - `scroll_offset`: lines the caller should scroll the terminal up +fn compute_popup_placement( + cursor_row: u16, + cursor_col: u16, + term_rows: u16, + term_cols: u16, + card_width: u16, +) -> (Rect, u16, bool) { + // Horizontal: anchor card near cursor, clamp to screen + let popup_w = card_width.min(term_cols); + let preferred_x = cursor_col.saturating_sub(2); + let max_x = term_cols.saturating_sub(popup_w); + let popup_x = preferred_x.min(max_x); + + // Vertical: use a reasonable height, not the full terminal + let max_h = MAX_POPUP_HEIGHT + .min(term_rows.saturating_sub(2)) + .max(MIN_POPUP_HEIGHT); + let space_above = cursor_row; + let space_below = term_rows.saturating_sub(cursor_row); + + if max_h <= space_below { + // Fits below cursor — input at top (close to prompt) + let popup_y = cursor_row; + (Rect::new(popup_x, popup_y, popup_w, max_h), 0, false) + } else if max_h <= space_above { + // Fits above cursor — input at bottom (close to prompt) + let popup_y = cursor_row.saturating_sub(max_h); + (Rect::new(popup_x, popup_y, popup_w, max_h), 0, true) + } else { + // Neither side fits fully — use whichever side has more space, + // scrolling the terminal if needed to reach MIN_POPUP_HEIGHT. + let render_above = space_above > space_below; + let available = if render_above { + space_above + } else { + space_below + }; + let h = available.max(MIN_POPUP_HEIGHT).min(max_h); + let scroll = h.saturating_sub(available); + let popup_y = if render_above { + cursor_row.saturating_sub(h + scroll) + } else { + cursor_row.saturating_sub(scroll) + }; + ( + Rect::new(popup_x, popup_y, popup_w, h), + scroll, + render_above, + ) + } +} diff --git a/crates/atuin-ai/src/tui/render.rs b/crates/atuin-ai/src/tui/render.rs new file mode 100644 index 00000000..9326b0df --- /dev/null +++ b/crates/atuin-ai/src/tui/render.rs @@ -0,0 +1,712 @@ +use atuin_client::theme::{Meaning, Theme}; +use pulldown_cmark::{Event, Parser, Tag, TagEnd}; +use ratatui::{ + Frame, + backend::FromCrossterm, + layout::{Alignment, Constraint, Direction, Layout, Rect}, + style::{Modifier, Style}, + text::{Line, Span}, + widgets::{Block as RatatuiBlock, Borders, Padding, Paragraph, Wrap}, +}; +use tui_textarea::TextArea; + +use super::spinner::active_frame; +use super::state::AppState; +use super::view_model::{Blocks, Content, WarningKind}; + +/// Fixed card width for the TUI +pub(crate) const CARD_WIDTH: u16 = 64; + +pub struct RenderContext<'a> { + pub theme: &'a Theme, + pub anchor_col: u16, + pub textarea: Option<&'a TextArea<'static>>, + /// Maximum viewport height (for scroll calculations) + pub max_height: u16, + /// When true, the viewport is a fixed rect already positioned for the card. + /// The card fills the entire viewport instead of positioning via anchor_col. + pub popup_mode: bool, + /// When true, blocks are rendered in reverse order so that the input field + /// appears at the bottom of the card (close to the prompt when the popup + /// is above the cursor). + pub render_above: bool, +} + +/// Calculate the height needed to render the current state. +/// Used to dynamically resize the viewport before rendering. +/// `card_width` is the outer card width (including borders); pass 0 to use CARD_WIDTH default. +pub fn calculate_needed_height(state: &AppState, card_width: u16) -> u16 { + let view = Blocks::from_state(state); + let w = if card_width > 0 { + card_width + } else { + CARD_WIDTH + }; + let content_width = usize::from(w.saturating_sub(4)).max(1); + + let mut total_height = 0u16; + for (idx, block) in view.items.iter().enumerate() { + if idx > 0 { + total_height = total_height.saturating_add(1); // separator + total_height = total_height.saturating_add(1); // leading blank after separator + } + total_height = + total_height.saturating_add(calculate_block_height(&block.content, content_width)); + } + + // Add borders (2) + top padding (1), minimum 5 + total_height.saturating_add(3).max(5) +} + +/// Main render function: derives view model from state, then renders it +pub fn render(frame: &mut Frame, state: &AppState, ctx: &RenderContext) { + // PURE DERIVATION: view model is always rebuilt from state + let view = Blocks::from_state(state); + + // Render the derived view model + render_view(frame, &view, ctx); +} + +fn render_view(frame: &mut Frame, view: &Blocks, ctx: &RenderContext) { + let full_area = frame.area(); + + // In popup mode, the viewport is already positioned and sized for the card. + // Clear it to prevent background bleed-through, then inset by margin for the card. + let (area, card_x, desired_width) = if ctx.popup_mode { + #[cfg(unix)] + use super::popup::POPUP_MARGIN; + #[cfg(not(unix))] + const POPUP_MARGIN: u16 = 0; + frame.render_widget(ratatui::widgets::Clear, full_area); + let inset = full_area.inner(ratatui::layout::Margin { + horizontal: POPUP_MARGIN, + vertical: POPUP_MARGIN, + }); + (inset, inset.x, inset.width) + } else { + let dw = CARD_WIDTH.min(full_area.width.saturating_sub(2)).max(32); + let max_x = full_area.x + full_area.width.saturating_sub(dw); + let preferred_x = full_area.x + ctx.anchor_col.saturating_sub(2); + (full_area, preferred_x.min(max_x), dw) + }; + let content_width = usize::from(desired_width.saturating_sub(4)).max(1); + + // Build ordered items list — the active content (input/LLM response) + // should always be closest to the cursor/prompt: + // - Popup below cursor (render_above=false): reverse so active is at top + // - Popup above cursor (render_above=true): normal order, active is at bottom + // - Inline mode: normal order (no reversal) + let items: Vec<&super::view_model::Block> = if ctx.popup_mode && !ctx.render_above { + view.items.iter().rev().collect() + } else { + view.items.iter().collect() + }; + + // Calculate height from view model + let mut total_height = 0u16; + for (idx, block) in items.iter().enumerate() { + if idx > 0 { + total_height = total_height.saturating_add(1); // separator + total_height = total_height.saturating_add(1); // leading blank after separator + } + total_height = + total_height.saturating_add(calculate_block_height(&block.content, content_width)); + } + + let desired_height = total_height + .saturating_add(3) // borders (2) + top padding (1), no bottom padding + .max(5); + + // Cap card height at viewport height to prevent overflow + let actual_height = desired_height.min(area.height); + + // Calculate scroll offset to keep the active content visible when overflowing. + // When render_above=false (popup below cursor), items are reversed so the active + // content (input/spinner) is at the top — scroll_offset stays 0 to show the top. + // Otherwise, scroll to show the bottom where the active content lives. + let scroll_offset = if ctx.popup_mode && !ctx.render_above { + 0 + } else { + desired_height.saturating_sub(actual_height) + }; + + let card = Rect { + x: card_x, + y: area.y, + width: desired_width, + height: actual_height, + }; + + // Get title from first block in ORIGINAL order (always the input block) + let title = view + .items + .first() + .and_then(|b| b.title.as_deref()) + .unwrap_or("Describe the command you'd like to generate:"); + + // Create bordered frame + // Padding: left=1, right=1, top=1, bottom=0 (blocks have trailing blanks) + let mut outer_block = RatatuiBlock::default() + .borders(Borders::ALL) + .title(title) + .title_bottom(Line::from(view.footer).alignment(Alignment::Right)) + .padding(Padding::new(1, 1, 1, 0)); + + // Status bar: transient status on the bottom border, left-aligned + if let Some(ref sb) = view.status_bar { + let style = Style::from_crossterm(ctx.theme.as_style(Meaning::Annotation)); + let spinner = active_frame(sb.frame); + let status_text = format!(" {} {} ", spinner, sb.text); + outer_block = outer_block + .title_bottom(Line::from(Span::styled(status_text, style)).alignment(Alignment::Left)); + } + + let inner_area = outer_block.inner(card); + frame.render_widget(outer_block, card); + + // Render blocks (with scroll offset for overflowing content) + render_blocks_content(frame, &items, ctx, inner_area, card.width, scroll_offset); +} + +fn render_blocks_content( + frame: &mut Frame, + items: &[&super::view_model::Block], + ctx: &RenderContext, + area: Rect, + card_width: u16, + scroll_offset: u16, +) { + let content_width = usize::from(area.width).max(1); + + // Build layout constraints for full content + let mut constraints = Vec::new(); + let mut block_heights = Vec::new(); + for (idx, block) in items.iter().enumerate() { + if idx > 0 { + constraints.push(Constraint::Length(1)); // separator + constraints.push(Constraint::Length(1)); // leading blank after separator + block_heights.push(1); + block_heights.push(1); + } + let height = calculate_block_height(&block.content, content_width); + constraints.push(Constraint::Length(height)); + block_heights.push(height); + } + + if constraints.is_empty() { + return; + } + + // Calculate cumulative heights to find which blocks are visible after scrolling + let mut cumulative: Vec<u16> = Vec::with_capacity(block_heights.len() + 1); + cumulative.push(0); + for h in &block_heights { + cumulative.push(cumulative.last().unwrap() + h); + } + + // Render each chunk, offsetting by scroll_offset and clipping to visible area + let chunks = Layout::default() + .direction(Direction::Vertical) + .constraints(constraints) + .split(area); + + let mut chunk_idx = 0; + for (idx, block) in items.iter().enumerate() { + if idx > 0 { + // Check if separator is visible (its position minus scroll_offset) + let sep_start = cumulative[chunk_idx]; + if sep_start >= scroll_offset && sep_start < scroll_offset + area.height { + let adjusted_chunk = Rect { + y: area.y + sep_start - scroll_offset, + ..chunks[chunk_idx] + }; + render_separator(frame, adjusted_chunk, ctx, card_width); + } + chunk_idx += 1; + chunk_idx += 1; // skip leading blank + } + + // Check if this block is at least partially visible + let block_start = cumulative[chunk_idx]; + let block_end = cumulative[chunk_idx + 1]; + + // Block is visible if it starts before viewport end and ends after viewport start + if block_start < scroll_offset + area.height && block_end > scroll_offset { + // Calculate visible portion + let visible_start = block_start.max(scroll_offset); + let visible_end = block_end.min(scroll_offset + area.height); + + let adjusted_chunk = Rect { + x: area.x, + y: area.y + visible_start - scroll_offset, + width: area.width, + height: visible_end - visible_start, + }; + + render_block_content(frame, &block.content, adjusted_chunk, ctx); + } + + chunk_idx += 1; + } +} + +/// Render all content items in a block +fn render_block_content(frame: &mut Frame, content: &[Content], area: Rect, ctx: &RenderContext) { + if content.is_empty() { + return; + } + + let content_width = usize::from(area.width).max(1); + + // Build layout constraints for each content item WITH spacing between items + let mut constraints = Vec::new(); + for (idx, c) in content.iter().enumerate() { + if idx > 0 { + constraints.push(Constraint::Length(1)); // blank line between items + } + constraints.push(Constraint::Length(calculate_single_content_height( + c, + content_width, + ))); + } + + let chunks = Layout::default() + .direction(Direction::Vertical) + .constraints(constraints) + .split(area); + + let mut chunk_idx = 0; + for (idx, item) in content.iter().enumerate() { + if idx > 0 { + chunk_idx += 1; // skip the blank line chunk + } + render_single_content(frame, item, chunks[chunk_idx], ctx); + chunk_idx += 1; + } +} + +/// Render a single content item using ratatui's native wrapping. +/// Symbol is rendered at column 0, text wraps in columns 2+ (offset area). +fn render_single_content(frame: &mut Frame, content: &Content, area: Rect, ctx: &RenderContext) { + // Helper to create offset text area (2 chars for symbol column) + let text_area = Rect { + x: area.x.saturating_add(2), + y: area.y, + width: area.width.saturating_sub(2), + height: area.height, + }; + + match content { + Content::Input { text, active, .. } => { + let symbol_style = Style::from_crossterm(ctx.theme.as_style(Meaning::Guidance)); + let text_style = Style::from_crossterm(ctx.theme.as_style(Meaning::Base)); + + // Render ">" symbol at column 0 + render_symbol(frame, ">", symbol_style, area); + + if *active { + // Active input: render TextArea widget (handles cursor display) + if let Some(textarea) = ctx.textarea { + frame.render_widget(textarea, text_area); + } + } else { + // Inactive input: render as plain paragraph + let paragraph = Paragraph::new(text.as_str()) + .style(text_style) + .wrap(Wrap { trim: false }); + frame.render_widget(paragraph, text_area); + } + } + + Content::Command { text, faded } => { + let symbol_style = Style::from_crossterm(ctx.theme.as_style(Meaning::Important)); + let mut text_style = Style::from_crossterm(ctx.theme.as_style(Meaning::Base)); + if *faded { + text_style = text_style.add_modifier(Modifier::DIM); + } + + render_symbol(frame, "$", symbol_style, area); + + let paragraph = Paragraph::new(text.as_str()) + .style(text_style) + .wrap(Wrap { trim: false }); + frame.render_widget(paragraph, text_area); + } + + Content::Text { markdown } => { + // No symbol, just indent - render directly in offset area + let text_style = Style::from_crossterm(ctx.theme.as_style(Meaning::Base)); + + let paragraph = Paragraph::new(markdown.as_str()) + .style(text_style) + .wrap(Wrap { trim: false }); + frame.render_widget(paragraph, text_area); + } + + Content::Error { message } => { + let symbol_style = Style::from_crossterm(ctx.theme.as_style(Meaning::AlertError)); + let text_style = Style::from_crossterm(ctx.theme.as_style(Meaning::Base)); + + render_symbol(frame, "!", symbol_style, area); + + let paragraph = Paragraph::new(message.as_str()) + .style(text_style) + .wrap(Wrap { trim: false }); + frame.render_widget(paragraph, text_area); + } + + Content::Warning { + kind, + text, + pending_confirm, + } => { + let (symbol, meaning) = match kind { + WarningKind::Danger => ("!", Meaning::AlertError), + WarningKind::LowConfidence => ("?", Meaning::AlertWarn), + }; + let symbol_style = Style::from_crossterm(ctx.theme.as_style(meaning)); + let text_style = Style::from_crossterm(ctx.theme.as_style(Meaning::Base)); + + let display_text = if *pending_confirm { + "Press Enter again to run this dangerous command" + } else { + text.as_str() + }; + + render_symbol(frame, symbol, symbol_style, area); + + let paragraph = Paragraph::new(display_text) + .style(text_style) + .wrap(Wrap { trim: false }); + frame.render_widget(paragraph, text_area); + } + + Content::Spinner { + frame: spinner_frame, + status_text, + } => { + let style = Style::from_crossterm(ctx.theme.as_style(Meaning::Annotation)); + let symbol = active_frame(*spinner_frame); + + render_symbol(frame, symbol, style, area); + + let paragraph = Paragraph::new(status_text.as_str()).style(style); + frame.render_widget(paragraph, text_area); + } + + Content::ToolStatus { + completed_count, + current_label, + frame: spinner_frame, + } => { + let style = Style::from_crossterm(ctx.theme.as_style(Meaning::Annotation)); + + let (symbol, text) = if let Some(label) = current_label { + let spinner = active_frame(*spinner_frame); + let text = if *completed_count > 0 { + format!( + "{} (used {} tool{})", + label, + completed_count, + if *completed_count == 1 { "" } else { "s" } + ) + } else { + label.clone() + }; + (spinner, text) + } else { + ( + "\u{2713}", + format!( + "Used {} tool{}", + completed_count, + if *completed_count == 1 { "" } else { "s" } + ), + ) + }; + + render_symbol(frame, symbol, style, area); + + let paragraph = Paragraph::new(text).style(style); + frame.render_widget(paragraph, text_area); + } + } +} + +/// Render a single-character symbol at the start of an area +fn render_symbol(frame: &mut Frame, symbol: &str, style: Style, area: Rect) { + let symbol_area = Rect { + x: area.x, + y: area.y, + width: 1, + height: 1, + }; + frame.render_widget(Paragraph::new(symbol).style(style), symbol_area); +} + +fn render_separator(frame: &mut Frame, area: Rect, ctx: &RenderContext, card_width: u16) { + let style = Style::from_crossterm(ctx.theme.as_style(Meaning::Muted)); + + // Build separator: ├ + ─ repeated + ┤ spanning the full card width + // -2 for the ├ and ┤ characters themselves + let inner_width = card_width.saturating_sub(2) as usize; + let separator = format!( + "\u{251c}{}\u{2524}", // ├ ... ┤ + "\u{2500}".repeat(inner_width) // ─ + ); + + let paragraph = Paragraph::new(Span::styled(separator, style)); + + // Render at x offset to overlap the border (area is inside padding, border is 2 chars left) + let sep_area = Rect { + x: area.x.saturating_sub(2), // move left to overlap left border + y: area.y, + width: card_width, + height: 1, + }; + frame.render_widget(paragraph, sep_area); +} + +/// Calculate total height for all content items in a block +fn calculate_block_height(content: &[Content], width: usize) -> u16 { + let content_height: u16 = content + .iter() + .map(|c| calculate_single_content_height(c, width)) + .sum(); + + // Add spacing between items (n-1 blank lines for n items) + let spacing = if content.len() > 1 { + (content.len() - 1) as u16 + } else { + 0 + }; + + // Add 1 for trailing blank line (padding after content) + content_height.saturating_add(spacing).saturating_add(1) +} + +/// Calculate height for a single content item. +/// Uses ratatui's Paragraph::line_count for consistency with rendering. +fn calculate_single_content_height(content: &Content, width: usize) -> u16 { + // Text area is offset by 2 for symbol column + let text_width = width.saturating_sub(2); + + match content { + // Input uses word wrapping (WrapMode::Word) in TextArea, which can produce + // more lines than character wrapping since it won't break words mid-word + Content::Input { text, active, .. } => { + if *active { + // For active input, use word-wrap line counting to match TextArea behavior + let (lines, last_line_width) = + word_wrap_line_count_with_last_width(text, text_width); + // Only add extra line for cursor if the last line is full + if last_line_width >= text_width { + lines.saturating_add(1) + } else { + lines + } + } else { + line_count_wrapped(text, text_width) + } + } + Content::Command { text, .. } => line_count_wrapped(text, text_width), + Content::Text { markdown } => line_count_wrapped(markdown, text_width), + Content::Error { message } => line_count_wrapped(message, text_width), + Content::Warning { + text, + pending_confirm, + .. + } => { + let display_text = if *pending_confirm { + "Press Enter again to run this dangerous command" + } else { + text.as_str() + }; + line_count_wrapped(display_text, text_width) + } + Content::Spinner { .. } => 1, + Content::ToolStatus { .. } => 1, + } +} + +/// Count lines when text is wrapped at given width. +/// Uses ratatui's Paragraph::line_count for accurate wrapping calculation. +fn line_count_wrapped(text: &str, width: usize) -> u16 { + if width == 0 { + return 1; + } + + let paragraph = Paragraph::new(text).wrap(Wrap { trim: false }); + paragraph.line_count(width as u16).max(1) as u16 +} + +/// Count lines using word-wrap algorithm (matches TextArea's WrapMode::Word). +/// Words won't be broken mid-word, so this may produce more lines than character wrapping. +/// Returns (line_count, last_line_width) so caller can determine if cursor needs extra space. +fn word_wrap_line_count_with_last_width(text: &str, width: usize) -> (u16, usize) { + if width == 0 || text.is_empty() { + return (1, 0); + } + + let mut line_count = 0u16; + let mut current_line_width = 0usize; + + for line in text.lines() { + if line.is_empty() { + line_count += 1; + current_line_width = 0; + continue; + } + + let mut line_started = false; + + for word in line.split_whitespace() { + let word_width = unicode_width::UnicodeWidthStr::width(word); + + if !line_started { + // First word on line + if word_width > width { + // Word is longer than width, it will be split by character + // Count how many lines it takes + line_count += word_width.div_ceil(width) as u16; + current_line_width = word_width % width; + if current_line_width == 0 { + current_line_width = 0; + line_started = false; + } else { + line_started = true; + } + } else { + current_line_width = word_width; + line_started = true; + } + } else { + // Subsequent word - need space before it + let needed = current_line_width + 1 + word_width; + if needed > width { + // Word doesn't fit, start new line + line_count += 1; + if word_width > width { + // Word itself is too long, will be split + line_count += word_width.div_ceil(width) as u16; + current_line_width = word_width % width; + if current_line_width == 0 { + line_started = false; + } + } else { + current_line_width = word_width; + } + } else { + current_line_width = needed; + } + } + } + + // Count the last line of this logical line + if line_started { + line_count += 1; + } + } + + // Handle case where text has no lines() output (empty or just whitespace) + if line_count == 0 { + line_count = 1; + current_line_width = 0; + } + + (line_count, current_line_width) +} + +/// Convert markdown to styled spans (existing function, kept as-is) +pub fn markdown_to_spans<'a>(text: &'a str, theme: &'a Theme) -> Vec<Line<'a>> { + let parser = Parser::new(text); + let mut lines: Vec<Vec<Span<'a>>> = vec![Vec::new()]; + let mut current_line = 0; + + let base_style = Style::from_crossterm(theme.as_style(Meaning::Base)); + let code_style = Style::from_crossterm(theme.as_style(Meaning::Important)); + let mut style_stack: Vec<Style> = vec![base_style]; + let mut in_code_block = false; + + for event in parser { + match event { + Event::Start(Tag::Strong) => { + let bold_style = style_stack + .last() + .copied() + .unwrap_or(base_style) + .add_modifier(Modifier::BOLD); + style_stack.push(bold_style); + } + Event::End(TagEnd::Strong) => { + style_stack.pop(); + } + Event::Start(Tag::Emphasis) => { + let underline_style = style_stack + .last() + .copied() + .unwrap_or(base_style) + .add_modifier(Modifier::UNDERLINED); + style_stack.push(underline_style); + } + Event::End(TagEnd::Emphasis) => { + style_stack.pop(); + } + Event::Start(Tag::CodeBlock(_)) => { + in_code_block = true; + // Start new line for code block + if !lines[current_line].is_empty() { + current_line += 1; + lines.push(Vec::new()); + } + } + Event::End(TagEnd::CodeBlock) => { + in_code_block = false; + // Ensure blank line after code block + if !lines[current_line].is_empty() { + current_line += 1; + lines.push(Vec::new()); + } + } + Event::Code(code) => { + lines[current_line].push(Span::styled(format!("`{}`", code), code_style)); + } + Event::Text(text) => { + let current_style = if in_code_block { + // Use Important style for code block content + code_style + } else { + style_stack.last().copied().unwrap_or(base_style) + }; + let parts: Vec<&str> = text.split('\n').collect(); + for (i, part) in parts.iter().enumerate() { + if i > 0 { + current_line += 1; + lines.push(Vec::new()); + } + if !part.is_empty() { + lines[current_line].push(Span::styled(part.to_string(), current_style)); + } + } + } + Event::SoftBreak => { + let current_style = style_stack.last().copied().unwrap_or(base_style); + lines[current_line].push(Span::styled(" ", current_style)); + } + Event::HardBreak => { + current_line += 1; + lines.push(Vec::new()); + } + Event::Start(Tag::Paragraph) => { + if current_line > 0 || !lines[0].is_empty() { + current_line += 1; + lines.push(Vec::new()); + } + } + Event::End(TagEnd::Paragraph) => {} + _ => {} + } + } + + lines.into_iter().map(Line::from).collect() +} diff --git a/crates/atuin-ai/src/tui/spinner.rs b/crates/atuin-ai/src/tui/spinner.rs new file mode 100644 index 00000000..138e0269 --- /dev/null +++ b/crates/atuin-ai/src/tui/spinner.rs @@ -0,0 +1,99 @@ +//! Spinner styles and configuration for TUI animations +//! +//! To experiment with different spinners, change `ACTIVE_SPINNER` below. + +use std::time::Duration; + +/// Active spinner style - change this to experiment with different styles +pub const ACTIVE_SPINNER: SpinnerStyle = SpinnerStyle::Dots; + +/// Spinner style definitions +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum SpinnerStyle { + /// Classic ASCII line spinner: / - \ | + Line, + /// Braille dots pattern + Dots, + /// Growing/shrinking dots + Pulse, + /// Simple arrow rotation + Arrow, + /// Block building + Block, +} + +impl SpinnerStyle { + /// Get the frames for this spinner style + pub const fn frames(&self) -> &'static [&'static str] { + match self { + SpinnerStyle::Line => &["/", "-", "\\", "|"], + SpinnerStyle::Dots => &["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"], + SpinnerStyle::Pulse => &["·", "•", "●", "•"], + SpinnerStyle::Arrow => &["←", "↖", "↑", "↗", "→", "↘", "↓", "↙"], + SpinnerStyle::Block => &[ + "▏", "▎", "▍", "▌", "▋", "▊", "▉", "█", "▉", "▊", "▋", "▌", "▍", "▎", "▏", + ], + } + } + + /// Get the recommended tick interval for this spinner style + /// Faster spinners need shorter intervals to look smooth + pub const fn tick_interval(&self) -> Duration { + match self { + SpinnerStyle::Line => Duration::from_millis(150), + SpinnerStyle::Dots => Duration::from_millis(80), + SpinnerStyle::Pulse => Duration::from_millis(200), + SpinnerStyle::Arrow => Duration::from_millis(100), + SpinnerStyle::Block => Duration::from_millis(80), + } + } + + /// Get the frame at the given index (wraps around) + pub fn frame_at(&self, index: usize) -> &'static str { + let frames = self.frames(); + frames[index % frames.len()] + } + + /// Get the number of frames in this spinner + pub fn frame_count(&self) -> usize { + self.frames().len() + } +} + +/// Get the active spinner's frame at the given index +pub fn active_frame(index: usize) -> &'static str { + ACTIVE_SPINNER.frame_at(index) +} + +/// Get the active spinner's tick interval +pub fn active_tick_interval() -> Duration { + ACTIVE_SPINNER.tick_interval() +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_frame_wrapping() { + let style = SpinnerStyle::Line; + assert_eq!(style.frame_at(0), "/"); + assert_eq!(style.frame_at(4), "/"); // wraps + assert_eq!(style.frame_at(5), "-"); + } + + #[test] + fn test_all_styles_have_frames() { + let styles = [ + SpinnerStyle::Line, + SpinnerStyle::Dots, + SpinnerStyle::Pulse, + SpinnerStyle::Arrow, + SpinnerStyle::Block, + ]; + for style in styles { + assert!(!style.frames().is_empty()); + assert!(style.tick_interval().as_millis() > 0); + } + } +} diff --git a/crates/atuin-ai/src/tui/state.rs b/crates/atuin-ai/src/tui/state.rs new file mode 100644 index 00000000..ba9c8ac6 --- /dev/null +++ b/crates/atuin-ai/src/tui/state.rs @@ -0,0 +1,530 @@ +//! Domain state types for the TUI application +//! +//! This module contains the core state types that represent the application's +//! domain model. Conversation events match the API protocol format. + +use std::time::Instant; +use tui_textarea::TextArea; + +use super::spinner::{ACTIVE_SPINNER, active_tick_interval}; + +/// Streaming status indicators from server +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum StreamingStatus { + Processing, + Searching, + Thinking, + WaitingForTools, +} + +impl StreamingStatus { + pub fn from_status_str(s: &str) -> Self { + match s { + "processing" => Self::Processing, + "searching" => Self::Searching, + "waiting_for_tools" => Self::WaitingForTools, + _ => Self::Thinking, // Default to thinking for "thinking" and unknown + } + } + + pub fn display_text(&self) -> &'static str { + match self { + Self::Processing => "Processing...", + Self::Searching => "Searching...", + Self::Thinking => "Thinking...", + Self::WaitingForTools => "Waiting for tools...", + } + } +} + +/// Conversation event types matching the API protocol +#[derive(Debug, Clone)] +pub enum ConversationEvent { + /// User message (what the user typed) + UserMessage { content: String }, + /// Text content from assistant (streamed or complete) + Text { content: String }, + /// Tool call from assistant + ToolCall { + id: String, + name: String, + input: serde_json::Value, + }, + /// Tool result (usually from server-side execution) + ToolResult { + tool_use_id: String, + content: String, + is_error: bool, + }, +} + +impl ConversationEvent { + /// Convert to JSON for API calls + pub fn to_json(&self) -> serde_json::Value { + match self { + ConversationEvent::UserMessage { content } => serde_json::json!({ + "type": "user_message", + "content": content + }), + ConversationEvent::Text { content } => serde_json::json!({ + "type": "text", + "content": content + }), + ConversationEvent::ToolCall { id, name, input } => serde_json::json!({ + "type": "tool_call", + "id": id, + "name": name, + "input": input + }), + ConversationEvent::ToolResult { + tool_use_id, + content, + is_error, + } => serde_json::json!({ + "type": "tool_result", + "tool_use_id": tool_use_id, + "content": content, + "is_error": is_error + }), + } + } + + /// Extract command from a suggest_command tool call + pub fn as_command(&self) -> Option<&str> { + if let ConversationEvent::ToolCall { name, input, .. } = self + && name == "suggest_command" + { + // command can be null for pure conversational turns + return input.get("command").and_then(|v| v.as_str()); + } + None + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum AppMode { + /// User is typing input + Input, + /// Waiting for generation (showing spinner) + Generating, + /// Streaming SSE response + Streaming, + /// Reviewing generated command + Review, + /// Error state, can retry + Error, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ExitAction { + /// Run the command + Execute(String), + /// Insert command without running + Insert(String), + /// User canceled + Cancel, +} + +/// Application state - the domain model +/// +/// Conversation is stored as a sequence of events matching the API protocol. +/// The view model is derived from this state via `Blocks::from_state()`. +pub struct AppState { + /// Current application mode + pub mode: AppMode, + /// Conversation events (source of truth, matches API protocol) + pub events: Vec<ConversationEvent>, + /// Text being streamed (accumulated, flushed to Text event on completion) + pub streaming_text: String, + /// Active text input (uses tui-textarea for proper cursor handling) + pub textarea: TextArea<'static>, + /// Current error message (renders at end of blocks) + pub error: Option<String>, + /// Whether app should exit + pub should_exit: bool, + /// Exit action (set when exiting) + pub exit_action: Option<ExitAction>, + /// Session ID from server (store after first response, send on subsequent) + pub session_id: Option<String>, + /// Current streaming status (for spinner text) + pub streaming_status: Option<StreamingStatus>, + /// Whether current turn was interrupted by user + pub was_interrupted: bool, + /// Spinner animation state + pub spinner_frame: usize, + /// When spinner frame last advanced (for timing control) + pub last_spinner_tick: Instant, + /// When streaming started (for spinner delay) + pub streaming_started: Option<Instant>, + /// True when user has pressed Enter once on a dangerous command + pub confirmation_pending: bool, +} + +/// Create a TextArea with our preferred configuration +fn create_textarea() -> TextArea<'static> { + let mut textarea = TextArea::default(); + // Disable underline on cursor line - it's distracting + textarea.set_cursor_line_style(ratatui::style::Style::default()); + // Enable word wrapping + textarea.set_wrap_mode(tui_textarea::WrapMode::Word); + textarea +} + +impl AppState { + pub fn new() -> Self { + Self { + mode: AppMode::Input, + events: Vec::new(), + streaming_text: String::new(), + textarea: create_textarea(), + error: None, + should_exit: false, + exit_action: None, + session_id: None, + streaming_status: None, + was_interrupted: false, + spinner_frame: 0, + last_spinner_tick: Instant::now(), + streaming_started: None, + confirmation_pending: false, + } + } + + /// Get the current input text + pub fn input(&self) -> String { + self.textarea.lines().join("\n") + } + + /// Check if input is empty + pub fn input_is_empty(&self) -> bool { + self.textarea.is_empty() + } + + /// Clear the input + pub fn clear_input(&mut self) { + self.textarea = create_textarea(); + } + + /// Convert conversation events to Claude API message format + /// Groups consecutive tool calls, handles role alternation + pub fn events_to_messages(&self) -> Vec<serde_json::Value> { + let mut messages = Vec::new(); + let mut i = 0; + let events = &self.events; + + while i < events.len() { + match &events[i] { + ConversationEvent::UserMessage { content } => { + messages.push(serde_json::json!({ + "role": "user", + "content": content + })); + i += 1; + } + ConversationEvent::Text { content } => { + messages.push(serde_json::json!({ + "role": "assistant", + "content": content + })); + i += 1; + } + ConversationEvent::ToolCall { .. } => { + // Group consecutive tool calls into single assistant message + let mut tool_uses = Vec::new(); + while i < events.len() { + if let ConversationEvent::ToolCall { id, name, input } = &events[i] { + tool_uses.push(serde_json::json!({ + "type": "tool_use", + "id": id, + "name": name, + "input": input + })); + i += 1; + } else { + break; + } + } + messages.push(serde_json::json!({ + "role": "assistant", + "content": tool_uses + })); + } + ConversationEvent::ToolResult { + tool_use_id, + content, + is_error, + } => { + messages.push(serde_json::json!({ + "role": "user", + "content": [{ + "type": "tool_result", + "tool_use_id": tool_use_id, + "content": content, + "is_error": is_error + }] + })); + i += 1; + } + } + } + + messages + } + + // ===== Generation lifecycle methods ===== + + /// Start generating from current input + pub fn start_generating(&mut self) { + // Add user message event + self.events.push(ConversationEvent::UserMessage { + content: self.input(), + }); + + // Clear input, switch mode + self.clear_input(); + self.mode = AppMode::Generating; + } + + /// Generation complete with command (legacy method, kept for compatibility) + pub fn generation_complete( + &mut self, + command: String, + explanation: Option<String>, + dangerous: bool, + warnings: Vec<String>, + ) { + // Add explanation as text event if present + if let Some(ref exp) = explanation { + self.events.push(ConversationEvent::Text { + content: exp.clone(), + }); + } + + // Add tool_call event for suggest_command + let tool_id = format!("gen_{}", uuid::Uuid::new_v4().simple()); + let mut tool_input = serde_json::json!({ + "command": command, + "conversation_only": false, + "confidence": "high" + }); + if let Some(ref exp) = explanation { + tool_input["message"] = serde_json::json!(exp); + } + if dangerous { + tool_input["danger"] = serde_json::json!("high"); + } + if !warnings.is_empty() { + tool_input["warning"] = serde_json::json!(warnings.join("; ")); + } + + self.events.push(ConversationEvent::ToolCall { + id: tool_id, + name: "suggest_command".to_string(), + input: tool_input, + }); + + self.mode = AppMode::Review; + } + + /// Generation error occurred + pub fn generation_error(&mut self, error: String) { + self.error = Some(error); + self.mode = AppMode::Error; + } + + /// Cancel during generation + pub fn cancel_generation(&mut self) { + // Remove the last user message since generation was cancelled + if let Some(ConversationEvent::UserMessage { .. }) = self.events.last() { + self.events.pop(); + } + self.mode = AppMode::Input; + self.clear_input(); + } + + // ===== Streaming lifecycle methods ===== + + /// Start streaming response + pub fn start_streaming(&mut self) { + self.streaming_text.clear(); + self.streaming_status = None; + self.was_interrupted = false; + self.streaming_started = Some(Instant::now()); + self.mode = AppMode::Streaming; + } + + /// Store session ID from server response + pub fn store_session_id(&mut self, session_id: String) { + self.session_id = Some(session_id); + } + + /// Update streaming status from SSE event + pub fn update_streaming_status(&mut self, status: &str) { + self.streaming_status = Some(StreamingStatus::from_status_str(status)); + } + + /// Cancel streaming with context preservation + pub fn cancel_streaming(&mut self) { + // Mark as interrupted + self.was_interrupted = true; + + // Flush partial text with interruption marker if any + // Trim leading whitespace since LLM responses often start with \n\n + let content = std::mem::take(&mut self.streaming_text); + let trimmed = content.trim_start(); + if !trimmed.is_empty() { + let interrupted_text = format!("{trimmed}\n\n[User cancelled this generation]"); + self.events.push(ConversationEvent::Text { + content: interrupted_text, + }); + } + + // Clear status and return to input + self.streaming_status = None; + self.confirmation_pending = false; + self.mode = AppMode::Input; + } + + /// Append text chunk during streaming + /// Trims leading whitespace from the first chunk(s) since LLM responses often start with \n\n + pub fn append_streaming_text(&mut self, chunk: &str) { + if self.streaming_text.is_empty() { + // First chunk(s): trim leading whitespace + let trimmed = chunk.trim_start(); + if !trimmed.is_empty() { + self.streaming_text.push_str(trimmed); + } + } else { + // Subsequent chunks: append as-is + self.streaming_text.push_str(chunk); + } + } + + /// Add a tool call event during streaming + /// Flushes any pending streaming text first to maintain correct event order + /// For suggest_command, also transitions to Review mode since that ends the LLM turn + pub fn add_tool_call(&mut self, id: String, name: String, input: serde_json::Value) { + // Flush streaming text before adding tool call to maintain correct order + let content = std::mem::take(&mut self.streaming_text); + let trimmed = content.trim_start(); + if !trimmed.is_empty() { + self.events.push(ConversationEvent::Text { + content: trimmed.to_string(), + }); + } + + // suggest_command marks the end of the LLM turn - transition to Review + let is_suggest_command = name == "suggest_command"; + + self.events + .push(ConversationEvent::ToolCall { id, name, input }); + + if is_suggest_command { + self.streaming_status = None; + self.streaming_started = None; + self.mode = AppMode::Review; + } + } + + /// Add a tool result event during streaming + pub fn add_tool_result(&mut self, tool_use_id: String, content: String, is_error: bool) { + self.events.push(ConversationEvent::ToolResult { + tool_use_id, + content, + is_error, + }); + } + + /// Finalize streaming - flush accumulated text to event + pub fn finalize_streaming(&mut self) { + // Flush streaming text to a Text event if non-empty + // Trim leading whitespace since LLM responses often start with \n\n + let content = std::mem::take(&mut self.streaming_text); + let trimmed = content.trim_start(); + if !trimmed.is_empty() { + self.events.push(ConversationEvent::Text { + content: trimmed.to_string(), + }); + } + self.streaming_status = None; + self.streaming_started = None; + self.mode = AppMode::Review; + } + + /// Streaming error + pub fn streaming_error(&mut self, error: String) { + // Discard any partial streaming text + self.streaming_text.clear(); + self.streaming_started = None; + self.error = Some(error); + self.mode = AppMode::Error; + } + + // ===== Edit mode and exit methods ===== + + /// Start edit mode for refinement + pub fn start_edit_mode(&mut self) { + self.confirmation_pending = false; + self.clear_input(); + self.mode = AppMode::Input; + } + + /// Exit with action + pub fn exit(&mut self, action: ExitAction) { + self.exit_action = Some(action); + self.should_exit = true; + } + + /// Retry after error + pub fn retry(&mut self) { + self.error = None; + self.mode = AppMode::Generating; + } + + // ===== Utility methods ===== + + /// Advance spinner frame if enough time has passed + /// Called on every event loop tick (50ms), but only advances spinner + /// when the active spinner's interval has elapsed + pub fn tick(&mut self) { + let interval = active_tick_interval(); + if self.last_spinner_tick.elapsed() >= interval { + self.spinner_frame = (self.spinner_frame + 1) % ACTIVE_SPINNER.frame_count(); + self.last_spinner_tick = Instant::now(); + } + } + + /// Get the most recent command from events + pub fn current_command(&self) -> Option<&str> { + self.events.iter().rev().find_map(|e| e.as_command()) + } + + /// Check if the most recent command suggestion is marked dangerous + /// Checks the `danger` field for "high", "medium", or "med" values + pub fn is_current_command_dangerous(&self) -> bool { + self.events + .iter() + .rev() + .find_map(|e| { + if let ConversationEvent::ToolCall { name, input, .. } = e + && name == "suggest_command" + { + let danger_level = input + .get("danger") + .and_then(|v| v.as_str()) + .unwrap_or("low"); + return Some( + danger_level == "high" || danger_level == "medium" || danger_level == "med", + ); + } + None + }) + .unwrap_or(false) + } +} + +impl Default for AppState { + fn default() -> Self { + Self::new() + } +} diff --git a/crates/atuin-ai/src/tui/terminal.rs b/crates/atuin-ai/src/tui/terminal.rs new file mode 100644 index 00000000..f8089323 --- /dev/null +++ b/crates/atuin-ai/src/tui/terminal.rs @@ -0,0 +1,278 @@ +use crossterm::{ + cursor, + terminal::{disable_raw_mode, enable_raw_mode}, +}; +use eyre::{Context, Result, bail}; +use ratatui::{Terminal, TerminalOptions, Viewport, backend::CrosstermBackend, layout::Rect}; +use std::io::{IsTerminal, Stdout, stdout}; + +/// Install a panic hook that ensures the terminal is restored to a usable state +/// even if the application panics. +/// +/// This must be called before creating the TerminalGuard to ensure proper cleanup +/// during panics. The hook will: +/// 1. Disable raw mode (restoring normal terminal behavior) +/// 2. Call the original panic hook to display panic information +/// +/// # Implementation Note +/// This satisfies TUI-07: Terminal remains usable after panic by ensuring +/// disable_raw_mode() is called before the panic message is displayed. +pub fn install_panic_hook() { + let original_hook = std::panic::take_hook(); + std::panic::set_hook(Box::new(move |panic_info| { + // Attempt to restore terminal - ignore errors since we're already panicking + let _ = disable_raw_mode(); + // Call original hook to display panic with backtrace + original_hook(panic_info); + })); +} + +/// Minimum viewport height +const MIN_VIEWPORT_HEIGHT: u16 = 10; + +/// Margin to leave below viewport for shell prompt +const VIEWPORT_BOTTOM_MARGIN: u16 = 2; + +/// Guards terminal lifecycle, ensuring proper setup and cleanup. +/// +/// # Lifecycle +/// - **Setup** (`new()`): Captures cursor position, enables raw mode, creates inline viewport +/// - **Cleanup** (`Drop`): Clears terminal, disables raw mode +/// +/// # Dynamic Viewport Sizing +/// The viewport starts at 15 lines (enough for simple commands) and grows +/// dynamically when content requires more space. Use `ensure_height()` before +/// rendering to grow the viewport if needed. +/// +/// # Safety Features +/// - Non-TTY detection: Returns error early if stdout is not a terminal +/// - Panic recovery: Works with `install_panic_hook()` to restore terminal after panic +/// - Drop-based cleanup: Ensures terminal is restored on normal exit +/// +/// # Example +/// ```no_run +/// use atuin_ai::tui::{install_panic_hook, TerminalGuard}; +/// +/// install_panic_hook(); // Once at program start +/// let mut guard = TerminalGuard::new(true)?; +/// let terminal = guard.terminal(); +/// // ... use terminal ... +/// // Drop automatically cleans up +/// # Ok::<(), eyre::Report>(()) +/// ``` +pub struct TerminalGuard { + terminal: Terminal<CrosstermBackend<Stdout>>, + anchor_col: u16, + keep_output: bool, + viewport_height: u16, + popup_mode: bool, +} + +impl TerminalGuard { + /// Create a new TerminalGuard, initializing the terminal for inline TUI mode. + /// + /// # Arguments + /// * `keep_output` - If true, preserve TUI output on exit; if false, clear it + /// + /// # Process + /// 1. Check if stdout is a terminal (non-TTY detection) + /// 2. Capture cursor position for inline rendering anchor + /// 3. Enable raw mode for keyboard input + /// 4. Create terminal with inline viewport + /// + /// # Errors + /// - Returns error if stdout is not a terminal (e.g., piped or redirected) + /// - Returns error if terminal initialization fails + /// + /// # Implementation Note + /// Cursor position is captured BEFORE enabling raw mode because some terminals + /// may report position differently after raw mode is enabled. + pub fn new(keep_output: bool) -> Result<Self> { + // Non-TTY check: fail early if stdout is not a terminal + if !stdout().is_terminal() { + bail!( + "atuin-ai requires a terminal (TTY) but stdout is not a terminal. \ + This typically happens when output is piped or redirected." + ); + } + + // Get terminal size and calculate viewport height + let (_, term_height) = crossterm::terminal::size().unwrap_or((80, 24)); + let viewport_height = term_height + .saturating_sub(VIEWPORT_BOTTOM_MARGIN) + .max(MIN_VIEWPORT_HEIGHT); + + // Capture cursor position BEFORE raw mode for accurate anchor + let anchor_col = cursor::position().map(|(x, _)| x).unwrap_or(0); + + // Enable raw mode for keyboard input + enable_raw_mode().context("failed to enable raw mode")?; + + // Create terminal with fixed viewport based on terminal size + let backend = CrosstermBackend::new(stdout()); + let terminal = Terminal::with_options( + backend, + TerminalOptions { + viewport: Viewport::Inline(viewport_height), + }, + ) + .context("failed to create terminal with inline viewport")?; + + Ok(Self { + terminal, + anchor_col, + keep_output, + viewport_height, + popup_mode: false, + }) + } + + /// Create a new TerminalGuard for popup overlay mode. + /// + /// In popup mode: + /// - Raw mode is not managed (atuin-hex owns it) + /// - The viewport is a fixed rect positioned over existing terminal content + /// - The popup area is pre-cleared to prevent background bleed-through + /// - Drop does not clear the viewport or disable raw mode + pub fn new_popup(popup_rect: Rect, anchor_col: u16) -> Result<Self> { + // Pre-clear the popup area before creating the ratatui terminal. + // Ratatui's diff-based rendering won't write "default" (space) cells on + // the first frame because its previous buffer is also all-default. By + // writing spaces to the terminal now, we ensure those positions are + // visually blank even if ratatui skips them. + { + use crossterm::cursor::MoveTo; + use crossterm::execute; + use crossterm::style::{Attribute, SetAttribute}; + use std::io::Write; + + let mut out = stdout(); + for row in popup_rect.y..popup_rect.y.saturating_add(popup_rect.height) { + let _ = execute!( + out, + MoveTo(popup_rect.x, row), + SetAttribute(Attribute::Reset) + ); + let _ = write!(out, "{:width$}", "", width = popup_rect.width as usize); + } + let _ = out.flush(); + } + + let backend = CrosstermBackend::new(stdout()); + let terminal = Terminal::with_options( + backend, + TerminalOptions { + viewport: Viewport::Fixed(popup_rect), + }, + ) + .context("failed to create terminal with fixed viewport")?; + + Ok(Self { + terminal, + anchor_col, + keep_output: false, + viewport_height: popup_rect.height, + popup_mode: true, + }) + } + + /// Returns the current viewport height. + /// + /// The viewport is fixed at creation time based on terminal size. + /// Content that exceeds this height will be scrolled automatically. + /// + /// The `_needed` parameter is kept for API compatibility but ignored - + /// we no longer attempt to resize the viewport dynamically since that + /// operation can fail unpredictably with inline viewports. + pub fn ensure_height(&mut self, _needed: u16) -> Result<u16> { + Ok(self.viewport_height) + } + + /// Get the current viewport height. + pub fn viewport_height(&self) -> u16 { + self.viewport_height + } + + /// Get mutable reference to the underlying terminal. + /// + /// Use this to perform rendering operations. + pub fn terminal(&mut self) -> &mut Terminal<CrosstermBackend<Stdout>> { + &mut self.terminal + } + + /// Resize the popup viewport to a new rect. + /// + /// Creates a fresh terminal with the updated Fixed viewport. The caller + /// is responsible for pre-clearing any newly exposed rows before calling + /// this (see `PopupState::grow_to`). + pub fn resize_popup(&mut self, new_rect: Rect) -> Result<()> { + self.viewport_height = new_rect.height; + let backend = CrosstermBackend::new(stdout()); + self.terminal = Terminal::with_options( + backend, + TerminalOptions { + viewport: Viewport::Fixed(new_rect), + }, + ) + .context("failed to resize popup terminal")?; + Ok(()) + } + + /// Get the anchor column where the inline UI should be positioned. + /// + /// This is the column position where the cursor was located when + /// the terminal was initialized. + pub fn anchor_col(&self) -> u16 { + self.anchor_col + } +} + +/// Cleanup terminal state when TerminalGuard is dropped. +/// +/// This implements TUI-08: Terminal restores correctly after normal exit. +/// +/// # Cleanup Process +/// 1. Conditionally clear terminal content (based on keep_output flag) +/// 2. Disable raw mode (restore normal terminal behavior) +/// +/// # Error Handling +/// Errors are intentionally ignored during cleanup since: +/// - We're already exiting and can't meaningfully handle errors +/// - Best-effort restoration is better than panicking during Drop +/// - The panic hook provides a second layer of safety for abnormal exits +impl Drop for TerminalGuard { + fn drop(&mut self) { + if self.popup_mode { + // Popup mode: screen restoration handled by caller before drop. + // Raw mode is owned by atuin-hex, don't touch it. + return; + } + + // Clear terminal content only if keep_output is false - ignore errors (best-effort) + if !self.keep_output { + let _ = self.terminal.clear(); + } + + // Disable raw mode to restore normal terminal behavior - ignore errors + let _ = disable_raw_mode(); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_panic_hook_installation() { + // Test that panic hook can be installed without error + install_panic_hook(); + // Installing again should work (replaces previous hook) + install_panic_hook(); + } + + // Note: Cannot easily test TerminalGuard::new() in CI since it requires a TTY. + // Manual testing required for: + // 1. Non-TTY detection: echo "" | cargo run -p atuin-ai -- inline + // 2. Drop cleanup: Run inline command, press Esc, verify terminal is normal + // 3. Panic recovery: Add panic!("test") after TerminalGuard::new(), verify terminal is usable +} diff --git a/crates/atuin-ai/src/tui/view_model.rs b/crates/atuin-ai/src/tui/view_model.rs new file mode 100644 index 00000000..0a296065 --- /dev/null +++ b/crates/atuin-ai/src/tui/view_model.rs @@ -0,0 +1,413 @@ +//! View model types for the TUI application +//! +//! This module contains the view model types that represent the rendering +//! specification. These types are derived from the domain state (conversation +//! events) via the `Blocks::from_state()` function. + +use super::state::{AppMode, AppState, ConversationEvent}; + +/// Warning classification for command suggestions +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum WarningKind { + /// Dangerous command (! indicator, AlertError color) + Danger, + /// Low confidence answer (? indicator, AlertWarn color) + LowConfidence, +} + +/// Content variants for blocks - each variant is fully self-describing +#[derive(Debug, Clone)] +pub enum Content { + Input { + text: String, + active: bool, + cursor_pos: usize, + }, + /// Command suggestion (from suggest_command tool call) + Command { + text: String, + faded: bool, // Phase 5 feature + }, + Text { + markdown: String, + }, + Error { + message: String, + }, + /// Warning for dangerous or low-confidence commands + Warning { + kind: WarningKind, + text: String, + pending_confirm: bool, // true when awaiting second Enter + }, + Spinner { + frame: usize, // 0-3 for animation + status_text: String, // Status-based text (Processing..., Thinking..., etc.) + }, + /// Tool call status display (in-flight or completed summary) + ToolStatus { + /// Number of non-suggest_command tools completed + completed_count: usize, + /// Current in-flight tool description (None if all done) + current_label: Option<String>, + /// Spinner frame for in-flight display + frame: usize, + }, +} + +impl Content { + /// Get the prefix symbol for this content type + pub fn prefix_symbol(&self) -> &'static str { + match self { + Content::Input { .. } => ">", + Content::Command { .. } => "$", + Content::Text { .. } => " ", + Content::Error { .. } => "!", + Content::Warning { kind, .. } => match kind { + WarningKind::Danger => "!", + WarningKind::LowConfidence => "?", + }, + Content::Spinner { .. } => "/", + Content::ToolStatus { current_label, .. } => { + if current_label.is_some() { + "/" + } else { + "\u{2713}" + } // spinner or checkmark + } + } + } +} + +/// A visual block in the UI +#[derive(Debug, Clone)] +pub struct Block { + pub content: Vec<Content>, + pub separator_above: bool, + pub title: Option<String>, +} + +/// Status bar content shown on the bottom border during processing +#[derive(Debug, Clone)] +pub struct StatusBar { + /// Spinner animation frame + pub frame: usize, + /// Status text to display (e.g., "Thinking...", "run_bash (used 2 tools)") + pub text: String, +} + +/// Complete view model - the rendering specification +#[derive(Debug, Clone)] +pub struct Blocks { + pub items: Vec<Block>, + pub footer: &'static str, + /// Transient status shown on bottom border during streaming/generating + pub status_bar: Option<StatusBar>, +} + +/// Count non-suggest_command tool calls since the last user message +fn count_tool_calls_since_last_user(events: &[ConversationEvent]) -> (usize, Option<String>) { + let last_user_idx = events + .iter() + .rposition(|e| matches!(e, ConversationEvent::UserMessage { .. })) + .unwrap_or(0); + + let mut completed = 0; + let mut in_flight: Option<String> = None; + + for event in &events[last_user_idx..] { + match event { + ConversationEvent::ToolCall { name, .. } if name != "suggest_command" => { + // New tool call starts as in-flight + if in_flight.is_some() { + // Previous tool is now completed + completed += 1; + } + in_flight = Some(name.clone()); + } + ConversationEvent::ToolResult { .. } => { + // Tool completed + if in_flight.is_some() { + completed += 1; + in_flight = None; + } + } + _ => {} + } + } + + (completed, in_flight) +} + +/// Check if any turn in the conversation has a command +fn has_any_command(events: &[ConversationEvent]) -> bool { + events.iter().any(|e| { + if let ConversationEvent::ToolCall { name, input, .. } = e { + name == "suggest_command" && input.get("command").and_then(|v| v.as_str()).is_some() + } else { + false + } + }) +} + +impl Blocks { + /// Pure function: derive the complete view model from state + /// + /// Iterates through conversation events and builds visual blocks. + /// Also handles streaming text and mode-dependent UI. + pub fn from_state(state: &AppState) -> Self { + let mut items = Vec::new(); + let mut status_bar = None; + + // 1. Build blocks from conversation events + for event in &state.events { + match event { + ConversationEvent::UserMessage { content } => { + items.push(Block { + content: vec![Content::Input { + text: content.clone(), + active: false, + cursor_pos: 0, + }], + separator_above: false, + title: None, + }); + } + ConversationEvent::Text { content } => { + // In Review mode with completed tool calls, prepend ToolStatus to this Text block + let (completed, _) = count_tool_calls_since_last_user(&state.events); + let mut block_content = Vec::new(); + + if state.mode == AppMode::Review && completed > 0 { + block_content.push(Content::ToolStatus { + completed_count: completed, + current_label: None, + frame: 0, + }); + } + + block_content.push(Content::Text { + markdown: content.clone(), + }); + + items.push(Block { + content: block_content, + separator_above: false, + title: None, + }); + } + ConversationEvent::ToolCall { name, input, .. } => { + // Only render suggest_command tool calls with a command + if name == "suggest_command" { + let command = input.get("command").and_then(|v| v.as_str()); + + // Build block content - only render if command is present + // When command is null, this is a conversation-only turn and the + // response text comes via a separate Text event + let mut block_content = Vec::new(); + + if let Some(cmd) = command { + block_content.push(Content::Command { + text: cmd.to_string(), + faded: false, + }); + } + + // Extract warning data from tool call input + // danger: "high" | "medium" | "med" | "low" - high/medium/med trigger warning + let danger_level = input + .get("danger") + .and_then(|v| v.as_str()) + .unwrap_or("low"); + let is_dangerous = danger_level == "high" + || danger_level == "medium" + || danger_level == "med"; + let danger_notes = input.get("danger_notes").and_then(|v| v.as_str()); + + // confidence: "high" | "medium" | "low" - low triggers warning + let confidence_level = input + .get("confidence") + .and_then(|v| v.as_str()) + .unwrap_or("high"); + let is_low_confidence = confidence_level == "low"; + let confidence_notes = + input.get("confidence_notes").and_then(|v| v.as_str()); + + // Add warning content if applicable (danger takes precedence) + if is_dangerous { + if let Some(notes) = danger_notes { + block_content.push(Content::Warning { + kind: WarningKind::Danger, + text: notes.to_string(), + pending_confirm: state.confirmation_pending, + }); + } + } else if is_low_confidence && let Some(notes) = confidence_notes { + block_content.push(Content::Warning { + kind: WarningKind::LowConfidence, + text: notes.to_string(), + pending_confirm: false, // low confidence doesn't require confirm + }); + } + + // Only add block if there's content + if !block_content.is_empty() { + items.push(Block { + content: block_content, + separator_above: false, + title: None, + }); + } + } + // Other tool calls are not rendered (internal protocol) + } + ConversationEvent::ToolResult { .. } => { + // Tool results are not rendered (internal protocol) + } + } + } + + // 2. AI response block (streaming text only) - shown during Streaming only + // Transient status (spinner, tool progress) goes to status_bar on the bottom border. + // In Review mode, ToolStatus is handled inline with ConversationEvent::Text above. + if state.mode == AppMode::Streaming { + let (completed, in_flight) = count_tool_calls_since_last_user(&state.events); + + // Tool status -> status bar + if let Some(ref label) = in_flight { + let text = if completed > 0 { + format!( + "{} (used {} tool{})", + label, + completed, + if completed == 1 { "" } else { "s" } + ) + } else { + label.clone() + }; + status_bar = Some(StatusBar { + frame: state.spinner_frame, + text, + }); + } + + // Spinner -> status bar (only when no text yet and no tool in-flight) + if state.streaming_text.is_empty() { + let should_show_spinner = state.streaming_status.is_some() + || state + .streaming_started + .map(|start| start.elapsed() >= std::time::Duration::from_millis(200)) + .unwrap_or(true); + + if should_show_spinner && in_flight.is_none() { + let status_text = state + .streaming_status + .as_ref() + .map(|s| s.display_text().to_string()) + .unwrap_or_else(|| "Generating...".to_string()); + + status_bar = Some(StatusBar { + frame: state.spinner_frame, + text: status_text, + }); + } + } else { + // Show streaming text as content + items.push(Block { + content: vec![Content::Text { + markdown: state.streaming_text.clone(), + }], + separator_above: false, + title: None, + }); + } + } + + // 3. Mode-dependent UI + match state.mode { + AppMode::Input => { + // Active input uses TextArea widget, rendered directly + // We add a placeholder block that will be replaced by textarea rendering + items.push(Block { + content: vec![Content::Input { + text: state.input(), + active: true, + cursor_pos: 0, // Not used for active input - textarea handles cursor + }], + separator_above: false, + title: None, + }); + } + AppMode::Generating => { + let status_text = state + .streaming_status + .as_ref() + .map(|s| s.display_text().to_string()) + .unwrap_or_else(|| "Generating...".to_string()); + + status_bar = Some(StatusBar { + frame: state.spinner_frame, + text: status_text, + }); + } + AppMode::Streaming => { + // Handled above in streaming text section + } + AppMode::Review | AppMode::Error => { + // No additional UI elements + } + } + + // 4. Error if present (renders at end) + if let Some(ref err) = state.error { + items.push(Block { + content: vec![Content::Error { + message: err.clone(), + }], + separator_above: false, + title: None, + }); + } + + // 5. Set separator flags (first has no separator) + for (idx, block) in items.iter_mut().enumerate() { + block.separator_above = idx > 0; + } + + // 6. Set title on first block only + if let Some(first) = items.first_mut() { + first.title = Some("Ask questions or generate a command:".to_string()); + } + + // 7. Derive footer from mode and events + let footer = Self::footer_for_mode(&state.mode, &state.events, state.confirmation_pending); + + Self { + items, + footer, + status_bar, + } + } + + /// Derive footer text from current mode and conversation state + fn footer_for_mode( + mode: &AppMode, + events: &[ConversationEvent], + confirmation_pending: bool, + ) -> &'static str { + match mode { + AppMode::Input => "[Enter]: Accept [Esc]: Cancel", + AppMode::Generating | AppMode::Streaming => "[Esc]: Cancel", + AppMode::Review => { + if confirmation_pending { + "[Enter]: Confirm dangerous command [Esc]: Cancel" + } else if has_any_command(events) { + "[Enter]: Run [Tab]: Insert [f]: Follow-up [Esc]: Cancel" + } else { + "[f]: Follow-up [Esc]: Cancel" + } + } + AppMode::Error => "[Enter]/[r]: Retry [Esc]: Cancel", + } + } +} diff --git a/crates/atuin-ai/test-renders.json b/crates/atuin-ai/test-renders.json new file mode 100644 index 00000000..31c180fa --- /dev/null +++ b/crates/atuin-ai/test-renders.json @@ -0,0 +1,295 @@ +[ + { + "name": "01_empty_input", + "description": "Initial state with empty input prompt", + "state": { + "events": [], + "mode": "Input", + "input": "", + "cursor_pos": 0 + } + }, + { + "name": "02_typing_input", + "description": "User typing in input field", + "state": { + "events": [], + "mode": "Input", + "input": "list all files", + "cursor_pos": 14 + } + }, + { + "name": "03_generating_spinner", + "description": "Waiting for API response (spinner)", + "state": { + "events": [ + {"type": "user_message", "content": "list all files"} + ], + "mode": "Generating", + "spinner_frame": 0 + } + }, + { + "name": "04_streaming_text", + "description": "Text streaming in from API", + "state": { + "events": [ + {"type": "user_message", "content": "what is rust?"} + ], + "mode": "Streaming", + "streaming_text": "Rust is a systems programming language focused on safety, speed, and", + "spinner_frame": 2 + } + }, + { + "name": "05_simple_command", + "description": "Simple command suggestion", + "state": { + "events": [ + {"type": "user_message", "content": "list all files"}, + {"type": "tool_call", "id": "1", "name": "suggest_command", "input": {"command": "ls -la"}} + ], + "mode": "Review" + } + }, + { + "name": "06_command_with_long_text", + "description": "Command that wraps to multiple lines", + "state": { + "events": [ + {"type": "user_message", "content": "find large files"}, + {"type": "tool_call", "id": "1", "name": "suggest_command", "input": {"command": "find /home -type f -size +100M -exec ls -lh {} \\; 2>/dev/null | sort -k5 -h"}} + ], + "mode": "Review" + } + }, + { + "name": "07_conversation_only_response", + "description": "Response without command (conversation mode)", + "state": { + "events": [ + {"type": "user_message", "content": "what does the -la flag do?"}, + {"type": "text", "content": "The `-la` flags combine two options:\n\n- `-l` shows long format with permissions, owner, size, and date\n- `-a` shows all files including hidden ones (starting with .)"} + ], + "mode": "Review" + } + }, + { + "name": "08_multi_turn_conversation", + "description": "Multiple turns of conversation", + "state": { + "events": [ + {"type": "user_message", "content": "list all files"}, + {"type": "tool_call", "id": "1", "name": "suggest_command", "input": {"command": "ls -la"}}, + {"type": "user_message", "content": "can you explain those flags?"}, + {"type": "text", "content": "The -l flag shows long format with permissions, -a shows hidden files."} + ], + "mode": "Review" + } + }, + { + "name": "09_tool_call_in_progress", + "description": "Tool being executed (spinner)", + "state": { + "events": [ + {"type": "user_message", "content": "what is the latest version of node?"}, + {"type": "tool_call", "id": "1", "name": "web_search", "input": {"query": "nodejs latest version"}} + ], + "mode": "Streaming", + "streaming_text": "", + "spinner_frame": 1 + } + }, + { + "name": "10_tool_calls_completed_with_text", + "description": "Tools finished, text streaming", + "state": { + "events": [ + {"type": "user_message", "content": "what is the latest version of node?"}, + {"type": "tool_call", "id": "1", "name": "web_search", "input": {"query": "nodejs latest version"}}, + {"type": "tool_result", "tool_use_id": "1", "content": "Node.js v22.0.0"} + ], + "mode": "Streaming", + "streaming_text": "The latest version of Node.js is v22.0.0, released in April 2024.", + "spinner_frame": 0 + } + }, + { + "name": "11_tool_calls_in_review", + "description": "Completed tools shown in review mode", + "state": { + "events": [ + {"type": "user_message", "content": "what is the latest version of node?"}, + {"type": "tool_call", "id": "1", "name": "web_search", "input": {"query": "nodejs latest version"}}, + {"type": "tool_result", "tool_use_id": "1", "content": "Node.js v22.0.0"}, + {"type": "tool_call", "id": "2", "name": "web_fetch", "input": {"url": "https://nodejs.org"}}, + {"type": "tool_result", "tool_use_id": "2", "content": "..."}, + {"type": "text", "content": "The latest version of Node.js is **v22.0.0**, released in April 2024. Key features include:\n\n- Native WebSocket client\n- Improved ES modules support\n- Better performance"} + ], + "mode": "Review" + } + }, + { + "name": "12_error_state", + "description": "Error message displayed", + "state": { + "events": [ + {"type": "user_message", "content": "do something"} + ], + "mode": "Error", + "error": "Failed to connect to API: connection timeout" + } + }, + { + "name": "13_dangerous_command", + "description": "Dangerous command with warning", + "state": { + "events": [ + {"type": "user_message", "content": "delete all files in home"}, + {"type": "tool_call", "id": "1", "name": "suggest_command", "input": { + "command": "rm -rf ~/*", + "dangerous": true, + "warning": "This will permanently delete all files in your home directory including documents, configurations, and SSH keys." + }} + ], + "mode": "Review", + "confirmation_pending": false + } + }, + { + "name": "14_dangerous_command_confirming", + "description": "Dangerous command awaiting second Enter", + "state": { + "events": [ + {"type": "user_message", "content": "delete all files in home"}, + {"type": "tool_call", "id": "1", "name": "suggest_command", "input": { + "command": "rm -rf ~/*", + "dangerous": true, + "warning": "This will permanently delete all files in your home directory." + }} + ], + "mode": "Review", + "confirmation_pending": true + } + }, + { + "name": "15_low_confidence", + "description": "Low confidence command with warning", + "state": { + "events": [ + {"type": "user_message", "content": "do that thing with the files"}, + {"type": "tool_call", "id": "1", "name": "suggest_command", "input": { + "command": "ls -la", + "confidence": "low", + "warning": "I'm not entirely sure what you mean by 'that thing'. This lists files - is that what you wanted?" + }} + ], + "mode": "Review" + } + }, + { + "name": "16_long_user_input", + "description": "User input that wraps", + "state": { + "events": [ + {"type": "user_message", "content": "I need a command that will find all JavaScript files in my project, excluding node_modules, and count the total lines of code"} + ], + "mode": "Generating", + "spinner_frame": 0 + } + }, + { + "name": "17_long_text_response", + "description": "Long text response that wraps multiple times", + "state": { + "events": [ + {"type": "user_message", "content": "explain git"}, + {"type": "text", "content": "Git is a distributed version control system created by Linus Torvalds in 2005. It tracks changes to files and enables collaboration between developers. Key concepts include:\n\n- **Repository**: A directory containing your project and its history\n- **Commit**: A snapshot of your changes with a message\n- **Branch**: An independent line of development\n- **Merge**: Combining changes from different branches\n- **Remote**: A version of your repository hosted elsewhere (like GitHub)"} + ], + "mode": "Review" + } + }, + { + "name": "18_streaming_with_tool_in_progress", + "description": "Tool in progress while streaming", + "state": { + "events": [ + {"type": "user_message", "content": "search for rust async patterns"}, + {"type": "text", "content": "Let me search for that..."}, + {"type": "tool_call", "id": "1", "name": "web_search", "input": {"query": "rust async patterns"}} + ], + "mode": "Streaming", + "streaming_text": "", + "spinner_frame": 2 + } + }, + { + "name": "19_multiple_commands_in_conversation", + "description": "Multiple command suggestions across turns", + "state": { + "events": [ + {"type": "user_message", "content": "create a new directory called test"}, + {"type": "tool_call", "id": "1", "name": "suggest_command", "input": {"command": "mkdir test"}}, + {"type": "user_message", "content": "now cd into it"}, + {"type": "tool_call", "id": "2", "name": "suggest_command", "input": {"command": "cd test"}}, + {"type": "user_message", "content": "create a file"}, + {"type": "tool_call", "id": "3", "name": "suggest_command", "input": {"command": "touch file.txt"}} + ], + "mode": "Review" + } + }, + { + "name": "20_empty_command_with_description", + "description": "Tool call with null command (conversation only)", + "state": { + "events": [ + {"type": "user_message", "content": "what's the weather like?"}, + {"type": "tool_call", "id": "1", "name": "suggest_command", "input": { + "command": null, + "description": "I can't check the weather directly, but you could use: curl wttr.in" + }} + ], + "mode": "Review" + } + }, + { + "name": "21_status_processing", + "description": "Streaming with Processing status", + "state": { + "events": [ + {"type": "user_message", "content": "analyze this code"} + ], + "mode": "Streaming", + "streaming_text": "", + "streaming_status": "Processing", + "spinner_frame": 0 + } + }, + { + "name": "22_status_thinking", + "description": "Streaming with Thinking status", + "state": { + "events": [ + {"type": "user_message", "content": "how do I optimize this query?"} + ], + "mode": "Streaming", + "streaming_text": "", + "streaming_status": "Thinking", + "spinner_frame": 1 + } + }, + { + "name": "23_follow_up_input", + "description": "Follow-up input after command", + "state": { + "events": [ + {"type": "user_message", "content": "list files"}, + {"type": "tool_call", "id": "1", "name": "suggest_command", "input": {"command": "ls -la"}} + ], + "mode": "Input", + "input": "but only show directories", + "cursor_pos": 24 + } + } +] diff --git a/crates/atuin-client/Cargo.toml b/crates/atuin-client/Cargo.toml new file mode 100644 index 00000000..c4f5679c --- /dev/null +++ b/crates/atuin-client/Cargo.toml @@ -0,0 +1,84 @@ +[package] +name = "atuin-client" +edition = "2024" +description = "client library for atuin" + +rust-version = { workspace = true } +version = { workspace = true } +authors = { workspace = true } +license = { workspace = true } +homepage = { workspace = true } +repository = { workspace = true } + +# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html + +[features] +default = ["sync", "hub", "daemon"] +sync = ["urlencoding", "reqwest", "sha2", "hex"] +hub = ["reqwest"] +daemon = [] +check-update = [] + +[dependencies] +atuin-common = { path = "../atuin-common", version = "18.13.2" } + +log = { workspace = true } +base64 = { workspace = true } +time = { workspace = true, features = ["macros", "formatting", "parsing"] } +clap = { workspace = true } +eyre = { workspace = true } +directories = { workspace = true } +uuid = { workspace = true } +whoami = { workspace = true } +interim = { workspace = true } +config = { workspace = true } +serde = { workspace = true } +serde_json = { workspace = true } +humantime = "2.1.0" +async-trait = { workspace = true } +itertools = { workspace = true } +rand = { workspace = true } +shellexpand = "3" +sqlx = { workspace = true, features = ["sqlite", "regexp"] } +minspan = "0.1.5" +regex = "1.10.5" +serde_regex = "1.1.0" +fs-err = { workspace = true } +sql-builder = { workspace = true } +memchr = "2.7" +rmp = { version = "0.8.14" } +typed-builder = { workspace = true } +tokio = { workspace = true } +semver = { workspace = true } +thiserror = { workspace = true } +futures = "0.3" +notify = "7" +crypto_secretbox = "0.1.1" +generic-array = { version = "0.14", features = ["serde"] } +serde_with = "3.8.1" + +# encryption +rusty_paseto = { version = "0.8.0", default-features = false } +rusty_paserk = { version = "0.5.0", default-features = false, features = [ + "v4", + "serde", +] } + +# sync +urlencoding = { version = "2.1.0", optional = true } +reqwest = { workspace = true, optional = true } +hex = { version = "0.4", optional = true } +sha2 = { version = "0.10", optional = true } +indicatif = "0.18.0" +tiny-bip39 = "2.0.0" + +# theme +crossterm = { workspace = true, features = ["serde"] } +palette = { version = "0.7.5", features = ["serializing"] } +strum_macros = "0.27" +strum = { version = "0.27", features = ["strum_macros"] } + +[dev-dependencies] +tokio = { version = "1", features = ["full"] } +pretty_assertions = { workspace = true } +testing_logger = "0.1.1" diff --git a/crates/atuin-client/config.toml b/crates/atuin-client/config.toml new file mode 100644 index 00000000..6e67a4e1 --- /dev/null +++ b/crates/atuin-client/config.toml @@ -0,0 +1,366 @@ +## Base directory for Atuin data files (databases, keys, session, etc.) +## All data file paths default to being relative to this directory. +## linux/mac: ~/.local/share/atuin (or XDG_DATA_HOME/atuin) +## windows: %USERPROFILE%/.local/share/atuin +# data_dir = "~/.local/share/atuin" + +## where to store your database, default is your system data directory +## linux/mac: ~/.local/share/atuin/history.db +## windows: %USERPROFILE%/.local/share/atuin/history.db +# db_path = "~/.history.db" + +## where to store your encryption key, default is your system data directory +## linux/mac: ~/.local/share/atuin/key +## windows: %USERPROFILE%/.local/share/atuin/key +# key_path = "~/.key" + +## where to store your auth session token, default is your system data directory +## linux/mac: ~/.local/share/atuin/session +## windows: %USERPROFILE%/.local/share/atuin/session +# session_path = "~/.session" + +## date format used, either "us" or "uk" +# dialect = "us" + +## default timezone to use when displaying time +## either "l", "local" to use the system's current local timezone, or an offset +## from UTC in the format of "<+|->H[H][:M[M][:S[S]]]" +## for example: "+9", "-05", "+03:30", "-01:23:45", etc. +# timezone = "local" + +## enable or disable automatic sync +# auto_sync = true + +## enable or disable automatic update checks +# update_check = true + +## address of the sync server +# sync_address = "https://api.atuin.sh" + +## how often to sync history. note that this is only triggered when a command +## is ran, so sync intervals may well be longer +## set it to 0 to sync after every command +# sync_frequency = "10m" + +## which search mode to use +## possible values: prefix, fulltext, fuzzy, skim +# search_mode = "fuzzy" + +## which filter mode to use by default +## possible values: "global", "host", "session", "session-preload", "directory", "workspace" +## consider using search.filters to customize the enablement and order of filter modes +# filter_mode = "global" + +## With workspace filtering enabled, Atuin will filter for commands executed +## in any directory within a git repository tree (default: false). +## +## To use workspace mode by default when available, set this to true and +## set filter_mode to "workspace" or leave it unspecified and +## set search.filters to include "workspace" before other filter modes. +# workspaces = false + +## which filter mode to use when atuin is invoked from a shell up-key binding +## the accepted values are identical to those of "filter_mode" +## leave unspecified to use same mode set in "filter_mode" +# filter_mode_shell_up_key_binding = "global" + +## which search mode to use when atuin is invoked from a shell up-key binding +## the accepted values are identical to those of "search_mode" +## leave unspecified to use same mode set in "search_mode" +# search_mode_shell_up_key_binding = "fuzzy" + +## which style to use +## possible values: auto, full, compact +# style = "auto" + +## the maximum number of lines the interface should take up +## set it to 0 to always go full screen +# inline_height = 0 + +## the maximum number of lines the interface should take up +## when atuin is invoked from a shell up-key binding +## the accepted values are identical to those of "inline_height" +# inline_height_shell_up_key_binding = 0 + +## Invert the UI - put the search bar at the top , Default to `false` +# invert = false + +## enable or disable showing a preview of the selected command +## useful when the command is longer than the terminal width and is cut off +# show_preview = true + +## what to do when the escape key is pressed when searching +## possible values: return-original, return-query +# exit_mode = "return-original" + +## possible values: emacs, subl +# word_jump_mode = "emacs" + +## characters that count as a part of a word +# word_chars = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789" + +## number of context lines to show when scrolling by pages +# scroll_context_lines = 1 + +## use ctrl instead of alt as the shortcut modifier key for numerical UI shortcuts +## alt-0 .. alt-9 +# ctrl_n_shortcuts = false + +## Show numeric shortcuts (1..9) beside list items in the TUI +## set to false to hide the moving numbers if you find them distracting +# show_numeric_shortcuts = true + +## default history list format - can also be specified with the --format arg +# history_format = "{time}\t{command}\t{duration}" + +## prevent commands matching any of these regexes from being written to history. +## Note that these regular expressions are unanchored, i.e. if they don't start +## with ^ or end with $, they'll match anywhere in the command. +## For details on the supported regular expression syntax, see +## https://docs.rs/regex/latest/regex/#syntax +# history_filter = [ +# "^secret-cmd", +# "^innocuous-cmd .*--secret=.+", +# ] + +## prevent commands run with cwd matching any of these regexes from being written +## to history. Note that these regular expressions are unanchored, i.e. if they don't +## start with ^ or end with $, they'll match anywhere in CWD. +## For details on the supported regular expression syntax, see +## https://docs.rs/regex/latest/regex/#syntax +# cwd_filter = [ +# "^/very/secret/area", +# ] + +## Configure the maximum height of the preview to show. +## Useful when you have long scripts in your history that you want to distinguish +## by more than the first few lines. +# max_preview_height = 4 + +## Configure whether or not to show the help row, which includes the current Atuin +## version (and whether an update is available), a keymap hint, and the total +## amount of commands in your history. +# show_help = true + +## Configure whether or not to show tabs for search and inspect +# show_tabs = true + +## Configure whether or not the tabs row may be auto-hidden, which includes the current Atuin +## tab, such as Search or Inspector, and other tabs you may wish to see. This will +## only be hidden if there are fewer than this count of lines available, and does not affect the use +## of keyboard shortcuts to switch tab. 0 to never auto-hide, default is 8 (lines). +## This is ignored except in `compact` mode. +# auto_hide_height = 8 + +## Defaults to true. This matches history against a set of default regex, and will not save it if we get a match. Defaults include +## 1. AWS key id +## 2. Github pat (old and new) +## 3. Slack oauth tokens (bot, user) +## 4. Slack webhooks +## 5. Stripe live/test keys +# secrets_filter = true + +## Defaults to true. If enabled, upon hitting enter Atuin will immediately execute the command, +## whereas tab will put the command in the prompt for editing. +## If set to false, both enter and tab will place the command in the prompt for editing. +## This applies for new installs. Old installs will keep the old behaviour unless configured otherwise. +enter_accept = true + +## Defaults to false. If enabled, when triggered after &&, || or |, Atuin will complete commands to chain rather than replace the current line. +# command_chaining = false + +## Defaults to "emacs". This specifies the keymap on the startup of `atuin +## search`. If this is set to "auto", the startup keymap mode in the Atuin +## search is automatically selected based on the shell's keymap where the +## keybinding is defined. If this is set to "emacs", "vim-insert", or +## "vim-normal", the startup keymap mode in the Atuin search is forced to be +## the specified one. +# keymap_mode = "auto" + +## Cursor style in each keymap mode. If specified, the cursor style is changed +## in entering the cursor shape. Available values are "default" and +## "{blink,steady}-{block,underline,bar}". +# keymap_cursor = { emacs = "blink-block", vim_insert = "blink-block", vim_normal = "steady-block" } + +# network_connect_timeout = 5 +# network_timeout = 5 + +## Timeout (in seconds) for acquiring a local database connection (sqlite) +# local_timeout = 5 + +## Set this to true and Atuin will minimize motion in the UI - timers will not update live, etc. +## Alternatively, set env NO_MOTION=true +# prefers_reduced_motion = false + +[stats] +## Set commands where we should consider the subcommand for statistics. Eg, kubectl get vs just kubectl +# common_subcommands = [ +# "apt", +# "cargo", +# "composer", +# "dnf", +# "docker", +# "dotnet", +# "git", +# "go", +# "ip", +# "jj", +# "kubectl", +# "nix", +# "nmcli", +# "npm", +# "pecl", +# "pnpm", +# "podman", +# "port", +# "systemctl", +# "tmux", +# "yarn", +# ] + +## Set commands that should be totally stripped and ignored from stats +# common_prefix = ["sudo"] + +## Set commands that will be completely ignored from stats +# ignored_commands = [ +# "cd", +# "ls", +# "vi" +# ] + +[keys] +# Defaults to true. If disabled, using the up/down key won't exit the TUI when scrolled past the first/last entry. +# scroll_exits = true + +# Defaults to true. The left arrow key will exit the TUI when scrolling before the first character +# exit_past_line_start = true + +# Defaults to true. The right arrow key performs the same functionality as Tab and copies the selected line to the command line to be modified. +# accept_past_line_end = true + +# Defaults to false. The left arrow key performs the same functionality as Tab and copies the selected line to the command line to be modified. +# accept_past_line_start = false + +# Defaults to false. The backspace key performs the same functionality as Tab and copies the selected line to the command line to be modified when at the start of the line. +# accept_with_backspace = false + +[sync] +# Enable sync v2 by default +# This ensures that sync v2 is enabled for new installs only +# In a later release it will become the default across the board +records = true + +[preview] +## which preview strategy to use to calculate the preview height (respects max_preview_height). +## possible values: auto, static +## auto: length of the selected command. +## static: length of the longest command stored in the history. +## fixed: use max_preview_height as fixed height. +# strategy = "auto" + +[daemon] +## Enables using the daemon to sync. +# enabled = false + +## Automatically start and manage the daemon when needed. +## Not compatible with `systemd_socket = true`. +# autostart = false + +## How often the daemon should sync in seconds +# sync_frequency = 300 + +## The path to the unix socket used by the daemon (on unix systems) +## linux/mac: ~/.local/share/atuin/atuin.sock +## windows: Not Supported +# socket_path = "~/.local/share/atuin/atuin.sock" + +## The daemon pidfile used for lifecycle management. +## Defaults to the Atuin data directory. +# pidfile_path = "~/.local/share/atuin/atuin-daemon.pid" + +## Use systemd socket activation rather than opening the given path (the path must still be correct for the client) +## linux: false +## mac/windows: Not Supported +# systemd_socket = false + +## The port that should be used for TCP on non unix systems +# tcp_port = 8889 + +# [theme] +## Color theme to use for rendering in the terminal. +## There are some built-in themes, including the base theme ("default"), +## "autumn" and "marine". You can add your own themes to the "./themes" subdirectory of your +## Atuin config (or ATUIN_THEME_DIR, if provided) as TOML files whose keys should be one or +## more of AlertInfo, AlertWarn, AlertError, Annotation, Base, Guidance, Important, and +## the string values as lowercase entries from this list: +## https://ogeon.github.io/docs/palette/master/palette/named/index.html +## If you provide a custom theme file, it should be called "NAME.toml" and the theme below +## should be the stem, i.e. `theme = "NAME"` for your chosen NAME. +# name = "autumn" + +## Whether the theme manager should output normal or extra information to help fix themes. +## Boolean, true or false. If unset, left up to the theme manager. +# debug = true + +[search] +## The list of enabled filter modes, in order of priority. +## The "workspace" mode is skipped when not in a workspace or workspaces = false. +## Default filter mode can be overridden with the filter_mode setting. +# filters = [ "global", "host", "session", "session-preload", "workspace", "directory" ] + +[tmux] +## Enable using atuin with tmux popup (requires tmux >= 3.2) +## When enabled and running inside tmux, Atuin will use a popup window for interactive search. +## Set to false to disable the popup. +## This can also be controlled with the ATUIN_TMUX_POPUP environment variable. +## Note: The tmux popup is currently supported in zsh, bash, and fish shells. This currently doesn't work with iTerm native tmux integration. +# enabled = false + +## Width of the tmux popup window +## Can be a percentage, or integer (e.g. "100" means 100 characters wide) +# width = "80%" + +## Height of the tmux popup window +## Can be a percentage, or integer (e.g. "100" means 100 lines tall) +# height = "60%" + +[ui] +## Columns to display in the interactive search, from left to right. +## The selection indicator (" > ") is always shown first implicitly. +## +## Each column can be specified as a simple string (uses default width) +## or as an object with type, width, and expand: +## { type = "directory", width = 30, expand = true } +## +## Available column types (with default widths): +## duration (5) - Command execution duration (e.g., "123ms") +## time (8) - Relative time since execution (e.g., "59m ago") +## datetime (16) - Absolute timestamp (e.g., "2025-01-22 14:35") +## directory (20) - Working directory (truncated if too long) +## host (15) - Hostname where command was run +## user (10) - Username +## exit (3) - Exit code (colored by success/failure) +## command (*) - The command itself (expands by default) +## +## The "expand" option (default: true for command, false for others) makes a +## column fill remaining space. Only one column should have expand = true. +## +## Default: +# columns = ["duration", "time", "command"] +## +## Examples: +## +## Minimal - more space for commands: +# columns = ["duration", "command"] +## +## With wider directory column: +# columns = ["duration", { type = "directory", width = 30 }, "command"] +## +## Show host for multi-machine sync users: +# columns = ["duration", "time", "host", "command"] +## +## Show exit codes prominently: +# columns = ["exit", "duration", "command"] +## +## Make directory expand instead of command: +# columns = ["duration", "time", { type = "directory", expand = true }, { type = "command", expand = false }] diff --git a/crates/atuin-client/meta-migrations/20260203030924_create_meta.sql b/crates/atuin-client/meta-migrations/20260203030924_create_meta.sql new file mode 100644 index 00000000..26c3c142 --- /dev/null +++ b/crates/atuin-client/meta-migrations/20260203030924_create_meta.sql @@ -0,0 +1,5 @@ +create table if not exists meta ( + key text not null primary key, + value text not null, + updated_at integer not null default (strftime('%s', 'now')) +); diff --git a/crates/atuin-client/migrations/20210422143411_create_history.sql b/crates/atuin-client/migrations/20210422143411_create_history.sql new file mode 100644 index 00000000..1f3f8686 --- /dev/null +++ b/crates/atuin-client/migrations/20210422143411_create_history.sql @@ -0,0 +1,16 @@ +-- Add migration script here +create table if not exists history ( + id text primary key, + timestamp integer not null, + duration integer not null, + exit integer not null, + command text not null, + cwd text not null, + session text not null, + hostname text not null, + + unique(timestamp, cwd, command) +); + +create index if not exists idx_history_timestamp on history(timestamp); +create index if not exists idx_history_command on history(command); diff --git a/crates/atuin-client/migrations/20220505083406_create-events.sql b/crates/atuin-client/migrations/20220505083406_create-events.sql new file mode 100644 index 00000000..f6cafeba --- /dev/null +++ b/crates/atuin-client/migrations/20220505083406_create-events.sql @@ -0,0 +1,11 @@ +create table if not exists events ( + id text primary key, + timestamp integer not null, + hostname text not null, + event_type text not null, + + history_id text not null +); + +-- Ensure there is only ever one of each event type per history item +create unique index history_event_idx ON events(event_type, history_id); diff --git a/crates/atuin-client/migrations/20220806155627_interactive_search_index.sql b/crates/atuin-client/migrations/20220806155627_interactive_search_index.sql new file mode 100644 index 00000000..b5770e62 --- /dev/null +++ b/crates/atuin-client/migrations/20220806155627_interactive_search_index.sql @@ -0,0 +1,6 @@ +-- Interactive search filters by command then by the max(timestamp) for that +-- command. Create an index that covers those +create index if not exists idx_history_command_timestamp on history( + command, + timestamp +); diff --git a/crates/atuin-client/migrations/20230315220114_drop-events.sql b/crates/atuin-client/migrations/20230315220114_drop-events.sql new file mode 100644 index 00000000..fe3cae17 --- /dev/null +++ b/crates/atuin-client/migrations/20230315220114_drop-events.sql @@ -0,0 +1,2 @@ +-- Add migration script here +drop table events; diff --git a/crates/atuin-client/migrations/20230319185725_deleted_at.sql b/crates/atuin-client/migrations/20230319185725_deleted_at.sql new file mode 100644 index 00000000..6c422abc --- /dev/null +++ b/crates/atuin-client/migrations/20230319185725_deleted_at.sql @@ -0,0 +1,2 @@ +-- Add migration script here +alter table history add column deleted_at integer; diff --git a/crates/atuin-client/migrations/20260224000100_history_author_intent.sql b/crates/atuin-client/migrations/20260224000100_history_author_intent.sql new file mode 100644 index 00000000..2bed17e9 --- /dev/null +++ b/crates/atuin-client/migrations/20260224000100_history_author_intent.sql @@ -0,0 +1,2 @@ +alter table history add column author text; +alter table history add column intent text; diff --git a/crates/atuin-client/record-migrations/20230531212437_create-records.sql b/crates/atuin-client/record-migrations/20230531212437_create-records.sql new file mode 100644 index 00000000..4f4b304a --- /dev/null +++ b/crates/atuin-client/record-migrations/20230531212437_create-records.sql @@ -0,0 +1,16 @@ +-- Add migration script here +create table if not exists records ( + id text primary key, + parent text unique, -- null if this is the first one + host text not null, + + timestamp integer not null, + tag text not null, + version text not null, + data blob not null, + cek blob not null +); + +create index host_idx on records (host); +create index tag_idx on records (tag); +create index host_tag_idx on records (host, tag); diff --git a/crates/atuin-client/record-migrations/20231127090831_create-store.sql b/crates/atuin-client/record-migrations/20231127090831_create-store.sql new file mode 100644 index 00000000..53d78860 --- /dev/null +++ b/crates/atuin-client/record-migrations/20231127090831_create-store.sql @@ -0,0 +1,15 @@ +-- Add migration script here +create table if not exists store ( + id text primary key, -- globally unique ID + + idx integer, -- incrementing integer ID unique per (host, tag) + host text not null, -- references the host row + tag text not null, + + timestamp integer not null, + version text not null, + data blob not null, + cek blob not null +); + +create unique index record_uniq ON store(host, tag, idx); diff --git a/crates/atuin-client/src/api_client.rs b/crates/atuin-client/src/api_client.rs new file mode 100644 index 00000000..066fecb5 --- /dev/null +++ b/crates/atuin-client/src/api_client.rs @@ -0,0 +1,461 @@ +use std::collections::HashMap; +use std::env; +use std::time::Duration; + +use eyre::{Result, bail, eyre}; +use reqwest::{ + Response, StatusCode, Url, + header::{AUTHORIZATION, HeaderMap, USER_AGENT}, +}; + +use atuin_common::{ + api::{ATUIN_CARGO_VERSION, ATUIN_HEADER_VERSION, ATUIN_VERSION}, + record::{EncryptedData, HostId, Record, RecordIdx}, + tls::ensure_crypto_provider, +}; +use atuin_common::{ + api::{ + AddHistoryRequest, ChangePasswordRequest, CountResponse, DeleteHistoryRequest, + ErrorResponse, LoginRequest, LoginResponse, MeResponse, RegisterResponse, StatusResponse, + SyncHistoryResponse, + }, + record::RecordStatus, +}; + +use semver::Version; +use time::OffsetDateTime; +use time::format_description::well_known::Rfc3339; + +use crate::{history::History, sync::hash_str, utils::get_host_user}; + +static APP_USER_AGENT: &str = concat!("atuin/", env!("CARGO_PKG_VERSION"),); + +/// Authentication token for sync API requests. +/// +/// The sync API supports two authentication methods: +/// - `Bearer`: Hub API tokens (for users authenticated via Atuin Hub) +/// - `Token`: Legacy CLI session tokens (for users registered via CLI or self-hosted) +/// +/// When both are available, Hub tokens are preferred as they provide unified +/// authentication across CLI and Hub features. +#[derive(Debug, Clone)] +pub enum AuthToken { + /// Hub API token, used with "Bearer {token}" header + Bearer(String), + /// Legacy CLI session token, used with "Token {token}" header + Token(String), +} + +impl AuthToken { + /// Format the token as an Authorization header value + fn to_header_value(&self) -> String { + match self { + AuthToken::Bearer(token) => format!("Bearer {token}"), + AuthToken::Token(token) => format!("Token {token}"), + } + } +} + +pub struct Client<'a> { + sync_addr: &'a str, + client: reqwest::Client, +} + +fn make_url(address: &str, path: &str) -> Result<String> { + // `join()` expects a trailing `/` in order to join paths + // e.g. it treats `http://host:port/subdir` as a file called `subdir` + let address = if address.ends_with("/") { + address + } else { + &format!("{address}/") + }; + + // passing a path with a leading `/` will cause `join()` to replace the entire URL path + let path = path.strip_prefix("/").unwrap_or(path); + + let url = Url::parse(address) + .map(|url| url.join(path))? + .map_err(|_| eyre!("invalid address"))?; + + Ok(url.to_string()) +} + +pub async fn register( + address: &str, + username: &str, + email: &str, + password: &str, +) -> Result<RegisterResponse> { + ensure_crypto_provider(); + let mut map = HashMap::new(); + map.insert("username", username); + map.insert("email", email); + map.insert("password", password); + + let url = make_url(address, &format!("/user/{username}"))?; + let resp = reqwest::get(url).await?; + + if resp.status().is_success() { + bail!("username already in use"); + } + + let url = make_url(address, "/register")?; + let client = reqwest::Client::new(); + let resp = client + .post(url) + .header(USER_AGENT, APP_USER_AGENT) + .header(ATUIN_HEADER_VERSION, ATUIN_CARGO_VERSION) + .json(&map) + .send() + .await?; + let resp = handle_resp_error(resp).await?; + + if !ensure_version(&resp)? { + bail!("could not register user due to version mismatch"); + } + + let session = resp.json::<RegisterResponse>().await?; + Ok(session) +} + +pub async fn login(address: &str, req: LoginRequest) -> Result<LoginResponse> { + ensure_crypto_provider(); + let url = make_url(address, "/login")?; + let client = reqwest::Client::new(); + + let resp = client + .post(url) + .header(USER_AGENT, APP_USER_AGENT) + .json(&req) + .send() + .await?; + let resp = handle_resp_error(resp).await?; + + if !ensure_version(&resp)? { + bail!("Could not login due to version mismatch"); + } + + let session = resp.json::<LoginResponse>().await?; + Ok(session) +} + +#[cfg(feature = "check-update")] +pub async fn latest_version() -> Result<Version> { + use atuin_common::api::IndexResponse; + + ensure_crypto_provider(); + let url = "https://api.atuin.sh"; + let client = reqwest::Client::new(); + + let resp = client + .get(url) + .header(USER_AGENT, APP_USER_AGENT) + .send() + .await?; + let resp = handle_resp_error(resp).await?; + + let index = resp.json::<IndexResponse>().await?; + let version = Version::parse(index.version.as_str())?; + + Ok(version) +} + +pub fn ensure_version(response: &Response) -> Result<bool> { + let version = response.headers().get(ATUIN_HEADER_VERSION); + + let version = if let Some(version) = version { + match version.to_str() { + Ok(v) => Version::parse(v), + Err(e) => bail!("failed to parse server version: {:?}", e), + } + } else { + bail!("Server not reporting its version: it is either too old or unhealthy"); + }?; + + // If the client is newer than the server + if version.major < ATUIN_VERSION.major { + println!( + "Atuin version mismatch! In order to successfully sync, the server needs to run a newer version of Atuin" + ); + println!("Client: {ATUIN_CARGO_VERSION}"); + println!("Server: {version}"); + + return Ok(false); + } + + Ok(true) +} + +async fn handle_resp_error(resp: Response) -> Result<Response> { + let status = resp.status(); + let url = resp.url().to_string(); + + if status == StatusCode::SERVICE_UNAVAILABLE { + bail!( + "Service unavailable: check https://status.atuin.sh (or get in touch with your host)" + ); + } + + if status == StatusCode::TOO_MANY_REQUESTS { + bail!("Rate limited; please wait before doing that again"); + } + + if !status.is_success() { + if let Ok(error) = resp.json::<ErrorResponse>().await { + let reason = error.reason; + + if status.is_client_error() { + bail!("Invalid request to the service at {url}, {status} - {reason}.") + } + + bail!( + "There was an error with the atuin sync service at {url}, server error {status}: {reason}.\nIf the problem persists, contact the host" + ) + } + + bail!( + "There was an error with the atuin sync service at {url}, Status {status:?}.\nIf the problem persists, contact the host" + ) + } + + Ok(resp) +} + +impl<'a> Client<'a> { + pub fn new( + sync_addr: &'a str, + auth: AuthToken, + connect_timeout: u64, + timeout: u64, + ) -> Result<Self> { + ensure_crypto_provider(); + let mut headers = HeaderMap::new(); + headers.insert(AUTHORIZATION, auth.to_header_value().parse()?); + + // used for semver server check + headers.insert(ATUIN_HEADER_VERSION, ATUIN_CARGO_VERSION.parse()?); + + Ok(Client { + sync_addr, + client: reqwest::Client::builder() + .user_agent(APP_USER_AGENT) + .default_headers(headers) + .connect_timeout(Duration::new(connect_timeout, 0)) + .timeout(Duration::new(timeout, 0)) + .build()?, + }) + } + + pub async fn count(&self) -> Result<i64> { + let url = make_url(self.sync_addr, "/sync/count")?; + let url = Url::parse(url.as_str())?; + + let resp = self.client.get(url).send().await?; + let resp = handle_resp_error(resp).await?; + + if !ensure_version(&resp)? { + bail!("could not sync due to version mismatch"); + } + + if resp.status() != StatusCode::OK { + bail!("failed to get count (are you logged in?)"); + } + + let count = resp.json::<CountResponse>().await?; + + Ok(count.count) + } + + pub async fn status(&self) -> Result<StatusResponse> { + let url = make_url(self.sync_addr, "/sync/status")?; + let url = Url::parse(url.as_str())?; + + let resp = self.client.get(url).send().await?; + let resp = handle_resp_error(resp).await?; + + if !ensure_version(&resp)? { + bail!("could not sync due to version mismatch"); + } + + let status = resp.json::<StatusResponse>().await?; + + Ok(status) + } + + pub async fn me(&self) -> Result<MeResponse> { + let url = make_url(self.sync_addr, "/api/v0/me")?; + let url = Url::parse(url.as_str())?; + + let resp = self.client.get(url).send().await?; + let resp = handle_resp_error(resp).await?; + + let status = resp.json::<MeResponse>().await?; + + Ok(status) + } + + pub async fn get_history( + &self, + sync_ts: OffsetDateTime, + history_ts: OffsetDateTime, + host: Option<String>, + ) -> Result<SyncHistoryResponse> { + let host = host.unwrap_or_else(|| hash_str(&get_host_user())); + + let url = make_url( + self.sync_addr, + &format!( + "/sync/history?sync_ts={}&history_ts={}&host={}", + urlencoding::encode(sync_ts.format(&Rfc3339)?.as_str()), + urlencoding::encode(history_ts.format(&Rfc3339)?.as_str()), + host, + ), + )?; + + let resp = self.client.get(url).send().await?; + let resp = handle_resp_error(resp).await?; + + let history = resp.json::<SyncHistoryResponse>().await?; + Ok(history) + } + + pub async fn post_history(&self, history: &[AddHistoryRequest]) -> Result<()> { + let url = make_url(self.sync_addr, "/history")?; + let url = Url::parse(url.as_str())?; + + let resp = self.client.post(url).json(history).send().await?; + handle_resp_error(resp).await?; + + Ok(()) + } + + pub async fn delete_history(&self, h: History) -> Result<()> { + let url = make_url(self.sync_addr, "/history")?; + let url = Url::parse(url.as_str())?; + + let resp = self + .client + .delete(url) + .json(&DeleteHistoryRequest { + client_id: h.id.to_string(), + }) + .send() + .await?; + + handle_resp_error(resp).await?; + + Ok(()) + } + + pub async fn delete_store(&self) -> Result<()> { + let url = make_url(self.sync_addr, "/api/v0/store")?; + let url = Url::parse(url.as_str())?; + + let resp = self.client.delete(url).send().await?; + + handle_resp_error(resp).await?; + + Ok(()) + } + + pub async fn post_records(&self, records: &[Record<EncryptedData>]) -> Result<()> { + let url = make_url(self.sync_addr, "/api/v0/record")?; + let url = Url::parse(url.as_str())?; + + debug!("uploading {} records to {url}", records.len()); + + let resp = self.client.post(url).json(records).send().await?; + handle_resp_error(resp).await?; + + Ok(()) + } + + pub async fn next_records( + &self, + host: HostId, + tag: String, + start: RecordIdx, + count: u64, + ) -> Result<Vec<Record<EncryptedData>>> { + debug!("fetching record/s from host {}/{}/{}", host.0, tag, start); + + let url = make_url( + self.sync_addr, + &format!( + "/api/v0/record/next?host={}&tag={}&count={}&start={}", + host.0, tag, count, start + ), + )?; + + let url = Url::parse(url.as_str())?; + + let resp = self.client.get(url).send().await?; + let resp = handle_resp_error(resp).await?; + + let records = resp.json::<Vec<Record<EncryptedData>>>().await?; + + Ok(records) + } + + pub async fn record_status(&self) -> Result<RecordStatus> { + let url = make_url(self.sync_addr, "/api/v0/record")?; + let url = Url::parse(url.as_str())?; + + let resp = self.client.get(url).send().await?; + let resp = handle_resp_error(resp).await?; + + if !ensure_version(&resp)? { + bail!("could not sync records due to version mismatch"); + } + + let index = resp.json().await?; + + debug!("got remote index {index:?}"); + + Ok(index) + } + + pub async fn delete(&self) -> Result<()> { + let url = make_url(self.sync_addr, "/account")?; + let url = Url::parse(url.as_str())?; + + let resp = self.client.delete(url).send().await?; + + if resp.status() == 403 { + bail!("invalid login details"); + } else if resp.status() == 200 { + Ok(()) + } else { + bail!("Unknown error"); + } + } + + pub async fn change_password( + &self, + current_password: String, + new_password: String, + ) -> Result<()> { + let url = make_url(self.sync_addr, "/account/password")?; + let url = Url::parse(url.as_str())?; + + let resp = self + .client + .patch(url) + .json(&ChangePasswordRequest { + current_password, + new_password, + }) + .send() + .await?; + + if resp.status() == 401 { + bail!("current password is incorrect") + } else if resp.status() == 403 { + bail!("invalid login details"); + } else if resp.status() == 200 { + Ok(()) + } else { + bail!("Unknown error"); + } + } +} diff --git a/crates/atuin-client/src/database.rs b/crates/atuin-client/src/database.rs new file mode 100644 index 00000000..7c63368d --- /dev/null +++ b/crates/atuin-client/src/database.rs @@ -0,0 +1,1486 @@ +use std::{ + env, + path::{Path, PathBuf}, + str::FromStr, + time::Duration, +}; + +use async_trait::async_trait; +use atuin_common::utils; +use fs_err as fs; +use itertools::Itertools; +use rand::{Rng, distributions::Alphanumeric}; +use sql_builder::{SqlBuilder, SqlName, bind::Bind, esc, quote}; +use sqlx::{ + Result, Row, + sqlite::{ + SqliteConnectOptions, SqliteJournalMode, SqlitePool, SqlitePoolOptions, SqliteRow, + SqliteSynchronous, + }, +}; +use time::OffsetDateTime; +use uuid::Uuid; + +use crate::{ + history::{HistoryId, HistoryStats}, + utils::get_host_user, +}; + +use super::{ + history::History, + ordering, + settings::{FilterMode, SearchMode, Settings}, +}; + +#[derive(Clone)] +pub struct Context { + pub session: String, + pub cwd: String, + pub hostname: String, + pub host_id: String, + pub git_root: Option<PathBuf>, +} + +#[derive(Default, Clone)] +pub struct OptFilters { + pub exit: Option<i64>, + pub exclude_exit: Option<i64>, + pub cwd: Option<String>, + pub exclude_cwd: Option<String>, + pub before: Option<String>, + pub after: Option<String>, + pub limit: Option<i64>, + pub offset: Option<i64>, + pub reverse: bool, + pub include_duplicates: bool, +} + +pub async fn current_context() -> eyre::Result<Context> { + let session = env::var("ATUIN_SESSION").map_err(|_| { + eyre::eyre!("Failed to find $ATUIN_SESSION in the environment. Check that you have correctly set up your shell.") + })?; + let hostname = get_host_user(); + let cwd = utils::get_current_dir(); + let host_id = Settings::host_id().await?; + let git_root = utils::in_git_repo(cwd.as_str()); + + Ok(Context { + session, + hostname, + cwd, + git_root, + host_id: host_id.0.as_simple().to_string(), + }) +} + +impl Context { + pub fn from_history(entry: &History) -> Self { + Context { + session: entry.session.to_string(), + cwd: entry.cwd.to_string(), + hostname: entry.hostname.to_string(), + host_id: String::new(), + git_root: utils::in_git_repo(entry.cwd.as_str()), + } + } +} + +fn get_session_start_time(session_id: &str) -> Option<i64> { + if let Ok(uuid) = Uuid::parse_str(session_id) + && let Some(timestamp) = uuid.get_timestamp() + { + let (seconds, nanos) = timestamp.to_unix(); + return Some(seconds as i64 * 1_000_000_000 + nanos as i64); + } + None +} + +#[async_trait] +pub trait Database: Send + Sync + 'static { + async fn save(&self, h: &History) -> Result<()>; + async fn save_bulk(&self, h: &[History]) -> Result<()>; + + async fn load(&self, id: &str) -> Result<Option<History>>; + async fn list( + &self, + filters: &[FilterMode], + context: &Context, + max: Option<usize>, + unique: bool, + include_deleted: bool, + ) -> Result<Vec<History>>; + async fn range(&self, from: OffsetDateTime, to: OffsetDateTime) -> Result<Vec<History>>; + + async fn update(&self, h: &History) -> Result<()>; + async fn history_count(&self, include_deleted: bool) -> Result<i64>; + + async fn last(&self) -> Result<Option<History>>; + async fn before(&self, timestamp: OffsetDateTime, count: i64) -> Result<Vec<History>>; + + async fn delete(&self, h: History) -> Result<()>; + async fn delete_rows(&self, ids: &[HistoryId]) -> Result<()>; + async fn deleted(&self) -> Result<Vec<History>>; + + // Yes I know, it's a lot. + // Could maybe break it down to a searchparams struct or smth but that feels a little... pointless. + // Been debating maybe a DSL for search? eg "before:time limit:1 the query" + #[allow(clippy::too_many_arguments)] + async fn search( + &self, + search_mode: SearchMode, + filter: FilterMode, + context: &Context, + query: &str, + filter_options: OptFilters, + ) -> Result<Vec<History>>; + + async fn query_history(&self, query: &str) -> Result<Vec<History>>; + + async fn all_with_count(&self) -> Result<Vec<(History, i32)>>; + + fn all_paged(&self, page_size: usize, include_deleted: bool, unique: bool) -> Paged; + + async fn stats(&self, h: &History) -> Result<HistoryStats>; + + async fn get_dups(&self, before: i64, dupkeep: u32) -> Result<Vec<History>>; + + fn clone_boxed(&self) -> Box<dyn Database + 'static>; +} + +// Intended for use on a developer machine and not a sync server. +// TODO: implement IntoIterator +#[derive(Debug, Clone)] +pub struct Sqlite { + pub pool: SqlitePool, +} + +impl Sqlite { + pub async fn new(path: impl AsRef<Path>, timeout: f64) -> Result<Self> { + let path = path.as_ref(); + debug!("opening sqlite database at {path:?}"); + + if utils::broken_symlink(path) { + eprintln!( + "Atuin: Sqlite db path ({path:?}) is a broken symlink. Unable to read or create replacement." + ); + std::process::exit(1); + } + + if !path.exists() + && let Some(dir) = path.parent() + { + fs::create_dir_all(dir)?; + } + + let opts = SqliteConnectOptions::from_str(path.as_os_str().to_str().unwrap())? + .journal_mode(SqliteJournalMode::Wal) + .optimize_on_close(true, None) + .synchronous(SqliteSynchronous::Normal) + .with_regexp() + .create_if_missing(true); + + let pool = SqlitePoolOptions::new() + .acquire_timeout(Duration::from_secs_f64(timeout)) + .connect_with(opts) + .await?; + + Self::setup_db(&pool).await?; + Ok(Self { pool }) + } + + pub async fn sqlite_version(&self) -> Result<String> { + sqlx::query_scalar("SELECT sqlite_version()") + .fetch_one(&self.pool) + .await + } + + async fn setup_db(pool: &SqlitePool) -> Result<()> { + debug!("running sqlite database setup"); + + sqlx::migrate!("./migrations").run(pool).await?; + + Ok(()) + } + + async fn save_raw(tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, h: &History) -> Result<()> { + sqlx::query( + "insert or ignore into history(id, timestamp, duration, exit, command, cwd, session, hostname, author, intent, deleted_at) + values(?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11)", + ) + .bind(h.id.0.as_str()) + .bind(h.timestamp.unix_timestamp_nanos() as i64) + .bind(h.duration) + .bind(h.exit) + .bind(h.command.as_str()) + .bind(h.cwd.as_str()) + .bind(h.session.as_str()) + .bind(h.hostname.as_str()) + .bind(h.author.as_str()) + .bind(h.intent.as_deref()) + .bind(h.deleted_at.map(|t|t.unix_timestamp_nanos() as i64)) + .execute(&mut **tx) + .await?; + + Ok(()) + } + + async fn delete_row_raw( + tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, + id: HistoryId, + ) -> Result<()> { + sqlx::query("delete from history where id = ?1") + .bind(id.0.as_str()) + .execute(&mut **tx) + .await?; + + Ok(()) + } + + fn query_history(row: SqliteRow) -> History { + let deleted_at: Option<i64> = row.get("deleted_at"); + let hostname: String = row.get("hostname"); + let author: Option<String> = row.try_get("author").ok().flatten(); + let author = author + .filter(|author| !author.trim().is_empty()) + .unwrap_or_else(|| History::author_from_hostname(hostname.as_str())); + let intent: Option<String> = row.try_get("intent").ok().flatten(); + let intent = intent.filter(|intent| !intent.trim().is_empty()); + + History::from_db() + .id(row.get("id")) + .timestamp( + OffsetDateTime::from_unix_timestamp_nanos(row.get::<i64, _>("timestamp") as i128) + .unwrap(), + ) + .duration(row.get("duration")) + .exit(row.get("exit")) + .command(row.get("command")) + .cwd(row.get("cwd")) + .session(row.get("session")) + .hostname(hostname) + .author(author) + .intent(intent) + .deleted_at( + deleted_at.and_then(|t| OffsetDateTime::from_unix_timestamp_nanos(t as i128).ok()), + ) + .build() + .into() + } +} + +#[async_trait] +impl Database for Sqlite { + async fn save(&self, h: &History) -> Result<()> { + debug!("saving history to sqlite"); + let mut tx = self.pool.begin().await?; + Self::save_raw(&mut tx, h).await?; + tx.commit().await?; + + Ok(()) + } + + async fn save_bulk(&self, h: &[History]) -> Result<()> { + debug!("saving history to sqlite"); + + let mut tx = self.pool.begin().await?; + + for i in h { + Self::save_raw(&mut tx, i).await?; + } + + tx.commit().await?; + + Ok(()) + } + + async fn load(&self, id: &str) -> Result<Option<History>> { + debug!("loading history item {}", id); + + let res = sqlx::query("select * from history where id = ?1") + .bind(id) + .map(Self::query_history) + .fetch_optional(&self.pool) + .await?; + + Ok(res) + } + + async fn update(&self, h: &History) -> Result<()> { + debug!("updating sqlite history"); + + sqlx::query( + "update history + set timestamp = ?2, duration = ?3, exit = ?4, command = ?5, cwd = ?6, session = ?7, hostname = ?8, author = ?9, intent = ?10, deleted_at = ?11 + where id = ?1", + ) + .bind(h.id.0.as_str()) + .bind(h.timestamp.unix_timestamp_nanos() as i64) + .bind(h.duration) + .bind(h.exit) + .bind(h.command.as_str()) + .bind(h.cwd.as_str()) + .bind(h.session.as_str()) + .bind(h.hostname.as_str()) + .bind(h.author.as_str()) + .bind(h.intent.as_deref()) + .bind(h.deleted_at.map(|t|t.unix_timestamp_nanos() as i64)) + .execute(&self.pool) + .await?; + + Ok(()) + } + + // make a unique list, that only shows the *newest* version of things + async fn list( + &self, + filters: &[FilterMode], + context: &Context, + max: Option<usize>, + unique: bool, + include_deleted: bool, + ) -> Result<Vec<History>> { + debug!("listing history"); + + let mut query = SqlBuilder::select_from(SqlName::new("history").alias("h").baquoted()); + query.field("*").order_desc("timestamp"); + if !include_deleted { + query.and_where_is_null("deleted_at"); + } + + let git_root = if let Some(git_root) = context.git_root.clone() { + git_root.to_str().unwrap_or("/").to_string() + } else { + context.cwd.clone() + }; + + let session_start = get_session_start_time(&context.session); + + for filter in filters { + match filter { + FilterMode::Global => &mut query, + FilterMode::Host => query.and_where_eq("hostname", quote(&context.hostname)), + FilterMode::Session => query.and_where_eq("session", quote(&context.session)), + FilterMode::SessionPreload => { + query.and_where_eq("session", quote(&context.session)); + if let Some(session_start) = session_start { + query.or_where_lt("timestamp", session_start); + } + &mut query + } + FilterMode::Directory => query.and_where_eq("cwd", quote(&context.cwd)), + FilterMode::Workspace => query.and_where_like_left("cwd", &git_root), + }; + } + + if unique { + query.group_by("command").having("max(timestamp)"); + } + + if let Some(max) = max { + query.limit(max); + } + + let query = query.sql().expect("bug in list query. please report"); + + let res = sqlx::query(&query) + .map(Self::query_history) + .fetch_all(&self.pool) + .await?; + + Ok(res) + } + + async fn range(&self, from: OffsetDateTime, to: OffsetDateTime) -> Result<Vec<History>> { + debug!("listing history from {:?} to {:?}", from, to); + + let res = sqlx::query( + "select * from history where timestamp >= ?1 and timestamp <= ?2 order by timestamp asc", + ) + .bind(from.unix_timestamp_nanos() as i64) + .bind(to.unix_timestamp_nanos() as i64) + .map(Self::query_history) + .fetch_all(&self.pool) + .await?; + + Ok(res) + } + + async fn last(&self) -> Result<Option<History>> { + let res = sqlx::query( + "select * from history where duration >= 0 order by timestamp desc limit 1", + ) + .map(Self::query_history) + .fetch_optional(&self.pool) + .await?; + + Ok(res) + } + + async fn before(&self, timestamp: OffsetDateTime, count: i64) -> Result<Vec<History>> { + let res = sqlx::query( + "select * from history where timestamp < ?1 order by timestamp desc limit ?2", + ) + .bind(timestamp.unix_timestamp_nanos() as i64) + .bind(count) + .map(Self::query_history) + .fetch_all(&self.pool) + .await?; + + Ok(res) + } + + async fn deleted(&self) -> Result<Vec<History>> { + let res = sqlx::query("select * from history where deleted_at is not null") + .map(Self::query_history) + .fetch_all(&self.pool) + .await?; + + Ok(res) + } + + async fn history_count(&self, include_deleted: bool) -> Result<i64> { + let query = if include_deleted { + "select count(1) from history" + } else { + "select count(1) from history where deleted_at is null" + }; + + let res: (i64,) = sqlx::query_as(query).fetch_one(&self.pool).await?; + Ok(res.0) + } + + async fn search( + &self, + search_mode: SearchMode, + filter: FilterMode, + context: &Context, + query: &str, + filter_options: OptFilters, + ) -> Result<Vec<History>> { + let mut sql = SqlBuilder::select_from("history"); + + if !filter_options.include_duplicates { + sql.group_by("command").having("max(timestamp)"); + } + + if let Some(limit) = filter_options.limit { + sql.limit(limit); + } + + if let Some(offset) = filter_options.offset { + sql.offset(offset); + } + + if filter_options.reverse { + sql.order_asc("timestamp"); + } else { + sql.order_desc("timestamp"); + } + + let git_root = if let Some(git_root) = context.git_root.clone() { + git_root.to_str().unwrap_or("/").to_string() + } else { + context.cwd.clone() + }; + + let session_start = get_session_start_time(&context.session); + + match filter { + FilterMode::Global => &mut sql, + FilterMode::Host => { + sql.and_where_eq("lower(hostname)", quote(context.hostname.to_lowercase())) + } + FilterMode::Session => sql.and_where_eq("session", quote(&context.session)), + FilterMode::SessionPreload => { + sql.and_where_eq("session", quote(&context.session)); + if let Some(session_start) = session_start { + sql.or_where_lt("timestamp", session_start); + } + &mut sql + } + FilterMode::Directory => sql.and_where_eq("cwd", quote(&context.cwd)), + FilterMode::Workspace => sql.and_where_like_left("cwd", git_root), + }; + + let orig_query = query; + + let mut regexes = Vec::new(); + match search_mode { + SearchMode::Prefix => sql.and_where_like_left("command", query.replace('*', "%")), + _ => { + let mut is_or = false; + for token in QueryTokenizer::new(query) { + // TODO smart case mode could be made configurable like in fzf + let (is_glob, glob) = if token.has_uppercase() { + (true, "*") + } else { + (false, "%") + }; + let param = match token { + QueryToken::Regex(r) => { + regexes.push(String::from(r)); + continue; + } + QueryToken::Or => { + if !is_or { + is_or = true; + continue; + } else { + format!("{glob}|{glob}") + } + } + QueryToken::MatchStart(term, _) => { + format!("{term}{glob}") + } + QueryToken::MatchEnd(term, _) => { + format!("{glob}{term}") + } + QueryToken::MatchFull(term, _) => { + format!("{glob}{term}{glob}") + } + QueryToken::Match(term, _) => { + if search_mode == SearchMode::FullText { + format!("{glob}{term}{glob}") + } else { + term.split("").join(glob) + } + } + }; + + sql.fuzzy_condition("command", param, token.is_inverse(), is_glob, is_or); + is_or = false; + } + + &mut sql + } + }; + + for regex in regexes { + sql.and_where("command regexp ?".bind(®ex)); + } + + filter_options + .exit + .map(|exit| sql.and_where_eq("exit", exit)); + + filter_options + .exclude_exit + .map(|exclude_exit| sql.and_where_ne("exit", exclude_exit)); + + filter_options + .cwd + .map(|cwd| sql.and_where_eq("cwd", quote(cwd))); + + filter_options + .exclude_cwd + .map(|exclude_cwd| sql.and_where_ne("cwd", quote(exclude_cwd))); + + filter_options.before.map(|before| { + interim::parse_date_string( + before.as_str(), + OffsetDateTime::now_utc(), + interim::Dialect::Uk, + ) + .map(|before| { + sql.and_where_lt("timestamp", quote(before.unix_timestamp_nanos() as i64)) + }) + }); + + filter_options.after.map(|after| { + interim::parse_date_string( + after.as_str(), + OffsetDateTime::now_utc(), + interim::Dialect::Uk, + ) + .map(|after| sql.and_where_gt("timestamp", quote(after.unix_timestamp_nanos() as i64))) + }); + + sql.and_where_is_null("deleted_at"); + + let query = sql.sql().expect("bug in search query. please report"); + + let res = sqlx::query(&query) + .map(Self::query_history) + .fetch_all(&self.pool) + .await?; + + Ok(ordering::reorder_fuzzy(search_mode, orig_query, res)) + } + + async fn query_history(&self, query: &str) -> Result<Vec<History>> { + let res = sqlx::query(query) + .map(Self::query_history) + .fetch_all(&self.pool) + .await?; + + Ok(res) + } + + async fn all_with_count(&self) -> Result<Vec<(History, i32)>> { + debug!("listing history"); + + let mut query = SqlBuilder::select_from(SqlName::new("history").alias("h").baquoted()); + + query + .fields(&[ + "id", + "max(timestamp) as timestamp", + "max(duration) as duration", + "exit", + "command", + "deleted_at", + "null as author", + "null as intent", + "group_concat(cwd, ':') as cwd", + "group_concat(session) as session", + "group_concat(hostname, ',') as hostname", + "count(*) as count", + ]) + .group_by("command") + .group_by("exit") + .and_where("deleted_at is null") + .order_desc("timestamp"); + + let query = query.sql().expect("bug in list query. please report"); + + let res = sqlx::query(&query) + .map(|row: SqliteRow| { + let count: i32 = row.get("count"); + (Self::query_history(row), count) + }) + .fetch_all(&self.pool) + .await?; + + Ok(res) + } + + fn all_paged(&self, page_size: usize, include_deleted: bool, unique: bool) -> Paged { + Paged::new(Box::new(self.clone()), page_size, include_deleted, unique) + } + + // deleted_at doesn't mean the actual time that the user deleted it, + // but the time that the system marks it as deleted + async fn delete(&self, mut h: History) -> Result<()> { + let now = OffsetDateTime::now_utc(); + h.command = rand::thread_rng() + .sample_iter(&Alphanumeric) + .take(32) + .map(char::from) + .collect(); // overwrite with random string + h.deleted_at = Some(now); // delete it + + self.update(&h).await?; // save it + + Ok(()) + } + + async fn delete_rows(&self, ids: &[HistoryId]) -> Result<()> { + let mut tx = self.pool.begin().await?; + + for id in ids { + Self::delete_row_raw(&mut tx, id.clone()).await?; + } + + tx.commit().await?; + + Ok(()) + } + + async fn stats(&self, h: &History) -> Result<HistoryStats> { + // We select the previous in the session by time + let mut prev = SqlBuilder::select_from("history"); + prev.field("*") + .and_where("timestamp < ?1") + .and_where("session = ?2") + .order_by("timestamp", true) + .limit(1); + + let mut next = SqlBuilder::select_from("history"); + next.field("*") + .and_where("timestamp > ?1") + .and_where("session = ?2") + .order_by("timestamp", false) + .limit(1); + + let mut total = SqlBuilder::select_from("history"); + total.field("count(1)").and_where("command = ?1"); + + let mut average = SqlBuilder::select_from("history"); + average.field("avg(duration)").and_where("command = ?1"); + + let mut exits = SqlBuilder::select_from("history"); + exits + .fields(&["exit", "count(1) as count"]) + .and_where("command = ?1") + .group_by("exit"); + + // rewrite the following with sqlbuilder + let mut day_of_week = SqlBuilder::select_from("history"); + day_of_week + .fields(&[ + "strftime('%w', ROUND(timestamp / 1000000000), 'unixepoch') AS day_of_week", + "count(1) as count", + ]) + .and_where("command = ?1") + .group_by("day_of_week"); + + // Intentionally format the string with 01 hardcoded. We want the average runtime for the + // _entire month_, but will later parse it as a datetime for sorting + // Sqlite has no datetime so we cannot do it there, and otherwise sorting will just be a + // string sort, which won't be correct. + let mut duration_over_time = SqlBuilder::select_from("history"); + duration_over_time + .fields(&[ + "strftime('01-%m-%Y', ROUND(timestamp / 1000000000), 'unixepoch') AS month_year", + "avg(duration) as duration", + ]) + .and_where("command = ?1") + .group_by("month_year") + .having("duration > 0"); + + let prev = prev.sql().expect("issue in stats previous query"); + let next = next.sql().expect("issue in stats next query"); + let total = total.sql().expect("issue in stats average query"); + let average = average.sql().expect("issue in stats previous query"); + let exits = exits.sql().expect("issue in stats exits query"); + let day_of_week = day_of_week.sql().expect("issue in stats day of week query"); + let duration_over_time = duration_over_time + .sql() + .expect("issue in stats duration over time query"); + + let prev = sqlx::query(&prev) + .bind(h.timestamp.unix_timestamp_nanos() as i64) + .bind(&h.session) + .map(Self::query_history) + .fetch_optional(&self.pool) + .await?; + + let next = sqlx::query(&next) + .bind(h.timestamp.unix_timestamp_nanos() as i64) + .bind(&h.session) + .map(Self::query_history) + .fetch_optional(&self.pool) + .await?; + + let total: (i64,) = sqlx::query_as(&total) + .bind(&h.command) + .fetch_one(&self.pool) + .await?; + + let average: (f64,) = sqlx::query_as(&average) + .bind(&h.command) + .fetch_one(&self.pool) + .await?; + + let exits: Vec<(i64, i64)> = sqlx::query_as(&exits) + .bind(&h.command) + .fetch_all(&self.pool) + .await?; + + let day_of_week: Vec<(String, i64)> = sqlx::query_as(&day_of_week) + .bind(&h.command) + .fetch_all(&self.pool) + .await?; + + let duration_over_time: Vec<(String, f64)> = sqlx::query_as(&duration_over_time) + .bind(&h.command) + .fetch_all(&self.pool) + .await?; + + let duration_over_time = duration_over_time + .iter() + .map(|f| (f.0.clone(), f.1.round() as i64)) + .collect(); + + Ok(HistoryStats { + next, + previous: prev, + total: total.0 as u64, + average_duration: average.0 as u64, + exits, + day_of_week, + duration_over_time, + }) + } + + async fn get_dups(&self, before: i64, dupkeep: u32) -> Result<Vec<History>> { + let res = sqlx::query( + "SELECT * FROM ( + SELECT *, ROW_NUMBER() + OVER (PARTITION BY command, cwd, hostname ORDER BY timestamp DESC) + AS rn + FROM history + ) sub + WHERE rn > ?1 and timestamp < ?2; + ", + ) + .bind(dupkeep) + .bind(before) + .map(Self::query_history) + .fetch_all(&self.pool) + .await?; + + Ok(res) + } + + fn clone_boxed(&self) -> Box<dyn Database + 'static> { + Box::new(self.clone()) + } +} + +pub struct Paged { + database: Box<dyn Database + 'static>, + page_size: usize, + last_id: Option<String>, + include_deleted: bool, + unique: bool, +} + +impl Paged { + pub fn new( + database: Box<dyn Database + 'static>, + page_size: usize, + include_deleted: bool, + unique: bool, + ) -> Self { + Self { + database, + page_size, + last_id: None, + include_deleted, + unique, + } + } + + pub async fn next(&mut self) -> Result<Option<Vec<History>>> { + let mut query = SqlBuilder::select_from(SqlName::new("history").alias("h").baquoted()); + + query.field("*").order_desc("id"); + + if !self.include_deleted { + query.and_where_is_null("deleted_at"); + } + + if self.unique { + // We want to deduplicate on command, but the user can search via cwd, hostname, and session. + // Without those fields, filter modes won't work right. With those fields, we get duplicates. + // This must be handled upstream. + query + .group_by("command, cwd, hostname, session") + .having("max(timestamp)"); + } + + query.limit(self.page_size); + + if let Some(last_id) = &self.last_id { + query.and_where_lt("id", quote(last_id)); + } + + let query = query.sql().expect("bug in list query. please report"); + let res = self.database.query_history(&query).await?; + + if res.is_empty() { + Ok(None) + } else { + self.last_id = Some(res.last().unwrap().id.0.clone()); + Ok(Some(res)) + } + } +} + +trait SqlBuilderExt { + fn fuzzy_condition<S: ToString, T: ToString>( + &mut self, + field: S, + mask: T, + inverse: bool, + glob: bool, + is_or: bool, + ) -> &mut Self; +} + +impl SqlBuilderExt for SqlBuilder { + /// adapted from the sql-builder *like functions + fn fuzzy_condition<S: ToString, T: ToString>( + &mut self, + field: S, + mask: T, + inverse: bool, + glob: bool, + is_or: bool, + ) -> &mut Self { + let mut cond = field.to_string(); + if inverse { + cond.push_str(" NOT"); + } + if glob { + cond.push_str(" GLOB '"); + } else { + cond.push_str(" LIKE '"); + } + cond.push_str(&esc(mask.to_string())); + cond.push('\''); + if is_or { + self.or_where(cond) + } else { + self.and_where(cond) + } + } +} + +#[cfg(test)] +mod test { + use crate::settings::test_local_timeout; + + use super::*; + use std::time::{Duration, Instant}; + + async fn assert_search_eq( + db: &impl Database, + mode: SearchMode, + filter_mode: FilterMode, + query: &str, + expected: usize, + ) -> Result<Vec<History>> { + let context = Context { + hostname: "test:host".to_string(), + session: "beepboopiamasession".to_string(), + cwd: "/home/ellie".to_string(), + host_id: "test-host".to_string(), + git_root: None, + }; + + let results = db + .search( + mode, + filter_mode, + &context, + query, + OptFilters { + ..Default::default() + }, + ) + .await?; + + assert_eq!( + results.len(), + expected, + "query \"{}\", commands: {:?}", + query, + results.iter().map(|a| &a.command).collect::<Vec<&String>>() + ); + Ok(results) + } + + async fn assert_search_commands( + db: &impl Database, + mode: SearchMode, + filter_mode: FilterMode, + query: &str, + expected_commands: Vec<&str>, + ) { + let results = assert_search_eq(db, mode, filter_mode, query, expected_commands.len()) + .await + .unwrap(); + let commands: Vec<&str> = results.iter().map(|a| a.command.as_str()).collect(); + assert_eq!(commands, expected_commands); + } + + async fn new_history_item(db: &mut impl Database, cmd: &str) -> Result<()> { + let mut captured: History = History::capture() + .timestamp(OffsetDateTime::now_utc()) + .command(cmd) + .cwd("/home/ellie") + .build() + .into(); + + captured.exit = 0; + captured.duration = 1; + captured.session = "beep boop".to_string(); + captured.hostname = "booop".to_string(); + + db.save(&captured).await + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_search_prefix() { + let mut db = Sqlite::new("sqlite::memory:", test_local_timeout()) + .await + .unwrap(); + new_history_item(&mut db, "ls /home/ellie").await.unwrap(); + + assert_search_eq(&db, SearchMode::Prefix, FilterMode::Global, "ls", 1) + .await + .unwrap(); + assert_search_eq(&db, SearchMode::Prefix, FilterMode::Global, "/home", 0) + .await + .unwrap(); + assert_search_eq(&db, SearchMode::Prefix, FilterMode::Global, "ls ", 0) + .await + .unwrap(); + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_search_fulltext() { + let mut db = Sqlite::new("sqlite::memory:", test_local_timeout()) + .await + .unwrap(); + new_history_item(&mut db, "ls /home/ellie").await.unwrap(); + + assert_search_eq(&db, SearchMode::FullText, FilterMode::Global, "ls", 1) + .await + .unwrap(); + assert_search_eq(&db, SearchMode::FullText, FilterMode::Global, "/home", 1) + .await + .unwrap(); + assert_search_eq(&db, SearchMode::FullText, FilterMode::Global, "ls ho", 1) + .await + .unwrap(); + assert_search_eq(&db, SearchMode::FullText, FilterMode::Global, "hm", 0) + .await + .unwrap(); + + // regex + assert_search_eq(&db, SearchMode::FullText, FilterMode::Global, "r/^ls ", 1) + .await + .unwrap(); + assert_search_eq( + &db, + SearchMode::FullText, + FilterMode::Global, + "r/ls / ie$", + 1, + ) + .await + .unwrap(); + assert_search_eq( + &db, + SearchMode::FullText, + FilterMode::Global, + "r/ls / !ie", + 0, + ) + .await + .unwrap(); + assert_search_eq( + &db, + SearchMode::FullText, + FilterMode::Global, + "meow r/ls/", + 0, + ) + .await + .unwrap(); + assert_search_eq(&db, SearchMode::FullText, FilterMode::Global, "r//hom/", 1) + .await + .unwrap(); + assert_search_eq( + &db, + SearchMode::FullText, + FilterMode::Global, + "r//home//", + 1, + ) + .await + .unwrap(); + assert_search_eq( + &db, + SearchMode::FullText, + FilterMode::Global, + "r//home///", + 0, + ) + .await + .unwrap(); + assert_search_eq(&db, SearchMode::FullText, FilterMode::Global, "/home.*e", 0) + .await + .unwrap(); + assert_search_eq( + &db, + SearchMode::FullText, + FilterMode::Global, + "r/home.*e", + 1, + ) + .await + .unwrap(); + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_search_fuzzy() { + let mut db = Sqlite::new("sqlite::memory:", test_local_timeout()) + .await + .unwrap(); + new_history_item(&mut db, "ls /home/ellie").await.unwrap(); + new_history_item(&mut db, "ls /home/frank").await.unwrap(); + new_history_item(&mut db, "cd /home/Ellie").await.unwrap(); + new_history_item(&mut db, "/home/ellie/.bin/rustup") + .await + .unwrap(); + + assert_search_eq(&db, SearchMode::Fuzzy, FilterMode::Global, "ls /", 3) + .await + .unwrap(); + assert_search_eq(&db, SearchMode::Fuzzy, FilterMode::Global, "ls/", 2) + .await + .unwrap(); + assert_search_eq(&db, SearchMode::Fuzzy, FilterMode::Global, "l/h/", 2) + .await + .unwrap(); + assert_search_eq(&db, SearchMode::Fuzzy, FilterMode::Global, "/h/e", 3) + .await + .unwrap(); + assert_search_eq(&db, SearchMode::Fuzzy, FilterMode::Global, "/hmoe/", 0) + .await + .unwrap(); + assert_search_eq(&db, SearchMode::Fuzzy, FilterMode::Global, "ellie/home", 0) + .await + .unwrap(); + assert_search_eq(&db, SearchMode::Fuzzy, FilterMode::Global, "lsellie", 1) + .await + .unwrap(); + assert_search_eq(&db, SearchMode::Fuzzy, FilterMode::Global, " ", 4) + .await + .unwrap(); + + // single term operators + assert_search_eq(&db, SearchMode::Fuzzy, FilterMode::Global, "^ls", 2) + .await + .unwrap(); + assert_search_eq(&db, SearchMode::Fuzzy, FilterMode::Global, "'ls", 2) + .await + .unwrap(); + assert_search_eq(&db, SearchMode::Fuzzy, FilterMode::Global, "ellie$", 2) + .await + .unwrap(); + assert_search_eq(&db, SearchMode::Fuzzy, FilterMode::Global, "!^ls", 2) + .await + .unwrap(); + assert_search_eq(&db, SearchMode::Fuzzy, FilterMode::Global, "!ellie", 1) + .await + .unwrap(); + assert_search_eq(&db, SearchMode::Fuzzy, FilterMode::Global, "!ellie$", 2) + .await + .unwrap(); + + // multiple terms + assert_search_eq(&db, SearchMode::Fuzzy, FilterMode::Global, "ls !ellie", 1) + .await + .unwrap(); + assert_search_eq(&db, SearchMode::Fuzzy, FilterMode::Global, "^ls !e$", 1) + .await + .unwrap(); + assert_search_eq(&db, SearchMode::Fuzzy, FilterMode::Global, "home !^ls", 2) + .await + .unwrap(); + assert_search_eq( + &db, + SearchMode::Fuzzy, + FilterMode::Global, + "'frank | 'rustup", + 2, + ) + .await + .unwrap(); + assert_search_eq( + &db, + SearchMode::Fuzzy, + FilterMode::Global, + "'frank | 'rustup 'ls", + 1, + ) + .await + .unwrap(); + + // case matching + assert_search_eq(&db, SearchMode::Fuzzy, FilterMode::Global, "Ellie", 1) + .await + .unwrap(); + + // regex + assert_search_eq(&db, SearchMode::Fuzzy, FilterMode::Global, "r/^ls ", 2) + .await + .unwrap(); + assert_search_eq(&db, SearchMode::Fuzzy, FilterMode::Global, "r/[Ee]llie", 3) + .await + .unwrap(); + assert_search_eq(&db, SearchMode::Fuzzy, FilterMode::Global, "/h/e r/^ls ", 1) + .await + .unwrap(); + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_search_reordered_fuzzy() { + let mut db = Sqlite::new("sqlite::memory:", test_local_timeout()) + .await + .unwrap(); + // test ordering of results: we should choose the first, even though it happened longer ago. + + new_history_item(&mut db, "curl").await.unwrap(); + new_history_item(&mut db, "corburl").await.unwrap(); + + // if fuzzy reordering is on, it should come back in a more sensible order + assert_search_commands( + &db, + SearchMode::Fuzzy, + FilterMode::Global, + "curl", + vec!["curl", "corburl"], + ) + .await; + + assert_search_eq(&db, SearchMode::Fuzzy, FilterMode::Global, "xxxx", 0) + .await + .unwrap(); + assert_search_eq(&db, SearchMode::Fuzzy, FilterMode::Global, "", 2) + .await + .unwrap(); + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_paged_basic() { + let mut db = Sqlite::new("sqlite::memory:", test_local_timeout()) + .await + .unwrap(); + + // Add 5 history items + for i in 0..5 { + new_history_item(&mut db, &format!("command{}", i)) + .await + .unwrap(); + } + + // Create a paged iterator with page_size of 2 + let mut paged = db.all_paged(2, false, false); + + // First page should have 2 items + let page1 = paged.next().await.unwrap(); + assert!(page1.is_some()); + assert_eq!(page1.unwrap().len(), 2); + + // Second page should have 2 items + let page2 = paged.next().await.unwrap(); + assert!(page2.is_some()); + assert_eq!(page2.unwrap().len(), 2); + + // Third page should have 1 item + let page3 = paged.next().await.unwrap(); + assert!(page3.is_some()); + assert_eq!(page3.unwrap().len(), 1); + + // Fourth page should be None (exhausted) + let page4 = paged.next().await.unwrap(); + assert!(page4.is_none()); + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_paged_empty() { + let db = Sqlite::new("sqlite::memory:", test_local_timeout()) + .await + .unwrap(); + + // Create a paged iterator on empty database + let mut paged = db.all_paged(10, false, false); + + // Should return None immediately + let page = paged.next().await.unwrap(); + assert!(page.is_none()); + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_paged_unique() { + let mut db = Sqlite::new("sqlite::memory:", test_local_timeout()) + .await + .unwrap(); + + // Add duplicate commands + new_history_item(&mut db, "duplicate").await.unwrap(); + new_history_item(&mut db, "duplicate").await.unwrap(); + new_history_item(&mut db, "unique1").await.unwrap(); + new_history_item(&mut db, "unique2").await.unwrap(); + + // Without unique flag - should get all 4 + let mut paged = db.all_paged(10, false, false); + let page = paged.next().await.unwrap().unwrap(); + assert_eq!(page.len(), 4); + + // With unique flag - should get 3 (duplicates collapsed) + let mut paged_unique = db.all_paged(10, false, true); + let page_unique = paged_unique.next().await.unwrap().unwrap(); + assert_eq!(page_unique.len(), 3); + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_paged_include_deleted() { + let mut db = Sqlite::new("sqlite::memory:", test_local_timeout()) + .await + .unwrap(); + + // Add items + new_history_item(&mut db, "keep1").await.unwrap(); + new_history_item(&mut db, "keep2").await.unwrap(); + new_history_item(&mut db, "delete_me").await.unwrap(); + + // Delete one item + let all = db + .list( + &[], + &Context { + hostname: "".to_string(), + session: "".to_string(), + cwd: "".to_string(), + host_id: "".to_string(), + git_root: None, + }, + None, + false, + false, + ) + .await + .unwrap(); + + let to_delete = all + .iter() + .find(|h| h.command == "delete_me") + .unwrap() + .clone(); + db.delete(to_delete).await.unwrap(); + + // Without include_deleted - should get 2 + let mut paged = db.all_paged(10, false, false); + let page = paged.next().await.unwrap().unwrap(); + assert_eq!(page.len(), 2); + + // With include_deleted - should get 3 + let mut paged_deleted = db.all_paged(10, true, false); + let page_deleted = paged_deleted.next().await.unwrap().unwrap(); + assert_eq!(page_deleted.len(), 3); + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_search_bench_dupes() { + let context = Context { + hostname: "test:host".to_string(), + session: "beepboopiamasession".to_string(), + cwd: "/home/ellie".to_string(), + host_id: "test-host".to_string(), + git_root: None, + }; + + let mut db = Sqlite::new("sqlite::memory:", test_local_timeout()) + .await + .unwrap(); + for _i in 1..10000 { + new_history_item(&mut db, "i am a duplicated command") + .await + .unwrap(); + } + let start = Instant::now(); + let _results = db + .search( + SearchMode::Fuzzy, + FilterMode::Global, + &context, + "", + OptFilters { + ..Default::default() + }, + ) + .await + .unwrap(); + let duration = start.elapsed(); + + assert!(duration < Duration::from_secs(15)); + } +} + +pub struct QueryTokenizer<'a> { + query: &'a str, + last_pos: usize, +} + +pub enum QueryToken<'a> { + Match(&'a str, bool), + MatchStart(&'a str, bool), + MatchEnd(&'a str, bool), + MatchFull(&'a str, bool), + Or, + Regex(&'a str), +} + +impl<'a> QueryToken<'a> { + pub fn has_uppercase(&self) -> bool { + match self { + Self::Match(term, _) + | Self::MatchStart(term, _) + | Self::MatchEnd(term, _) + | Self::MatchFull(term, _) => term.contains(char::is_uppercase), + _ => false, + } + } + + pub fn is_inverse(&self) -> bool { + match self { + Self::Match(_, inv) + | Self::MatchStart(_, inv) + | Self::MatchEnd(_, inv) + | Self::MatchFull(_, inv) => *inv, + _ => false, + } + } +} + +impl<'a> QueryTokenizer<'a> { + pub fn new(query: &'a str) -> Self { + Self { query, last_pos: 0 } + } +} + +impl<'a> Iterator for QueryTokenizer<'a> { + type Item = QueryToken<'a>; + fn next(&mut self) -> Option<Self::Item> { + let remaining = &self.query[self.last_pos..]; + if remaining.is_empty() { + return None; + } + + if let Some(remaining) = remaining.strip_prefix("r/") { + let (regex, next_pos) = if let Some(end) = remaining.find("/ ") { + (&remaining[..end], self.last_pos + 2 + end + 2) + } else if let Some(remaining) = remaining.strip_suffix('/') { + (remaining, self.query.len()) + } else { + (remaining, self.query.len()) + }; + self.last_pos = next_pos; + Some(QueryToken::Regex(regex)) + } else { + let (mut part, next_pos) = if let Some(sp) = remaining.find(' ') { + (&remaining[..sp], self.last_pos + sp + 1) + } else { + (remaining, self.query.len()) + }; + self.last_pos = next_pos; + + if part == "|" { + return Some(QueryToken::Or); + } + + let mut is_inverse = false; + if let Some(s) = part.strip_prefix('!') { + part = s; + is_inverse = true; + } + let token = if let Some(s) = part.strip_prefix('^') { + QueryToken::MatchStart(s, is_inverse) + } else if let Some(s) = part.strip_suffix('$') { + QueryToken::MatchEnd(s, is_inverse) + } else if let Some(s) = part.strip_prefix('\'') { + QueryToken::MatchFull(s, is_inverse) + } else { + QueryToken::Match(part, is_inverse) + }; + Some(token) + } + } +} diff --git a/crates/atuin-client/src/encryption.rs b/crates/atuin-client/src/encryption.rs new file mode 100644 index 00000000..f2032482 --- /dev/null +++ b/crates/atuin-client/src/encryption.rs @@ -0,0 +1,440 @@ +// The general idea is that we NEVER send cleartext history to the server +// This way the odds of anything private ending up where it should not are +// very low +// The server authenticates via the usual username and password. This has +// nothing to do with the encryption, and is purely authentication! The client +// generates its own secret key, and encrypts all shell history with libsodium's +// secretbox. The data is then sent to the server, where it is stored. All +// clients must share the secret in order to be able to sync, as it is needed +// to decrypt + +use std::{io::prelude::*, path::PathBuf}; + +use base64::prelude::{BASE64_STANDARD, Engine}; +pub use crypto_secretbox::Key; +use crypto_secretbox::{ + AeadCore, AeadInPlace, KeyInit, XSalsa20Poly1305, + aead::{Nonce, OsRng}, +}; +use eyre::{Context, Result, bail, ensure, eyre}; +use fs_err as fs; +use rmp::{Marker, decode::Bytes}; +use serde::{Deserialize, Serialize}; +use time::{OffsetDateTime, format_description::well_known::Rfc3339, macros::format_description}; + +use crate::{history::History, settings::Settings}; + +#[derive(Debug, Serialize, Deserialize)] +pub struct EncryptedHistory { + pub ciphertext: Vec<u8>, + pub nonce: Nonce<XSalsa20Poly1305>, +} + +pub fn generate_encoded_key() -> Result<(Key, String)> { + let key = XSalsa20Poly1305::generate_key(&mut OsRng); + let encoded = encode_key(&key)?; + + Ok((key, encoded)) +} + +pub fn new_key(settings: &Settings) -> Result<Key> { + let path = settings.key_path.as_str(); + let path = PathBuf::from(path); + + if path.exists() { + bail!("key already exists! cannot overwrite"); + } + + let (key, encoded) = generate_encoded_key()?; + + let mut file = fs::File::create(path)?; + file.write_all(encoded.as_bytes())?; + + Ok(key) +} + +// Loads the secret key, will create + save if it doesn't exist +pub fn load_key(settings: &Settings) -> Result<Key> { + let path = settings.key_path.as_str(); + + let key = if PathBuf::from(path).exists() { + let key = fs_err::read_to_string(path)?; + decode_key(key)? + } else { + new_key(settings)? + }; + + Ok(key) +} + +pub fn encode_key(key: &Key) -> Result<String> { + let mut buf = vec![]; + rmp::encode::write_array_len(&mut buf, key.len() as u32) + .wrap_err("could not encode key to message pack")?; + for b in key { + rmp::encode::write_uint(&mut buf, *b as u64) + .wrap_err("could not encode key to message pack")?; + } + let buf = BASE64_STANDARD.encode(buf); + + Ok(buf) +} + +pub fn decode_key(key: String) -> Result<Key> { + use rmp::decode; + + let buf = BASE64_STANDARD + .decode(key.trim_end()) + .wrap_err("encryption key is not a valid base64 encoding")?; + + // old code wrote the key as a fixed length array of 32 bytes + // new code writes the key with a length prefix + match <[u8; 32]>::try_from(&*buf) { + Ok(key) => Ok(key.into()), + Err(_) => { + let mut bytes = rmp::decode::Bytes::new(&buf); + + match Marker::from_u8(buf[0]) { + Marker::Bin8 => { + let len = decode::read_bin_len(&mut bytes).map_err(|err| eyre!("{err:?}"))?; + ensure!(len == 32, "encryption key is not the correct size"); + let key = <[u8; 32]>::try_from(bytes.remaining_slice()) + .context("could not decode encryption key")?; + Ok(key.into()) + } + Marker::Array16 => { + let len = decode::read_array_len(&mut bytes).map_err(|err| eyre!("{err:?}"))?; + ensure!(len == 32, "encryption key is not the correct size"); + + let mut key = Key::default(); + for i in &mut key { + *i = rmp::decode::read_int(&mut bytes).map_err(|err| eyre!("{err:?}"))?; + } + Ok(key) + } + _ => bail!("could not decode encryption key"), + } + } + } +} + +pub fn encrypt(history: &History, key: &Key) -> Result<EncryptedHistory> { + // serialize with msgpack + let mut buf = encode(history)?; + + let nonce = XSalsa20Poly1305::generate_nonce(&mut OsRng); + XSalsa20Poly1305::new(key) + .encrypt_in_place(&nonce, &[], &mut buf) + .map_err(|_| eyre!("could not encrypt"))?; + + Ok(EncryptedHistory { + ciphertext: buf, + nonce, + }) +} + +pub fn decrypt(mut encrypted_history: EncryptedHistory, key: &Key) -> Result<History> { + XSalsa20Poly1305::new(key) + .decrypt_in_place( + &encrypted_history.nonce, + &[], + &mut encrypted_history.ciphertext, + ) + .map_err(|_| eyre!("could not decrypt history"))?; + let plaintext = encrypted_history.ciphertext; + + let history = decode(&plaintext)?; + + Ok(history) +} + +fn format_rfc3339(ts: OffsetDateTime) -> Result<String> { + // horrible hack. chrono AutoSI limits to 0, 3, 6, or 9 decimal places for nanoseconds. + // time does not have this functionality. + static PARTIAL_RFC3339_0: &[time::format_description::FormatItem<'static>] = + format_description!("[year]-[month]-[day]T[hour]:[minute]:[second]Z"); + static PARTIAL_RFC3339_3: &[time::format_description::FormatItem<'static>] = + format_description!("[year]-[month]-[day]T[hour]:[minute]:[second].[subsecond digits:3]Z"); + static PARTIAL_RFC3339_6: &[time::format_description::FormatItem<'static>] = + format_description!("[year]-[month]-[day]T[hour]:[minute]:[second].[subsecond digits:6]Z"); + static PARTIAL_RFC3339_9: &[time::format_description::FormatItem<'static>] = + format_description!("[year]-[month]-[day]T[hour]:[minute]:[second].[subsecond digits:9]Z"); + + let fmt = match ts.nanosecond() { + 0 => PARTIAL_RFC3339_0, + ns if ns % 1_000_000 == 0 => PARTIAL_RFC3339_3, + ns if ns % 1_000 == 0 => PARTIAL_RFC3339_6, + _ => PARTIAL_RFC3339_9, + }; + + Ok(ts.format(fmt)?) +} + +fn encode(h: &History) -> Result<Vec<u8>> { + use rmp::encode; + + let mut output = vec![]; + // INFO: ensure this is updated when adding new fields + encode::write_array_len(&mut output, 9)?; + + encode::write_str(&mut output, &h.id.0)?; + encode::write_str(&mut output, &(format_rfc3339(h.timestamp)?))?; + encode::write_sint(&mut output, h.duration)?; + encode::write_sint(&mut output, h.exit)?; + encode::write_str(&mut output, &h.command)?; + encode::write_str(&mut output, &h.cwd)?; + encode::write_str(&mut output, &h.session)?; + encode::write_str(&mut output, &h.hostname)?; + match h.deleted_at { + Some(d) => encode::write_str(&mut output, &format_rfc3339(d)?)?, + None => encode::write_nil(&mut output)?, + } + + Ok(output) +} + +fn decode(bytes: &[u8]) -> Result<History> { + use rmp::decode::{self, DecodeStringError}; + + let mut bytes = Bytes::new(bytes); + + let nfields = decode::read_array_len(&mut bytes).map_err(error_report)?; + if nfields < 8 { + bail!("malformed decrypted history") + } + if nfields > 9 { + bail!("cannot decrypt history from a newer version of atuin"); + } + + let bytes = bytes.remaining_slice(); + let (id, bytes) = decode::read_str_from_slice(bytes).map_err(error_report)?; + let (timestamp, bytes) = decode::read_str_from_slice(bytes).map_err(error_report)?; + + let mut bytes = Bytes::new(bytes); + let duration = decode::read_int(&mut bytes).map_err(error_report)?; + let exit = decode::read_int(&mut bytes).map_err(error_report)?; + + let bytes = bytes.remaining_slice(); + let (command, bytes) = decode::read_str_from_slice(bytes).map_err(error_report)?; + let (cwd, bytes) = decode::read_str_from_slice(bytes).map_err(error_report)?; + let (session, bytes) = decode::read_str_from_slice(bytes).map_err(error_report)?; + let (hostname, bytes) = decode::read_str_from_slice(bytes).map_err(error_report)?; + + // if we have more fields, try and get the deleted_at + let mut deleted_at = None; + let mut bytes = bytes; + if nfields > 8 { + bytes = match decode::read_str_from_slice(bytes) { + Ok((d, b)) => { + deleted_at = Some(d); + b + } + // we accept null here + Err(DecodeStringError::TypeMismatch(Marker::Null)) => { + // consume the null marker + let mut c = Bytes::new(bytes); + decode::read_nil(&mut c).map_err(error_report)?; + c.remaining_slice() + } + Err(err) => return Err(error_report(err)), + }; + } + + if !bytes.is_empty() { + bail!("trailing bytes in encoded history. malformed") + } + + Ok(History { + id: id.to_owned().into(), + timestamp: OffsetDateTime::parse(timestamp, &Rfc3339)?, + duration, + exit, + command: command.to_owned(), + cwd: cwd.to_owned(), + session: session.to_owned(), + hostname: hostname.to_owned(), + author: History::author_from_hostname(hostname), + intent: None, + deleted_at: deleted_at + .map(|t| OffsetDateTime::parse(t, &Rfc3339)) + .transpose()?, + }) +} + +fn error_report<E: std::fmt::Debug>(err: E) -> eyre::Report { + eyre!("{err:?}") +} + +#[cfg(test)] +mod test { + use crypto_secretbox::{KeyInit, XSalsa20Poly1305, aead::OsRng}; + use pretty_assertions::assert_eq; + use time::{OffsetDateTime, macros::datetime}; + + use crate::history::History; + + use super::{decode, decrypt, encode, encrypt}; + + #[test] + fn test_encrypt_decrypt() { + let key1 = XSalsa20Poly1305::generate_key(&mut OsRng); + let key2 = XSalsa20Poly1305::generate_key(&mut OsRng); + + let history = History::from_db() + .id("1".into()) + .timestamp(OffsetDateTime::now_utc()) + .command("ls".into()) + .cwd("/home/ellie".into()) + .exit(0) + .duration(1) + .session("beep boop".into()) + .hostname("booop".into()) + .author("booop".into()) + .intent(None) + .deleted_at(None) + .build() + .into(); + + let e1 = encrypt(&history, &key1).unwrap(); + let e2 = encrypt(&history, &key2).unwrap(); + + assert_ne!(e1.ciphertext, e2.ciphertext); + assert_ne!(e1.nonce, e2.nonce); + + // test decryption works + // this should pass + match decrypt(e1, &key1) { + Err(e) => panic!("failed to decrypt, got {e}"), + Ok(h) => assert_eq!(h, history), + }; + + // this should err + let _ = decrypt(e2, &key1).expect_err("expected an error decrypting with invalid key"); + } + + #[test] + fn test_decode() { + let bytes = [ + 0x99, 0xD9, 32, 54, 54, 100, 49, 54, 99, 98, 101, 101, 55, 99, 100, 52, 55, 53, 51, 56, + 101, 53, 99, 53, 98, 56, 98, 52, 52, 101, 57, 48, 48, 54, 101, 187, 50, 48, 50, 51, 45, + 48, 53, 45, 50, 56, 84, 49, 56, 58, 51, 53, 58, 52, 48, 46, 54, 51, 51, 56, 55, 50, 90, + 206, 2, 238, 210, 240, 0, 170, 103, 105, 116, 32, 115, 116, 97, 116, 117, 115, 217, 42, + 47, 85, 115, 101, 114, 115, 47, 99, 111, 110, 114, 97, 100, 46, 108, 117, 100, 103, 97, + 116, 101, 47, 68, 111, 99, 117, 109, 101, 110, 116, 115, 47, 99, 111, 100, 101, 47, 97, + 116, 117, 105, 110, 217, 32, 98, 57, 55, 100, 57, 97, 51, 48, 54, 102, 50, 55, 52, 52, + 55, 51, 97, 50, 48, 51, 100, 50, 101, 98, 97, 52, 49, 102, 57, 52, 53, 55, 187, 102, + 118, 102, 103, 57, 51, 54, 99, 48, 107, 112, 102, 58, 99, 111, 110, 114, 97, 100, 46, + 108, 117, 100, 103, 97, 116, 101, 192, + ]; + let history = History { + id: "66d16cbee7cd47538e5c5b8b44e9006e".to_owned().into(), + timestamp: datetime!(2023-05-28 18:35:40.633872 +00:00), + duration: 49206000, + exit: 0, + command: "git status".to_owned(), + cwd: "/Users/conrad.ludgate/Documents/code/atuin".to_owned(), + session: "b97d9a306f274473a203d2eba41f9457".to_owned(), + hostname: "fvfg936c0kpf:conrad.ludgate".to_owned(), + author: "conrad.ludgate".to_owned(), + intent: None, + deleted_at: None, + }; + + let h = decode(&bytes).unwrap(); + assert_eq!(history, h); + + let b = encode(&h).unwrap(); + assert_eq!(&bytes, &*b); + } + + #[test] + fn test_decode_deleted() { + let history = History { + id: "66d16cbee7cd47538e5c5b8b44e9006e".to_owned().into(), + timestamp: datetime!(2023-05-28 18:35:40.633872 +00:00), + duration: 49206000, + exit: 0, + command: "git status".to_owned(), + cwd: "/Users/conrad.ludgate/Documents/code/atuin".to_owned(), + session: "b97d9a306f274473a203d2eba41f9457".to_owned(), + hostname: "fvfg936c0kpf:conrad.ludgate".to_owned(), + author: "conrad.ludgate".to_owned(), + intent: None, + deleted_at: Some(datetime!(2023-05-28 18:35:40.633872 +00:00)), + }; + + let b = encode(&history).unwrap(); + let h = decode(&b).unwrap(); + assert_eq!(history, h); + } + + #[test] + fn test_decode_old() { + let bytes = [ + 0x98, 0xD9, 32, 54, 54, 100, 49, 54, 99, 98, 101, 101, 55, 99, 100, 52, 55, 53, 51, 56, + 101, 53, 99, 53, 98, 56, 98, 52, 52, 101, 57, 48, 48, 54, 101, 187, 50, 48, 50, 51, 45, + 48, 53, 45, 50, 56, 84, 49, 56, 58, 51, 53, 58, 52, 48, 46, 54, 51, 51, 56, 55, 50, 90, + 206, 2, 238, 210, 240, 0, 170, 103, 105, 116, 32, 115, 116, 97, 116, 117, 115, 217, 42, + 47, 85, 115, 101, 114, 115, 47, 99, 111, 110, 114, 97, 100, 46, 108, 117, 100, 103, 97, + 116, 101, 47, 68, 111, 99, 117, 109, 101, 110, 116, 115, 47, 99, 111, 100, 101, 47, 97, + 116, 117, 105, 110, 217, 32, 98, 57, 55, 100, 57, 97, 51, 48, 54, 102, 50, 55, 52, 52, + 55, 51, 97, 50, 48, 51, 100, 50, 101, 98, 97, 52, 49, 102, 57, 52, 53, 55, 187, 102, + 118, 102, 103, 57, 51, 54, 99, 48, 107, 112, 102, 58, 99, 111, 110, 114, 97, 100, 46, + 108, 117, 100, 103, 97, 116, 101, + ]; + let history = History { + id: "66d16cbee7cd47538e5c5b8b44e9006e".to_owned().into(), + timestamp: datetime!(2023-05-28 18:35:40.633872 +00:00), + duration: 49206000, + exit: 0, + command: "git status".to_owned(), + cwd: "/Users/conrad.ludgate/Documents/code/atuin".to_owned(), + session: "b97d9a306f274473a203d2eba41f9457".to_owned(), + hostname: "fvfg936c0kpf:conrad.ludgate".to_owned(), + author: "conrad.ludgate".to_owned(), + intent: None, + deleted_at: None, + }; + + let h = decode(&bytes).unwrap(); + assert_eq!(history, h); + } + + #[test] + fn key_encodings() { + use super::{Key, decode_key, encode_key}; + + // a history of our key encodings. + // v11.0.0 xCAbWypb0msJ2Kq+8j4GVEWUlDX7deKnrTRSIopuqXxc5Q== + // v12.0.0 xCAbWypb0msJ2Kq+8j4GVEWUlDX7deKnrTRSIopuqXxc5Q== + // v13.0.0 xCAbWypb0msJ2Kq+8j4GVEWUlDX7deKnrTRSIopuqXxc5Q== + // v13.0.1 xCAbWypb0msJ2Kq+8j4GVEWUlDX7deKnrTRSIopuqXxc5Q== + // v14.0.0 xCAbWypb0msJ2Kq+8j4GVEWUlDX7deKnrTRSIopuqXxc5Q== + // v14.0.1 xCAbWypb0msJ2Kq+8j4GVEWUlDX7deKnrTRSIopuqXxc5Q== + // c7d89c1 3AAgG1sqW8zSawnM2MyqzL7M8j4GVEXMlMyUNcz7dczizKfMrTRSIsyKbsypfFzM5Q== (https://github.com/ellie/atuin/pull/805) + // b53ca35 3AAgG1sqW8zSawnM2MyqzL7M8j4GVEXMlMyUNcz7dczizKfMrTRSIsyKbsypfFzM5Q== (https://github.com/ellie/atuin/pull/974) + // v15.0.0 3AAgG1sqW8zSawnM2MyqzL7M8j4GVEXMlMyUNcz7dczizKfMrTRSIsyKbsypfFzM5Q== + // b8b57c8 xCAbWypb0msJ2Kq+8j4GVEWUlDX7deKnrTRSIopuqXxc5Q== (https://github.com/ellie/atuin/pull/1057) + // 8c94d79 3AAgG1sqW8zSawnM2MyqzL7M8j4GVEXMlMyUNcz7dczizKfMrTRSIsyKbsypfFzM5Q== (https://github.com/ellie/atuin/pull/1089) + + let key = Key::from([ + 27, 91, 42, 91, 210, 107, 9, 216, 170, 190, 242, 62, 6, 84, 69, 148, 148, 53, 251, 117, + 226, 167, 173, 52, 82, 34, 138, 110, 169, 124, 92, 229, + ]); + + assert_eq!( + encode_key(&key).unwrap(), + "3AAgG1sqW8zSawnM2MyqzL7M8j4GVEXMlMyUNcz7dczizKfMrTRSIsyKbsypfFzM5Q==" + ); + + // key encodings we have to support + let valid_encodings = [ + "xCAbWypb0msJ2Kq+8j4GVEWUlDX7deKnrTRSIopuqXxc5Q==", + "3AAgG1sqW8zSawnM2MyqzL7M8j4GVEXMlMyUNcz7dczizKfMrTRSIsyKbsypfFzM5Q==", + ]; + + for k in valid_encodings { + assert_eq!(decode_key(k.to_owned()).expect(k), key); + } + } +} diff --git a/crates/atuin-client/src/history.rs b/crates/atuin-client/src/history.rs new file mode 100644 index 00000000..a5adc233 --- /dev/null +++ b/crates/atuin-client/src/history.rs @@ -0,0 +1,722 @@ +use core::fmt::Formatter; +use rmp::decode::DecodeStringError; +use rmp::decode::ValueReadError; +use rmp::{Marker, decode::Bytes}; +use std::env; +use std::fmt::Display; + +use atuin_common::record::DecryptedData; +use atuin_common::utils::uuid_v7; + +use eyre::{Result, bail, eyre}; + +use crate::secrets::SECRET_PATTERNS_RE; +use crate::settings::Settings; +use crate::utils::get_host_user; +use time::OffsetDateTime; + +mod builder; +pub mod store; + +pub(crate) const HISTORY_VERSION_V0: &str = "v0"; +pub(crate) const HISTORY_VERSION_V1: &str = "v1"; +const HISTORY_RECORD_VERSION_V0: u16 = 0; +const HISTORY_RECORD_VERSION_V1: u16 = 1; +pub(crate) const HISTORY_VERSION: &str = HISTORY_VERSION_V1; +pub const HISTORY_TAG: &str = "history"; +const HISTORY_AUTHOR_ENV: &str = "ATUIN_HISTORY_AUTHOR"; +const HISTORY_INTENT_ENV: &str = "ATUIN_HISTORY_INTENT"; + +#[derive(Clone, Debug, Eq, PartialEq, Hash)] +pub struct HistoryId(pub String); + +impl Display for HistoryId { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "{}", self.0) + } +} + +impl From<String> for HistoryId { + fn from(s: String) -> Self { + Self(s) + } +} + +/// Client-side history entry. +/// +/// Client stores data unencrypted, and only encrypts it before sending to the server. +/// +/// To create a new history entry, use one of the builders: +/// - [`History::import()`] to import an entry from the shell history file +/// - [`History::capture()`] to capture an entry via hook +/// - [`History::from_db()`] to create an instance from the database entry +// +// ## Implementation Notes +// +// New fields must be added to `History::{serialize,deserialize}` in a backwards +// compatible way (sensible defaults and careful `nfields` handling). +#[derive(Debug, Clone, PartialEq, Eq, sqlx::FromRow)] +pub struct History { + /// A client-generated ID, used to identify the entry when syncing. + /// + /// Stored as `client_id` in the database. + pub id: HistoryId, + /// When the command was run. + pub timestamp: OffsetDateTime, + /// How long the command took to run. + pub duration: i64, + /// The exit code of the command. + pub exit: i64, + /// The command that was run. + pub command: String, + /// The current working directory when the command was run. + pub cwd: String, + /// The session ID, associated with a terminal session. + pub session: String, + /// The hostname of the machine the command was run on. + pub hostname: String, + /// Who wrote this command (human user or automation/agent identity). + pub author: String, + /// Optional rationale for why the command was executed. + pub intent: Option<String>, + /// Timestamp, which is set when the entry is deleted, allowing a soft delete. + pub deleted_at: Option<OffsetDateTime>, +} + +#[derive(Debug, Clone, PartialEq, Eq, sqlx::FromRow)] +pub struct HistoryStats { + /// The command that was ran after this one in the session + pub next: Option<History>, + /// + /// The command that was ran before this one in the session + pub previous: Option<History>, + + /// How many times has this command been ran? + pub total: u64, + + pub average_duration: u64, + + pub exits: Vec<(i64, i64)>, + + pub day_of_week: Vec<(String, i64)>, + + pub duration_over_time: Vec<(String, i64)>, +} + +impl History { + pub(crate) fn author_from_hostname(hostname: &str) -> String { + hostname + .split_once(':') + .map_or_else(|| hostname.to_owned(), |(_, user)| user.to_owned()) + } + + fn normalize_optional_field(field: Option<String>) -> Option<String> { + field.and_then(|value| { + let trimmed = value.trim(); + if trimmed.is_empty() { + None + } else { + Some(trimmed.to_owned()) + } + }) + } + + #[allow(clippy::too_many_arguments)] + fn new( + timestamp: OffsetDateTime, + command: String, + cwd: String, + exit: i64, + duration: i64, + session: Option<String>, + hostname: Option<String>, + author: Option<String>, + intent: Option<String>, + deleted_at: Option<OffsetDateTime>, + ) -> Self { + let session = session + .or_else(|| env::var("ATUIN_SESSION").ok()) + .unwrap_or_else(|| uuid_v7().as_simple().to_string()); + let hostname = hostname.unwrap_or_else(get_host_user); + let author = Self::normalize_optional_field(author) + .or_else(|| Self::normalize_optional_field(env::var(HISTORY_AUTHOR_ENV).ok())) + .unwrap_or_else(|| Self::author_from_hostname(hostname.as_str())); + let intent = Self::normalize_optional_field(intent) + .or_else(|| Self::normalize_optional_field(env::var(HISTORY_INTENT_ENV).ok())); + + Self { + id: uuid_v7().as_simple().to_string().into(), + timestamp, + command, + cwd, + exit, + duration, + session, + hostname, + author, + intent, + deleted_at, + } + } + + pub fn serialize(&self) -> Result<DecryptedData> { + // This is pretty much the same as what we used for the old history, with one difference - + // it uses integers for timestamps rather than a string format. + + use rmp::encode; + + let mut output = vec![]; + + // write the version + encode::write_u16(&mut output, HISTORY_RECORD_VERSION_V1)?; + let include_intent = self.intent.is_some(); + encode::write_array_len(&mut output, 10 + u32::from(include_intent))?; + + encode::write_str(&mut output, &self.id.0)?; + encode::write_u64(&mut output, self.timestamp.unix_timestamp_nanos() as u64)?; + encode::write_sint(&mut output, self.duration)?; + encode::write_sint(&mut output, self.exit)?; + encode::write_str(&mut output, &self.command)?; + encode::write_str(&mut output, &self.cwd)?; + encode::write_str(&mut output, &self.session)?; + encode::write_str(&mut output, &self.hostname)?; + + match self.deleted_at { + Some(d) => encode::write_u64(&mut output, d.unix_timestamp_nanos() as u64)?, + None => encode::write_nil(&mut output)?, + } + + encode::write_str(&mut output, self.author.as_str())?; + if let Some(intent) = &self.intent { + encode::write_str(&mut output, intent.as_str())?; + } + + Ok(DecryptedData(output)) + } + + fn read_optional_string(bytes: &[u8]) -> Result<(Option<String>, &[u8])> { + use rmp::decode; + + fn error_report<E: std::fmt::Debug>(err: E) -> eyre::Report { + eyre!("{err:?}") + } + + match decode::read_str_from_slice(bytes) { + Ok((value, bytes)) => Ok((Some(value.to_owned()), bytes)), + Err(DecodeStringError::TypeMismatch(Marker::Null)) => { + let mut cursor = Bytes::new(bytes); + decode::read_nil(&mut cursor).map_err(error_report)?; + + Ok((None, cursor.remaining_slice())) + } + Err(err) => Err(error_report(err)), + } + } + + fn deserialize_v0(bytes: &[u8]) -> Result<History> { + use rmp::decode; + + fn error_report<E: std::fmt::Debug>(err: E) -> eyre::Report { + eyre!("{err:?}") + } + + let mut bytes = Bytes::new(bytes); + + let version = decode::read_u16(&mut bytes).map_err(error_report)?; + + if version != HISTORY_RECORD_VERSION_V0 { + bail!("expected decoding v0 record, found v{version}"); + } + + let nfields = decode::read_array_len(&mut bytes).map_err(error_report)?; + + if nfields != 9 { + bail!("cannot decrypt history from a different version of Atuin"); + } + + let bytes = bytes.remaining_slice(); + let (id, bytes) = decode::read_str_from_slice(bytes).map_err(error_report)?; + + let mut bytes = Bytes::new(bytes); + let timestamp = decode::read_u64(&mut bytes).map_err(error_report)?; + let duration = decode::read_int(&mut bytes).map_err(error_report)?; + let exit = decode::read_int(&mut bytes).map_err(error_report)?; + + let bytes = bytes.remaining_slice(); + let (command, bytes) = decode::read_str_from_slice(bytes).map_err(error_report)?; + let (cwd, bytes) = decode::read_str_from_slice(bytes).map_err(error_report)?; + let (session, bytes) = decode::read_str_from_slice(bytes).map_err(error_report)?; + let (hostname, bytes) = decode::read_str_from_slice(bytes).map_err(error_report)?; + + let mut bytes = Bytes::new(bytes); + + let (deleted_at, bytes) = match decode::read_u64(&mut bytes) { + Ok(unix) => (Some(unix), bytes.remaining_slice()), + // we accept null here + Err(ValueReadError::TypeMismatch(Marker::Null)) => (None, bytes.remaining_slice()), + Err(err) => return Err(error_report(err)), + }; + if !bytes.is_empty() { + bail!("trailing bytes in encoded history. malformed") + } + + Ok(History { + id: id.to_owned().into(), + timestamp: OffsetDateTime::from_unix_timestamp_nanos(timestamp as i128)?, + duration, + exit, + command: command.to_owned(), + cwd: cwd.to_owned(), + session: session.to_owned(), + hostname: hostname.to_owned(), + author: Self::author_from_hostname(hostname), + intent: None, + deleted_at: deleted_at + .map(|t| OffsetDateTime::from_unix_timestamp_nanos(t as i128)) + .transpose()?, + }) + } + + fn deserialize_v1(bytes: &[u8]) -> Result<History> { + use rmp::decode; + + fn error_report<E: std::fmt::Debug>(err: E) -> eyre::Report { + eyre!("{err:?}") + } + + let mut bytes = Bytes::new(bytes); + + let version = decode::read_u16(&mut bytes).map_err(error_report)?; + + if version != HISTORY_RECORD_VERSION_V1 { + bail!("expected decoding v1 record, found v{version}"); + } + + let nfields = decode::read_array_len(&mut bytes).map_err(error_report)?; + + if !(10..=11).contains(&nfields) { + bail!("cannot decrypt history from a different version of Atuin"); + } + + let bytes = bytes.remaining_slice(); + let (id, bytes) = decode::read_str_from_slice(bytes).map_err(error_report)?; + + let mut bytes = Bytes::new(bytes); + let timestamp = decode::read_u64(&mut bytes).map_err(error_report)?; + let duration = decode::read_int(&mut bytes).map_err(error_report)?; + let exit = decode::read_int(&mut bytes).map_err(error_report)?; + + let bytes = bytes.remaining_slice(); + let (command, bytes) = decode::read_str_from_slice(bytes).map_err(error_report)?; + let (cwd, bytes) = decode::read_str_from_slice(bytes).map_err(error_report)?; + let (session, bytes) = decode::read_str_from_slice(bytes).map_err(error_report)?; + let (hostname, bytes) = decode::read_str_from_slice(bytes).map_err(error_report)?; + + let mut bytes = Bytes::new(bytes); + + let (deleted_at, bytes) = match decode::read_u64(&mut bytes) { + Ok(unix) => (Some(unix), bytes.remaining_slice()), + // we accept null here + Err(ValueReadError::TypeMismatch(Marker::Null)) => (None, bytes.remaining_slice()), + Err(err) => return Err(error_report(err)), + }; + let (author, bytes) = Self::read_optional_string(bytes)?; + let (intent, bytes) = if nfields > 10 { + Self::read_optional_string(bytes)? + } else { + (None, bytes) + }; + + if !bytes.is_empty() { + bail!("trailing bytes in encoded history. malformed") + } + + Ok(History { + id: id.to_owned().into(), + timestamp: OffsetDateTime::from_unix_timestamp_nanos(timestamp as i128)?, + duration, + exit, + command: command.to_owned(), + cwd: cwd.to_owned(), + session: session.to_owned(), + hostname: hostname.to_owned(), + author: author.unwrap_or_else(|| Self::author_from_hostname(hostname)), + intent, + deleted_at: deleted_at + .map(|t| OffsetDateTime::from_unix_timestamp_nanos(t as i128)) + .transpose()?, + }) + } + + pub fn deserialize(bytes: &[u8], version: &str) -> Result<History> { + match version { + HISTORY_VERSION_V0 => Self::deserialize_v0(bytes), + HISTORY_VERSION_V1 => Self::deserialize_v1(bytes), + + _ => bail!("unknown version {version:?}"), + } + } + + /// Builder for a history entry that is imported from shell history. + /// + /// The only two required fields are `timestamp` and `command`. + /// + /// ## Examples + /// ``` + /// use atuin_client::history::History; + /// + /// let history: History = History::import() + /// .timestamp(time::OffsetDateTime::now_utc()) + /// .command("ls -la") + /// .build() + /// .into(); + /// ``` + /// + /// If shell history contains more information, it can be added to the builder: + /// ``` + /// use atuin_client::history::History; + /// + /// let history: History = History::import() + /// .timestamp(time::OffsetDateTime::now_utc()) + /// .command("ls -la") + /// .cwd("/home/user") + /// .exit(0) + /// .duration(100) + /// .build() + /// .into(); + /// ``` + /// + /// Unknown command or command without timestamp cannot be imported, which + /// is forced at compile time: + /// + /// ```compile_fail + /// use atuin_client::history::History; + /// + /// // this will not compile because timestamp is missing + /// let history: History = History::import() + /// .command("ls -la") + /// .build() + /// .into(); + /// ``` + pub fn import() -> builder::HistoryImportedBuilder { + builder::HistoryImported::builder() + } + + /// Builder for a history entry that is captured via hook. + /// + /// This builder is used only at the `start` step of the hook, + /// so it doesn't have any fields which are known only after + /// the command is finished, such as `exit` or `duration`. + /// + /// ## Examples + /// ```rust + /// use atuin_client::history::History; + /// + /// let history: History = History::capture() + /// .timestamp(time::OffsetDateTime::now_utc()) + /// .command("ls -la") + /// .cwd("/home/user") + /// .build() + /// .into(); + /// ``` + /// + /// Command without any required info cannot be captured, which is forced at compile time: + /// + /// ```compile_fail + /// use atuin_client::history::History; + /// + /// // this will not compile because `cwd` is missing + /// let history: History = History::capture() + /// .timestamp(time::OffsetDateTime::now_utc()) + /// .command("ls -la") + /// .build() + /// .into(); + /// ``` + pub fn capture() -> builder::HistoryCapturedBuilder { + builder::HistoryCaptured::builder() + } + + /// Builder for a history entry that is captured via hook, and sent to the daemon. + /// + /// This builder is used only at the `start` step of the hook, + /// so it doesn't have any fields which are known only after + /// the command is finished, such as `exit` or `duration`. + /// + /// It does, however, include information that can usually be inferred. + /// + /// This is because the daemon we are sending a request to lacks the context of the command + /// + /// ## Examples + /// ```rust + /// use atuin_client::history::History; + /// + /// let history: History = History::daemon() + /// .timestamp(time::OffsetDateTime::now_utc()) + /// .command("ls -la") + /// .cwd("/home/user") + /// .session("018deb6e8287781f9973ef40e0fde76b") + /// .hostname("computer:ellie") + /// .build() + /// .into(); + /// ``` + /// + /// Command without any required info cannot be captured, which is forced at compile time: + /// + /// ```compile_fail + /// use atuin_client::history::History; + /// + /// // this will not compile because `hostname` is missing + /// let history: History = History::daemon() + /// .timestamp(time::OffsetDateTime::now_utc()) + /// .command("ls -la") + /// .cwd("/home/user") + /// .session("018deb6e8287781f9973ef40e0fde76b") + /// .build() + /// .into(); + /// ``` + pub fn daemon() -> builder::HistoryDaemonCaptureBuilder { + builder::HistoryDaemonCapture::builder() + } + + /// Builder for a history entry that is imported from the database. + /// + /// All fields are required, as they are all present in the database. + /// + /// ```compile_fail + /// use atuin_client::history::History; + /// + /// // this will not compile because `id` field is missing + /// let history: History = History::from_db() + /// .timestamp(time::OffsetDateTime::now_utc()) + /// .command("ls -la".to_string()) + /// .cwd("/home/user".to_string()) + /// .exit(0) + /// .duration(100) + /// .session("somesession".to_string()) + /// .hostname("localhost".to_string()) + /// .author("user".to_string()) + /// .intent(None) + /// .deleted_at(None) + /// .build() + /// .into(); + /// ``` + pub fn from_db() -> builder::HistoryFromDbBuilder { + builder::HistoryFromDb::builder() + } + + pub fn success(&self) -> bool { + self.exit == 0 || self.duration == -1 + } + + pub fn should_save(&self, settings: &Settings) -> bool { + !(self.command.starts_with(' ') + || self.command.is_empty() + || settings.history_filter.is_match(&self.command) + || settings.cwd_filter.is_match(&self.cwd) + || (settings.secrets_filter && SECRET_PATTERNS_RE.is_match(&self.command))) + } +} + +#[cfg(test)] +mod tests { + use regex::RegexSet; + use time::macros::datetime; + + use crate::{history::HISTORY_VERSION, settings::Settings}; + + use super::History; + + // Test that we don't save history where necessary + #[test] + fn privacy_test() { + let settings = Settings { + cwd_filter: RegexSet::new(["^/supasecret"]).unwrap(), + history_filter: RegexSet::new(["^psql"]).unwrap(), + ..Settings::utc() + }; + + let normal_command: History = History::capture() + .timestamp(time::OffsetDateTime::now_utc()) + .command("echo foo") + .cwd("/") + .build() + .into(); + + let with_space: History = History::capture() + .timestamp(time::OffsetDateTime::now_utc()) + .command(" echo bar") + .cwd("/") + .build() + .into(); + + let empty: History = History::capture() + .timestamp(time::OffsetDateTime::now_utc()) + .command("") + .cwd("/") + .build() + .into(); + + let stripe_key: History = History::capture() + .timestamp(time::OffsetDateTime::now_utc()) + .command("curl foo.com/bar?key=sk_test_1234567890abcdefghijklmnop") + .cwd("/") + .build() + .into(); + + let secret_dir: History = History::capture() + .timestamp(time::OffsetDateTime::now_utc()) + .command("echo ohno") + .cwd("/supasecret") + .build() + .into(); + + let with_psql: History = History::capture() + .timestamp(time::OffsetDateTime::now_utc()) + .command("psql") + .cwd("/supasecret") + .build() + .into(); + + assert!(normal_command.should_save(&settings)); + assert!(!with_space.should_save(&settings)); + assert!(!empty.should_save(&settings)); + assert!(!stripe_key.should_save(&settings)); + assert!(!secret_dir.should_save(&settings)); + assert!(!with_psql.should_save(&settings)); + } + + #[test] + fn disable_secrets() { + let settings = Settings { + secrets_filter: false, + ..Settings::utc() + }; + + let stripe_key: History = History::capture() + .timestamp(time::OffsetDateTime::now_utc()) + .command("curl foo.com/bar?key=sk_test_1234567890abcdefghijklmnop") + .cwd("/") + .build() + .into(); + + assert!(stripe_key.should_save(&settings)); + } + + #[test] + fn test_serialize_deserialize() { + let history = History { + id: "66d16cbee7cd47538e5c5b8b44e9006e".to_owned().into(), + timestamp: datetime!(2023-05-28 18:35:40.633872 +00:00), + duration: 49206000, + exit: 0, + command: "git status".to_owned(), + cwd: "/Users/conrad.ludgate/Documents/code/atuin".to_owned(), + session: "b97d9a306f274473a203d2eba41f9457".to_owned(), + hostname: "fvfg936c0kpf:conrad.ludgate".to_owned(), + author: "conrad.ludgate".to_owned(), + intent: None, + deleted_at: None, + }; + + let serialized = history.serialize().expect("failed to serialize history"); + assert_eq!( + &serialized.0[0..3], + [205, 0, 1], + "should encode as history v1" + ); + + let deserialized = History::deserialize(&serialized.0, HISTORY_VERSION) + .expect("failed to deserialize history"); + assert_eq!(history, deserialized); + } + + #[test] + fn test_serialize_deserialize_deleted() { + let history = History { + id: "66d16cbee7cd47538e5c5b8b44e9006e".to_owned().into(), + timestamp: datetime!(2023-05-28 18:35:40.633872 +00:00), + duration: 49206000, + exit: 0, + command: "git status".to_owned(), + cwd: "/Users/conrad.ludgate/Documents/code/atuin".to_owned(), + session: "b97d9a306f274473a203d2eba41f9457".to_owned(), + hostname: "fvfg936c0kpf:conrad.ludgate".to_owned(), + author: "conrad.ludgate".to_owned(), + intent: None, + deleted_at: Some(datetime!(2023-11-19 20:18 +00:00)), + }; + + let serialized = history.serialize().expect("failed to serialize history"); + + let deserialized = History::deserialize(&serialized.0, HISTORY_VERSION) + .expect("failed to deserialize history"); + + assert_eq!(history, deserialized); + } + + #[test] + fn test_serialize_deserialize_with_author_and_intent() { + let history = History { + id: "66d16cbee7cd47538e5c5b8b44e9006e".to_owned().into(), + timestamp: datetime!(2023-05-28 18:35:40.633872 +00:00), + duration: 49206000, + exit: 0, + command: "git status".to_owned(), + cwd: "/Users/conrad.ludgate/Documents/code/atuin".to_owned(), + session: "b97d9a306f274473a203d2eba41f9457".to_owned(), + hostname: "fvfg936c0kpf:conrad.ludgate".to_owned(), + author: "claude".to_owned(), + intent: Some("check repository status".to_owned()), + deleted_at: None, + }; + + let serialized = history.serialize().expect("failed to serialize history"); + let deserialized = History::deserialize(&serialized.0, HISTORY_VERSION) + .expect("failed to deserialize history"); + + assert_eq!(history, deserialized); + } + + #[test] + fn test_serialize_deserialize_version() { + // v0 + let bytes_v0 = [ + 205, 0, 0, 153, 217, 32, 54, 54, 100, 49, 54, 99, 98, 101, 101, 55, 99, 100, 52, 55, + 53, 51, 56, 101, 53, 99, 53, 98, 56, 98, 52, 52, 101, 57, 48, 48, 54, 101, 207, 23, 99, + 98, 117, 24, 210, 246, 128, 206, 2, 238, 210, 240, 0, 170, 103, 105, 116, 32, 115, 116, + 97, 116, 117, 115, 217, 42, 47, 85, 115, 101, 114, 115, 47, 99, 111, 110, 114, 97, 100, + 46, 108, 117, 100, 103, 97, 116, 101, 47, 68, 111, 99, 117, 109, 101, 110, 116, 115, + 47, 99, 111, 100, 101, 47, 97, 116, 117, 105, 110, 217, 32, 98, 57, 55, 100, 57, 97, + 51, 48, 54, 102, 50, 55, 52, 52, 55, 51, 97, 50, 48, 51, 100, 50, 101, 98, 97, 52, 49, + 102, 57, 52, 53, 55, 187, 102, 118, 102, 103, 57, 51, 54, 99, 48, 107, 112, 102, 58, + 99, 111, 110, 114, 97, 100, 46, 108, 117, 100, 103, 97, 116, 101, 192, + ]; + + let deserialized = History::deserialize(&bytes_v0, "v0"); + assert!(deserialized.is_ok()); + + let deserialized = History::deserialize(&bytes_v0, HISTORY_VERSION); + assert!(deserialized.is_err()); + + let current = History { + id: "66d16cbee7cd47538e5c5b8b44e9006e".to_owned().into(), + timestamp: datetime!(2023-05-28 18:35:40.633872 +00:00), + duration: 49206000, + exit: 0, + command: "git status".to_owned(), + cwd: "/Users/conrad.ludgate/Documents/code/atuin".to_owned(), + session: "b97d9a306f274473a203d2eba41f9457".to_owned(), + hostname: "fvfg936c0kpf:conrad.ludgate".to_owned(), + author: "conrad.ludgate".to_owned(), + intent: None, + deleted_at: None, + }; + + let bytes_v1 = current.serialize().expect("failed to serialize history"); + let deserialized = History::deserialize(&bytes_v1.0, HISTORY_VERSION); + assert!(deserialized.is_ok()); + + let deserialized = History::deserialize(&bytes_v1.0, "v0"); + assert!(deserialized.is_err()); + } +} diff --git a/crates/atuin-client/src/history/builder.rs b/crates/atuin-client/src/history/builder.rs new file mode 100644 index 00000000..72a505fd --- /dev/null +++ b/crates/atuin-client/src/history/builder.rs @@ -0,0 +1,154 @@ +use typed_builder::TypedBuilder; + +use super::History; + +/// Builder for a history entry that is imported from shell history. +/// +/// The only two required fields are `timestamp` and `command`. +#[derive(Debug, Clone, TypedBuilder)] +pub struct HistoryImported { + timestamp: time::OffsetDateTime, + #[builder(setter(into))] + command: String, + #[builder(default = "unknown".into(), setter(into))] + cwd: String, + #[builder(default = -1)] + exit: i64, + #[builder(default = -1)] + duration: i64, + #[builder(default, setter(strip_option, into))] + session: Option<String>, + #[builder(default, setter(strip_option, into))] + hostname: Option<String>, + #[builder(default, setter(strip_option, into))] + author: Option<String>, + #[builder(default, setter(strip_option, into))] + intent: Option<String>, +} + +impl From<HistoryImported> for History { + fn from(imported: HistoryImported) -> Self { + History::new( + imported.timestamp, + imported.command, + imported.cwd, + imported.exit, + imported.duration, + imported.session, + imported.hostname, + imported.author, + imported.intent, + None, + ) + } +} + +/// Builder for a history entry that is captured via hook. +/// +/// This builder is used only at the `start` step of the hook, +/// so it doesn't have any fields which are known only after +/// the command is finished, such as `exit` or `duration`. +#[derive(Debug, Clone, TypedBuilder)] +pub struct HistoryCaptured { + timestamp: time::OffsetDateTime, + #[builder(setter(into))] + command: String, + #[builder(setter(into))] + cwd: String, + #[builder(default, setter(strip_option, into))] + author: Option<String>, + #[builder(default, setter(strip_option, into))] + intent: Option<String>, +} + +impl From<HistoryCaptured> for History { + fn from(captured: HistoryCaptured) -> Self { + History::new( + captured.timestamp, + captured.command, + captured.cwd, + -1, + -1, + None, + None, + captured.author, + captured.intent, + None, + ) + } +} + +/// Builder for a history entry that is loaded from the database. +/// +/// All fields are required, as they are all present in the database. +#[derive(Debug, Clone, TypedBuilder)] +pub struct HistoryFromDb { + id: String, + timestamp: time::OffsetDateTime, + command: String, + cwd: String, + exit: i64, + duration: i64, + session: String, + hostname: String, + author: String, + intent: Option<String>, + deleted_at: Option<time::OffsetDateTime>, +} + +impl From<HistoryFromDb> for History { + fn from(from_db: HistoryFromDb) -> Self { + History { + id: from_db.id.into(), + timestamp: from_db.timestamp, + exit: from_db.exit, + command: from_db.command, + cwd: from_db.cwd, + duration: from_db.duration, + session: from_db.session, + hostname: from_db.hostname, + author: from_db.author, + intent: from_db.intent, + deleted_at: from_db.deleted_at, + } + } +} + +/// Builder for a history entry that is captured via hook and sent to the daemon +/// +/// This builder is similar to Capture, but we just require more information up front. +/// For the old setup, we could just rely on History::new to read some of the missing +/// data. This is no longer the case. +#[derive(Debug, Clone, TypedBuilder)] +pub struct HistoryDaemonCapture { + timestamp: time::OffsetDateTime, + #[builder(setter(into))] + command: String, + #[builder(setter(into))] + cwd: String, + #[builder(setter(into))] + session: String, + #[builder(setter(into))] + hostname: String, + #[builder(default, setter(strip_option, into))] + author: Option<String>, + #[builder(default, setter(strip_option, into))] + intent: Option<String>, +} + +impl From<HistoryDaemonCapture> for History { + fn from(captured: HistoryDaemonCapture) -> Self { + History::new( + captured.timestamp, + captured.command, + captured.cwd, + -1, + -1, + Some(captured.session), + Some(captured.hostname), + captured.author, + captured.intent, + None, + ) + } +} diff --git a/crates/atuin-client/src/history/store.rs b/crates/atuin-client/src/history/store.rs new file mode 100644 index 00000000..d166564f --- /dev/null +++ b/crates/atuin-client/src/history/store.rs @@ -0,0 +1,420 @@ +use std::{collections::HashSet, fmt::Write, time::Duration}; + +use eyre::{Result, bail, eyre}; +use indicatif::{ProgressBar, ProgressState, ProgressStyle}; +use rmp::decode::Bytes; + +use crate::{ + database::{Database, current_context}, + record::{encryption::PASETO_V4, sqlite_store::SqliteStore, store::Store}, +}; +use atuin_common::record::{DecryptedData, Host, HostId, Record, RecordId, RecordIdx}; + +use super::{HISTORY_TAG, HISTORY_VERSION, HISTORY_VERSION_V0, History, HistoryId}; + +#[derive(Debug, Clone)] +pub struct HistoryStore { + pub store: SqliteStore, + pub host_id: HostId, + pub encryption_key: [u8; 32], +} + +#[derive(Debug, Eq, PartialEq, Clone)] +pub enum HistoryRecord { + Create(History), // Create a history record + Delete(HistoryId), // Delete a history record, identified by ID +} + +impl HistoryRecord { + /// Serialize a history record, returning DecryptedData + /// The record will be of a certain type + /// We map those like so: + /// + /// HistoryRecord::Create -> 0 + /// HistoryRecord::Delete-> 1 + /// + /// This numeric identifier is then written as the first byte to the buffer. For history, we + /// append the serialized history right afterwards, to avoid having to handle serialization + /// twice. + /// + /// Deletion simply refers to the history by ID + pub fn serialize(&self) -> Result<DecryptedData> { + // probably don't actually need to use rmp here, but if we ever need to extend it, it's a + // nice wrapper around raw byte stuff + use rmp::encode; + + let mut output = vec![]; + + match self { + HistoryRecord::Create(history) => { + // 0 -> a history create + encode::write_u8(&mut output, 0)?; + + let bytes = history.serialize()?; + + encode::write_bin(&mut output, &bytes.0)?; + } + HistoryRecord::Delete(id) => { + // 1 -> a history delete + encode::write_u8(&mut output, 1)?; + encode::write_str(&mut output, id.0.as_str())?; + } + }; + + Ok(DecryptedData(output)) + } + + pub fn deserialize(bytes: &DecryptedData, version: &str) -> Result<Self> { + use rmp::decode; + + fn error_report<E: std::fmt::Debug>(err: E) -> eyre::Report { + eyre!("{err:?}") + } + + let mut bytes = Bytes::new(&bytes.0); + + let record_type = decode::read_u8(&mut bytes).map_err(error_report)?; + + match record_type { + // 0 -> HistoryRecord::Create + 0 => { + // not super useful to us atm, but perhaps in the future + // written by write_bin above + let _ = decode::read_bin_len(&mut bytes).map_err(error_report)?; + + let record = History::deserialize(bytes.remaining_slice(), version)?; + + Ok(HistoryRecord::Create(record)) + } + + // 1 -> HistoryRecord::Delete + 1 => { + let bytes = bytes.remaining_slice(); + let (id, bytes) = decode::read_str_from_slice(bytes).map_err(error_report)?; + + if !bytes.is_empty() { + bail!( + "trailing bytes decoding HistoryRecord::Delete - malformed? got {bytes:?}" + ); + } + + Ok(HistoryRecord::Delete(id.to_string().into())) + } + + n => { + bail!("unknown HistoryRecord type {n}") + } + } + } +} + +impl HistoryStore { + pub fn new(store: SqliteStore, host_id: HostId, encryption_key: [u8; 32]) -> Self { + HistoryStore { + store, + host_id, + encryption_key, + } + } + + async fn push_record(&self, record: HistoryRecord) -> Result<(RecordId, RecordIdx)> { + let bytes = record.serialize()?; + let idx = self + .store + .last(self.host_id, HISTORY_TAG) + .await? + .map_or(0, |p| p.idx + 1); + + let record = Record::builder() + .host(Host::new(self.host_id)) + .version(HISTORY_VERSION.to_string()) + .tag(HISTORY_TAG.to_string()) + .idx(idx) + .data(bytes) + .build(); + + let id = record.id; + + self.store + .push(&record.encrypt::<PASETO_V4>(&self.encryption_key)) + .await?; + + Ok((id, idx)) + } + + async fn push_batch(&self, records: impl Iterator<Item = HistoryRecord>) -> Result<()> { + let mut ret = Vec::new(); + + let idx = self + .store + .last(self.host_id, HISTORY_TAG) + .await? + .map_or(0, |p| p.idx + 1); + + // Could probably _also_ do this as an iterator, but let's see how this is for now. + // optimizing for minimal sqlite transactions, this code can be optimised later + for (n, record) in records.enumerate() { + let bytes = record.serialize()?; + + let record = Record::builder() + .host(Host::new(self.host_id)) + .version(HISTORY_VERSION.to_string()) + .tag(HISTORY_TAG.to_string()) + .idx(idx + n as u64) + .data(bytes) + .build(); + + let record = record.encrypt::<PASETO_V4>(&self.encryption_key); + + ret.push(record); + } + + self.store.push_batch(ret.iter()).await?; + + Ok(()) + } + + pub async fn delete(&self, id: HistoryId) -> Result<(RecordId, RecordIdx)> { + let record = HistoryRecord::Delete(id); + + self.push_record(record).await + } + + pub async fn push(&self, history: History) -> Result<(RecordId, RecordIdx)> { + // TODO(ellie): move the history store to its own file + // it's tiny rn so fine as is + let record = HistoryRecord::Create(history); + + self.push_record(record).await + } + + pub async fn history(&self) -> Result<Vec<HistoryRecord>> { + // Atm this loads all history into memory + // Not ideal as that is potentially quite a lot, although history will be small. + let records = self.store.all_tagged(HISTORY_TAG).await?; + let mut ret = Vec::with_capacity(records.len()); + + for record in records.into_iter() { + let hist = match record.version.as_str() { + HISTORY_VERSION_V0 | HISTORY_VERSION => { + let version = record.version.clone(); + let decrypted = record.decrypt::<PASETO_V4>(&self.encryption_key)?; + + HistoryRecord::deserialize(&decrypted.data, version.as_str()) + } + version => bail!("unknown history version {version:?}"), + }?; + + ret.push(hist); + } + + Ok(ret) + } + + pub async fn build(&self, database: &dyn Database) -> Result<()> { + // I'd like to change how we rebuild and not couple this with the database, but need to + // consider the structure more deeply. This will be easy to change. + + // TODO(ellie): page or iterate this + let history = self.history().await?; + + // In theory we could flatten this here + // The current issue is that the database may have history in it already, from the old sync + // This didn't actually delete old history + // If we're sure we have a DB only maintained by the new store, we can flatten + // create/delete before we even get to sqlite + let mut creates = Vec::new(); + let mut deletes = Vec::new(); + + for i in history { + match i { + HistoryRecord::Create(h) => { + creates.push(h); + } + HistoryRecord::Delete(id) => { + deletes.push(id); + } + } + } + + database.save_bulk(&creates).await?; + database.delete_rows(&deletes).await?; + + Ok(()) + } + + pub async fn incremental_build(&self, database: &dyn Database, ids: &[RecordId]) -> Result<()> { + for id in ids { + let record = self.store.get(*id).await; + + let record = match record { + Ok(record) => record, + _ => { + continue; + } + }; + + if record.tag != HISTORY_TAG { + continue; + } + + let version = record.version.clone(); + let decrypted = record.decrypt::<PASETO_V4>(&self.encryption_key)?; + let record = match version.as_str() { + HISTORY_VERSION_V0 | HISTORY_VERSION => { + HistoryRecord::deserialize(&decrypted.data, version.as_str())? + } + version => bail!("unknown history version {version:?}"), + }; + + match record { + HistoryRecord::Create(h) => { + // TODO: benchmark CPU time/memory tradeoff of batch commit vs one at a time + database.save(&h).await?; + } + HistoryRecord::Delete(id) => { + database.delete_rows(&[id]).await?; + } + } + } + + Ok(()) + } + + /// Get a list of history IDs that exist in the store + /// Note: This currently involves loading all history into memory. This is not going to be a + /// large amount in absolute terms, but do not all it in a hot loop. + pub async fn history_ids(&self) -> Result<HashSet<HistoryId>> { + let history = self.history().await?; + + let ret = HashSet::from_iter(history.iter().map(|h| match h { + HistoryRecord::Create(h) => h.id.clone(), + HistoryRecord::Delete(id) => id.clone(), + })); + + Ok(ret) + } + + pub async fn init_store(&self, db: &impl Database) -> Result<()> { + let pb = ProgressBar::new_spinner(); + pb.set_style( + ProgressStyle::with_template("{spinner:.blue} {msg}") + .unwrap() + .with_key("eta", |state: &ProgressState, w: &mut dyn Write| { + write!(w, "{:.1}s", state.eta().as_secs_f64()).unwrap() + }) + .progress_chars("#>-"), + ); + pb.enable_steady_tick(Duration::from_millis(500)); + + pb.set_message("Fetching history from old database"); + + let context = current_context().await?; + let history = db.list(&[], &context, None, false, true).await?; + + pb.set_message("Fetching history already in store"); + let store_ids = self.history_ids().await?; + + pb.set_message("Converting old history to new store"); + let mut records = Vec::new(); + + for i in history { + debug!("loaded {}", i.id); + + if store_ids.contains(&i.id) { + debug!("skipping {} - already exists", i.id); + continue; + } + + if i.deleted_at.is_some() { + records.push(HistoryRecord::Delete(i.id)); + } else { + records.push(HistoryRecord::Create(i)); + } + } + + pb.set_message("Writing to db"); + + if !records.is_empty() { + self.push_batch(records.into_iter()).await?; + } + + pb.finish_with_message("Import complete"); + + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use atuin_common::record::DecryptedData; + use time::macros::datetime; + + use crate::history::{HISTORY_VERSION, store::HistoryRecord}; + + use super::History; + + #[test] + fn test_serialize_deserialize_create() { + let bytes = [ + 204, 0, 196, 147, 205, 0, 1, 154, 217, 32, 48, 49, 56, 99, 100, 52, 102, 101, 56, 49, + 55, 53, 55, 99, 100, 50, 97, 101, 101, 54, 53, 99, 100, 55, 56, 54, 49, 102, 57, 99, + 56, 49, 207, 23, 166, 251, 212, 181, 82, 0, 0, 100, 0, 162, 108, 115, 217, 41, 47, 85, + 115, 101, 114, 115, 47, 101, 108, 108, 105, 101, 47, 115, 114, 99, 47, 103, 105, 116, + 104, 117, 98, 46, 99, 111, 109, 47, 97, 116, 117, 105, 110, 115, 104, 47, 97, 116, 117, + 105, 110, 217, 32, 48, 49, 56, 99, 100, 52, 102, 101, 97, 100, 56, 57, 55, 53, 57, 55, + 56, 53, 50, 53, 50, 55, 97, 51, 49, 99, 57, 57, 56, 48, 53, 57, 170, 98, 111, 111, 112, + 58, 101, 108, 108, 105, 101, 192, 165, 101, 108, 108, 105, 101, + ]; + + let history = History { + id: "018cd4fe81757cd2aee65cd7861f9c81".to_owned().into(), + timestamp: datetime!(2024-01-04 00:00:00.000000 +00:00), + duration: 100, + exit: 0, + command: "ls".to_owned(), + cwd: "/Users/ellie/src/github.com/atuinsh/atuin".to_owned(), + session: "018cd4fead897597852527a31c998059".to_owned(), + hostname: "boop:ellie".to_owned(), + author: "ellie".to_owned(), + intent: None, + deleted_at: None, + }; + + let record = HistoryRecord::Create(history); + + let serialized = record.serialize().expect("failed to serialize history"); + assert_eq!(serialized.0, bytes); + + let deserialized = HistoryRecord::deserialize(&serialized, HISTORY_VERSION) + .expect("failed to deserialize HistoryRecord"); + assert_eq!(deserialized, record); + + // check the snapshot too + let deserialized = + HistoryRecord::deserialize(&DecryptedData(Vec::from(bytes)), HISTORY_VERSION) + .expect("failed to deserialize HistoryRecord"); + assert_eq!(deserialized, record); + } + + #[test] + fn test_serialize_deserialize_delete() { + let bytes = [ + 204, 1, 217, 32, 48, 49, 56, 99, 100, 52, 102, 101, 56, 49, 55, 53, 55, 99, 100, 50, + 97, 101, 101, 54, 53, 99, 100, 55, 56, 54, 49, 102, 57, 99, 56, 49, + ]; + let record = HistoryRecord::Delete("018cd4fe81757cd2aee65cd7861f9c81".to_string().into()); + + let serialized = record.serialize().expect("failed to serialize history"); + assert_eq!(serialized.0, bytes); + + let deserialized = HistoryRecord::deserialize(&serialized, HISTORY_VERSION) + .expect("failed to deserialize HistoryRecord"); + assert_eq!(deserialized, record); + + let deserialized = + HistoryRecord::deserialize(&DecryptedData(Vec::from(bytes)), HISTORY_VERSION) + .expect("failed to deserialize HistoryRecord"); + assert_eq!(deserialized, record); + } +} diff --git a/crates/atuin-client/src/hub.rs b/crates/atuin-client/src/hub.rs new file mode 100644 index 00000000..2e40aad4 --- /dev/null +++ b/crates/atuin-client/src/hub.rs @@ -0,0 +1,304 @@ +//! Hub authentication support for Atuin +//! +//! This module provides programmatic access to the Atuin Hub authentication flow. +//! It can be used by other crates (like atuin-ai) to authenticate with the Hub +//! and obtain session tokens. +//! +//! Hub authentication is separate from sync authentication - users can have both +//! a sync session (for history sync) and a hub session (for Hub-specific features +//! like AI). + +use std::time::Duration; + +use eyre::{Context, Result, bail}; +use reqwest::{StatusCode, Url, header::USER_AGENT}; + +use atuin_common::{ + api::{ + ATUIN_CARGO_VERSION, ATUIN_HEADER_VERSION, CliCodeResponse, CliVerifyResponse, + ErrorResponse, + }, + tls::ensure_crypto_provider, +}; + +use crate::settings::Settings; + +static APP_USER_AGENT: &str = concat!("atuin/", env!("CARGO_PKG_VERSION")); + +/// The result of starting a hub authentication flow +#[derive(Debug, Clone)] +pub struct HubAuthSession { + /// The code to be verified + pub code: String, + /// The URL the user should visit to authenticate + pub auth_url: String, + /// The hub address being used + pub hub_address: String, +} + +/// The result of polling for hub auth completion +#[derive(Debug, Clone)] +pub enum HubAuthStatus { + /// Still waiting for user authorization + Pending, + /// Authorization complete, contains the session token + Complete(String), + /// Authorization failed with an error + Failed(String), +} + +/// Default poll interval for checking auth status +pub const DEFAULT_POLL_INTERVAL: Duration = Duration::from_secs(2); + +/// Default timeout for the entire auth flow +pub const DEFAULT_AUTH_TIMEOUT: Duration = Duration::from_secs(600); + +impl HubAuthSession { + /// Start a new hub authentication session + /// + /// Returns a session containing the code and auth URL that the user should visit. + pub async fn start(hub_address: &str) -> Result<Self> { + debug!("Starting Hub authentication process..."); + + let hub_address = hub_address.trim_end_matches('/'); + let code_response = request_code(hub_address) + .await + .context("Failed to request authentication code from Hub")?; + + debug!("Received code from Hub"); + + let code = code_response.code; + let auth_url = format!("{}/auth/cli?code={}", hub_address, code); + + Ok(Self { + code, + auth_url, + hub_address: hub_address.to_string(), + }) + } + + /// Poll for the authentication status + /// + /// Returns the current status of the authentication flow. + pub async fn poll(&self) -> Result<HubAuthStatus> { + match verify_code(&self.hub_address, &self.code).await { + Ok(response) => { + if let Some(token) = response.token { + debug!("Authentication complete, received token"); + Ok(HubAuthStatus::Complete(token)) + } else if let Some(error) = response.error { + error!("Authentication failed: {}", error); + Ok(HubAuthStatus::Failed(error)) + } else { + Ok(HubAuthStatus::Pending) + } + } + Err(e) => { + // Transient errors shouldn't fail the whole flow + log::debug!("Verification poll failed: {}", e); + Ok(HubAuthStatus::Pending) + } + } + } + + /// Poll until completion or timeout + /// + /// This is a convenience method that polls repeatedly until the auth completes + /// or times out. + pub async fn wait_for_completion( + &self, + timeout: Duration, + poll_interval: Duration, + ) -> Result<String> { + let start = std::time::Instant::now(); + + debug!("Polling for Hub authentication completion..."); + + loop { + if start.elapsed() > timeout { + warn!("Authentication loop exited due to timeout"); + bail!("Authentication timed out. Please try again."); + } + + match self.poll().await? { + HubAuthStatus::Complete(token) => return Ok(token), + HubAuthStatus::Failed(error) => bail!("Authentication failed: {}", error), + HubAuthStatus::Pending => { + tokio::time::sleep(poll_interval).await; + } + } + } + } +} + +/// Save a hub session token +/// +/// This saves the token to the meta store so it can be used for subsequent Hub API calls. +/// Note: This is separate from the sync session token. +pub async fn save_session(token: &str) -> Result<()> { + Settings::meta_store() + .await? + .save_hub_session(token) + .await + .context("Failed to save hub session") +} + +/// Delete the hub session token (logout from Hub) +pub async fn delete_session() -> Result<()> { + Settings::meta_store() + .await? + .delete_hub_session() + .await + .context("Failed to delete hub session") +} + +/// Check if the user is logged in with Hub authentication +/// +/// Returns true if the user has a valid Hub session token. +/// This is independent of whether they have a sync session. +pub async fn is_logged_in() -> Result<bool> { + Settings::meta_store().await?.hub_logged_in().await +} + +/// Get the hub session token if available +/// +/// Returns the Hub session token if the user is logged in with Hub auth, +/// or None if not logged in. +pub async fn get_session_token() -> Result<Option<String>> { + Settings::meta_store().await?.hub_session_token().await +} + +/// Link an existing CLI sync account to the current Hub user. +/// +/// This associates the CLI's sync records with the Hub account, enabling +/// unified authentication. After linking: +/// - The Hub token can be used for sync operations +/// - Records are migrated to be accessible via Hub auth +/// +/// Requires: +/// - A valid Hub session (user must be logged in to Hub) +/// - A valid CLI session token to link +/// +/// Returns Ok(()) on success, or an error if: +/// - Not logged in to Hub +/// - CLI token is invalid +/// - CLI account is already linked to a different Hub account +pub async fn link_account(hub_address: &str, cli_token: &str) -> Result<()> { + let hub_token = get_session_token() + .await? + .ok_or_else(|| eyre::eyre!("Not logged in to Hub - cannot link account"))?; + + let url = make_url(hub_address, "/api/v0/account/link")?; + + debug!("Linking CLI account to Hub at {}", hub_address); + + ensure_crypto_provider(); + let client = reqwest::Client::new(); + + let resp = client + .post(&url) + .header(USER_AGENT, APP_USER_AGENT) + .header(ATUIN_HEADER_VERSION, ATUIN_CARGO_VERSION) + .bearer_auth(&hub_token) + .json(&serde_json::json!({ "token": cli_token })) + .send() + .await?; + + let status = resp.status(); + + if status == StatusCode::CONFLICT { + // 409 means CLI account is already linked to a (possibly different) Hub account + debug!("CLI account already linked to a Hub account"); + return Ok(()); + } + + handle_resp_error(resp).await?; + + info!("Successfully linked CLI account to Hub"); + Ok(()) +} + +// --- Internal HTTP functions --- + +fn make_url(address: &str, path: &str) -> Result<String> { + let address = if address.ends_with('/') { + address.to_string() + } else { + format!("{address}/") + }; + + let path = path.strip_prefix('/').unwrap_or(path); + + let url = Url::parse(&address) + .context("failed to parse hub address")? + .join(path) + .context("failed to join hub URL path")?; + + Ok(url.to_string()) +} + +async fn handle_resp_error(resp: reqwest::Response) -> Result<reqwest::Response> { + let status = resp.status(); + + if status == StatusCode::SERVICE_UNAVAILABLE { + error!("Service unavailable: check https://status.atuin.sh"); + bail!("Service unavailable: check https://status.atuin.sh"); + } + + if status == StatusCode::TOO_MANY_REQUESTS { + error!("Rate limited; please wait before trying again"); + bail!("Rate limited; please wait before trying again"); + } + + if !status.is_success() { + if let Ok(error) = resp.json::<ErrorResponse>().await { + error!("Hub error: {} - {}", status, error.reason); + bail!("Hub error: {} - {}", status, error.reason); + } + error!("Hub request failed with status: {}", status); + bail!("Hub request failed with status: {}", status); + } + + Ok(resp) +} + +/// Request a CLI auth code from the Atuin Hub +async fn request_code(address: &str) -> Result<CliCodeResponse> { + ensure_crypto_provider(); + let url = make_url(address, "/auth/cli/code")?; + let client = reqwest::Client::new(); + + debug!("Requesting code from Hub at {url}"); + + let resp = client + .post(&url) + .header(USER_AGENT, APP_USER_AGENT) + .header(ATUIN_HEADER_VERSION, ATUIN_CARGO_VERSION) + .send() + .await?; + let resp = handle_resp_error(resp).await?; + + let code_response = resp.json::<CliCodeResponse>().await?; + Ok(code_response) +} + +/// Poll to verify the CLI auth code and get the session token +async fn verify_code(address: &str, code: &str) -> Result<CliVerifyResponse> { + ensure_crypto_provider(); + let base = make_url(address, "/auth/cli/verify")?; + let url = format!("{base}?code={code}"); + let client = reqwest::Client::new(); + + debug!("Verifying code with Hub at {base}?code=******"); + + let resp = client + .post(&url) + .header(USER_AGENT, APP_USER_AGENT) + .header(ATUIN_HEADER_VERSION, ATUIN_CARGO_VERSION) + .send() + .await?; + let resp = handle_resp_error(resp).await?; + + let verify_response = resp.json::<CliVerifyResponse>().await?; + Ok(verify_response) +} diff --git a/crates/atuin-client/src/import/bash.rs b/crates/atuin-client/src/import/bash.rs new file mode 100644 index 00000000..99a44a58 --- /dev/null +++ b/crates/atuin-client/src/import/bash.rs @@ -0,0 +1,220 @@ +use std::{path::PathBuf, str}; + +use async_trait::async_trait; +use directories::UserDirs; +use eyre::{Result, eyre}; +use itertools::Itertools; +use time::{Duration, OffsetDateTime}; + +use super::{Importer, Loader, get_histfile_path, unix_byte_lines}; +use crate::history::History; +use crate::import::read_to_end; + +#[derive(Debug)] +pub struct Bash { + bytes: Vec<u8>, +} + +fn default_histpath() -> Result<PathBuf> { + let user_dirs = UserDirs::new().ok_or_else(|| eyre!("could not find user directories"))?; + let home_dir = user_dirs.home_dir(); + + Ok(home_dir.join(".bash_history")) +} + +#[async_trait] +impl Importer for Bash { + const NAME: &'static str = "bash"; + + async fn new() -> Result<Self> { + let bytes = read_to_end(get_histfile_path(default_histpath)?)?; + Ok(Self { bytes }) + } + + async fn entries(&mut self) -> Result<usize> { + let count = unix_byte_lines(&self.bytes) + .map(LineType::from) + .filter(|line| matches!(line, LineType::Command(_))) + .count(); + Ok(count) + } + + async fn load(self, h: &mut impl Loader) -> Result<()> { + let lines = unix_byte_lines(&self.bytes) + .map(LineType::from) + .filter(|line| !matches!(line, LineType::NotUtf8)) // invalid utf8 are ignored + .collect_vec(); + + let (commands_before_first_timestamp, first_timestamp) = lines + .iter() + .enumerate() + .find_map(|(i, line)| match line { + LineType::Timestamp(t) => Some((i, *t)), + _ => None, + }) + // if no known timestamps, use now as base + .unwrap_or((lines.len(), OffsetDateTime::now_utc())); + + // if no timestamp is recorded, then use this increment to set an arbitrary timestamp + // to preserve ordering + // this increment is deliberately very small to prevent particularly fast fingers + // causing ordering issues; it also helps in handling the "here document" syntax, + // where several lines are recorded in succession without individual timestamps + let timestamp_increment = Duration::milliseconds(1); + + // make sure there is a minimum amount of time before the first known timestamp + // to fit all commands, given the default increment + let mut next_timestamp = + first_timestamp - timestamp_increment * commands_before_first_timestamp as i32; + + for line in lines.into_iter() { + match line { + LineType::NotUtf8 => unreachable!(), // already filtered + LineType::Empty => {} // do nothing + LineType::Timestamp(t) => { + if t < next_timestamp { + warn!( + "Time reversal detected in Bash history! Commands may be ordered incorrectly." + ); + } + next_timestamp = t; + } + LineType::Command(c) => { + let imported = History::import().timestamp(next_timestamp).command(c); + + h.push(imported.build().into()).await?; + next_timestamp += timestamp_increment; + } + } + } + + Ok(()) + } +} + +#[derive(Debug, Clone)] +enum LineType<'a> { + NotUtf8, + /// Can happen when using the "here document" syntax. + Empty, + /// A timestamp line start with a '#', followed immediately by an integer + /// that represents seconds since UNIX epoch. + Timestamp(OffsetDateTime), + /// Anything else. + Command(&'a str), +} +impl<'a> From<&'a [u8]> for LineType<'a> { + fn from(bytes: &'a [u8]) -> Self { + let Ok(line) = str::from_utf8(bytes) else { + return LineType::NotUtf8; + }; + if line.is_empty() { + return LineType::Empty; + } + + match try_parse_line_as_timestamp(line) { + Some(time) => LineType::Timestamp(time), + None => LineType::Command(line), + } + } +} + +fn try_parse_line_as_timestamp(line: &str) -> Option<OffsetDateTime> { + let seconds = line.strip_prefix('#')?.parse().ok()?; + OffsetDateTime::from_unix_timestamp(seconds).ok() +} + +#[cfg(test)] +mod test { + use std::cmp::Ordering; + + use itertools::{Itertools, assert_equal}; + + use crate::import::{Importer, tests::TestLoader}; + + use super::Bash; + + #[tokio::test] + async fn parse_no_timestamps() { + let bytes = r"cargo install atuin +cargo update +cargo :b̷i̶t̴r̵o̴t̴ ̵i̷s̴ ̷r̶e̵a̸l̷ +" + .as_bytes() + .to_owned(); + + let mut bash = Bash { bytes }; + assert_eq!(bash.entries().await.unwrap(), 3); + + let mut loader = TestLoader::default(); + bash.load(&mut loader).await.unwrap(); + + assert_equal( + loader.buf.iter().map(|h| h.command.as_str()), + [ + "cargo install atuin", + "cargo update", + "cargo :b̷i̶t̴r̵o̴t̴ ̵i̷s̴ ̷r̶e̵a̸l̷", + ], + ); + assert!(is_strictly_sorted(loader.buf.iter().map(|h| h.timestamp))) + } + + #[tokio::test] + async fn parse_with_timestamps() { + let bytes = b"#1672918999 +git reset +#1672919006 +git clean -dxf +#1672919020 +cd ../ +" + .to_vec(); + + let mut bash = Bash { bytes }; + assert_eq!(bash.entries().await.unwrap(), 3); + + let mut loader = TestLoader::default(); + bash.load(&mut loader).await.unwrap(); + + assert_equal( + loader.buf.iter().map(|h| h.command.as_str()), + ["git reset", "git clean -dxf", "cd ../"], + ); + assert_equal( + loader.buf.iter().map(|h| h.timestamp.unix_timestamp()), + [1672918999, 1672919006, 1672919020], + ) + } + + #[tokio::test] + async fn parse_with_partial_timestamps() { + let bytes = b"git reset +#1672919006 +git clean -dxf +cd ../ +" + .to_vec(); + + let mut bash = Bash { bytes }; + assert_eq!(bash.entries().await.unwrap(), 3); + + let mut loader = TestLoader::default(); + bash.load(&mut loader).await.unwrap(); + + assert_equal( + loader.buf.iter().map(|h| h.command.as_str()), + ["git reset", "git clean -dxf", "cd ../"], + ); + assert!(is_strictly_sorted(loader.buf.iter().map(|h| h.timestamp))) + } + + fn is_strictly_sorted<T>(iter: impl IntoIterator<Item = T>) -> bool + where + T: Clone + PartialOrd, + { + iter.into_iter() + .tuple_windows() + .all(|(a, b)| matches!(a.partial_cmp(&b), Some(Ordering::Less))) + } +} diff --git a/crates/atuin-client/src/import/fish.rs b/crates/atuin-client/src/import/fish.rs new file mode 100644 index 00000000..9fcf624c --- /dev/null +++ b/crates/atuin-client/src/import/fish.rs @@ -0,0 +1,179 @@ +// import old shell history! +// automatically hoover up all that we can find + +use std::path::PathBuf; + +use async_trait::async_trait; +use directories::BaseDirs; +use eyre::{Result, eyre}; +use time::OffsetDateTime; + +use super::{Importer, Loader, unix_byte_lines}; +use crate::history::History; +use crate::import::read_to_end; + +#[derive(Debug)] +pub struct Fish { + bytes: Vec<u8>, +} + +/// see https://fishshell.com/docs/current/interactive.html#searchable-command-history +fn default_histpath() -> Result<PathBuf> { + let base = BaseDirs::new().ok_or_else(|| eyre!("could not determine data directory"))?; + let data = std::env::var("XDG_DATA_HOME").map_or_else( + |_| base.home_dir().join(".local").join("share"), + PathBuf::from, + ); + + // fish supports multiple history sessions + // If `fish_history` var is missing, or set to `default`, use `fish` as the session + let session = std::env::var("fish_history").unwrap_or_else(|_| String::from("fish")); + let session = if session == "default" { + String::from("fish") + } else { + session + }; + + let mut histpath = data.join("fish"); + histpath.push(format!("{session}_history")); + + if histpath.exists() { + Ok(histpath) + } else { + Err(eyre!("Could not find history file.")) + } +} + +#[async_trait] +impl Importer for Fish { + const NAME: &'static str = "fish"; + + async fn new() -> Result<Self> { + let bytes = read_to_end(default_histpath()?)?; + Ok(Self { bytes }) + } + + async fn entries(&mut self) -> Result<usize> { + Ok(super::count_lines(&self.bytes)) + } + + async fn load(self, loader: &mut impl Loader) -> Result<()> { + let now = OffsetDateTime::now_utc(); + let mut time: Option<OffsetDateTime> = None; + let mut cmd: Option<String> = None; + + for b in unix_byte_lines(&self.bytes) { + let s = match std::str::from_utf8(b) { + Ok(s) => s, + Err(_) => continue, // we can skip past things like invalid utf8 + }; + + if let Some(c) = s.strip_prefix("- cmd: ") { + // first, we must deal with the prev cmd + if let Some(cmd) = cmd.take() { + let time = time.unwrap_or(now); + let entry = History::import().timestamp(time).command(cmd); + + loader.push(entry.build().into()).await?; + } + + // using raw strings to avoid needing escaping. + // replaces double backslashes with single backslashes + let c = c.replace(r"\\", r"\"); + // replaces escaped newlines + let c = c.replace(r"\n", "\n"); + // TODO: any other escape characters? + + cmd = Some(c); + } else if let Some(t) = s.strip_prefix(" when: ") { + // if t is not an int, just ignore this line + if let Ok(t) = t.parse::<i64>() { + time = Some(OffsetDateTime::from_unix_timestamp(t)?); + } + } else { + // ... ignore paths lines + } + } + + // we might have a trailing cmd + if let Some(cmd) = cmd.take() { + let time = time.unwrap_or(now); + let entry = History::import().timestamp(time).command(cmd); + + loader.push(entry.build().into()).await?; + } + + Ok(()) + } +} + +#[cfg(test)] +mod test { + + use crate::import::{Importer, tests::TestLoader}; + + use super::Fish; + + #[tokio::test] + async fn parse_complex() { + // complicated input with varying contents and escaped strings. + let bytes = r#"- cmd: history --help + when: 1639162832 +- cmd: cat ~/.bash_history + when: 1639162851 + paths: + - ~/.bash_history +- cmd: ls ~/.local/share/fish/fish_history + when: 1639162890 + paths: + - ~/.local/share/fish/fish_history +- cmd: cat ~/.local/share/fish/fish_history + when: 1639162893 + paths: + - ~/.local/share/fish/fish_history +ERROR +- CORRUPTED: ENTRY + CONTINUE: + - AS + - NORMAL +- cmd: echo "foo" \\\n'bar' baz + when: 1639162933 +- cmd: cat ~/.local/share/fish/fish_history + when: 1639162939 + paths: + - ~/.local/share/fish/fish_history +- cmd: echo "\\"" \\\\ "\\\\" + when: 1639163063 +- cmd: cat ~/.local/share/fish/fish_history + when: 1639163066 + paths: + - ~/.local/share/fish/fish_history +"# + .as_bytes() + .to_owned(); + + let fish = Fish { bytes }; + + let mut loader = TestLoader::default(); + fish.load(&mut loader).await.unwrap(); + let mut history = loader.buf.into_iter(); + + // simple wrapper for fish history entry + macro_rules! fishtory { + ($timestamp:expr_2021, $command:expr_2021) => { + let h = history.next().expect("missing entry in history"); + assert_eq!(h.command.as_str(), $command); + assert_eq!(h.timestamp.unix_timestamp(), $timestamp); + }; + } + + fishtory!(1639162832, "history --help"); + fishtory!(1639162851, "cat ~/.bash_history"); + fishtory!(1639162890, "ls ~/.local/share/fish/fish_history"); + fishtory!(1639162893, "cat ~/.local/share/fish/fish_history"); + fishtory!(1639162933, "echo \"foo\" \\\n'bar' baz"); + fishtory!(1639162939, "cat ~/.local/share/fish/fish_history"); + fishtory!(1639163063, r#"echo "\"" \\ "\\""#); + fishtory!(1639163066, "cat ~/.local/share/fish/fish_history"); + } +} diff --git a/crates/atuin-client/src/import/mod.rs b/crates/atuin-client/src/import/mod.rs new file mode 100644 index 00000000..4a1c6af6 --- /dev/null +++ b/crates/atuin-client/src/import/mod.rs @@ -0,0 +1,140 @@ +use std::fs::File; +use std::io::Read; +use std::path::PathBuf; + +use async_trait::async_trait; +use eyre::{Result, bail}; +use memchr::Memchr; + +use crate::history::History; + +pub mod bash; +pub mod fish; +pub mod nu; +pub mod nu_histdb; +pub mod powershell; +pub mod replxx; +pub mod resh; +pub mod xonsh; +pub mod xonsh_sqlite; +pub mod zsh; +pub mod zsh_histdb; + +#[async_trait] +pub trait Importer: Sized { + const NAME: &'static str; + async fn new() -> Result<Self>; + async fn entries(&mut self) -> Result<usize>; + async fn load(self, loader: &mut impl Loader) -> Result<()>; +} + +#[async_trait] +pub trait Loader: Sync + Send { + async fn push(&mut self, hist: History) -> eyre::Result<()>; +} + +fn unix_byte_lines(input: &[u8]) -> impl Iterator<Item = &[u8]> { + UnixByteLines { + iter: memchr::memchr_iter(b'\n', input), + bytes: input, + i: 0, + } +} + +struct UnixByteLines<'a> { + iter: Memchr<'a>, + bytes: &'a [u8], + i: usize, +} + +impl<'a> Iterator for UnixByteLines<'a> { + type Item = &'a [u8]; + + fn next(&mut self) -> Option<Self::Item> { + let j = self.iter.next()?; + let out = &self.bytes[self.i..j]; + self.i = j + 1; + Some(out) + } + + fn count(self) -> usize + where + Self: Sized, + { + self.iter.count() + } +} + +fn count_lines(input: &[u8]) -> usize { + unix_byte_lines(input).count() +} + +fn get_histpath<D>(def: D) -> Result<PathBuf> +where + D: FnOnce() -> Result<PathBuf>, +{ + if let Ok(p) = std::env::var("HISTFILE") { + Ok(PathBuf::from(p)) + } else { + def() + } +} + +fn get_histfile_path<D>(def: D) -> Result<PathBuf> +where + D: FnOnce() -> Result<PathBuf>, +{ + get_histpath(def).and_then(is_file) +} + +fn get_histdir_path<D>(def: D) -> Result<PathBuf> +where + D: FnOnce() -> Result<PathBuf>, +{ + get_histpath(def).and_then(is_dir) +} + +fn read_to_end(path: PathBuf) -> Result<Vec<u8>> { + let mut bytes = Vec::new(); + let mut f = File::open(path)?; + f.read_to_end(&mut bytes)?; + Ok(bytes) +} +fn is_file(p: PathBuf) -> Result<PathBuf> { + if p.is_file() { + Ok(p) + } else { + bail!( + "Could not find history file {:?}. Try setting and exporting $HISTFILE", + p + ) + } +} +fn is_dir(p: PathBuf) -> Result<PathBuf> { + if p.is_dir() { + Ok(p) + } else { + bail!( + "Could not find history directory {:?}. Try setting and exporting $HISTFILE", + p + ) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[derive(Default)] + pub struct TestLoader { + pub buf: Vec<History>, + } + + #[async_trait] + impl Loader for TestLoader { + async fn push(&mut self, hist: History) -> Result<()> { + self.buf.push(hist); + Ok(()) + } + } +} diff --git a/crates/atuin-client/src/import/nu.rs b/crates/atuin-client/src/import/nu.rs new file mode 100644 index 00000000..cae90ac4 --- /dev/null +++ b/crates/atuin-client/src/import/nu.rs @@ -0,0 +1,67 @@ +// import old shell history! +// automatically hoover up all that we can find + +use std::path::PathBuf; + +use async_trait::async_trait; +use directories::BaseDirs; +use eyre::{Result, eyre}; +use time::OffsetDateTime; + +use super::{Importer, Loader, unix_byte_lines}; +use crate::history::History; +use crate::import::read_to_end; + +#[derive(Debug)] +pub struct Nu { + bytes: Vec<u8>, +} + +fn get_histpath() -> Result<PathBuf> { + let base = BaseDirs::new().ok_or_else(|| eyre!("could not determine data directory"))?; + let config_dir = base.config_dir().join("nushell"); + + let histpath = config_dir.join("history.txt"); + if histpath.exists() { + Ok(histpath) + } else { + Err(eyre!("Could not find history file.")) + } +} + +#[async_trait] +impl Importer for Nu { + const NAME: &'static str = "nu"; + + async fn new() -> Result<Self> { + let bytes = read_to_end(get_histpath()?)?; + Ok(Self { bytes }) + } + + async fn entries(&mut self) -> Result<usize> { + Ok(super::count_lines(&self.bytes)) + } + + async fn load(self, h: &mut impl Loader) -> Result<()> { + let now = OffsetDateTime::now_utc(); + + let mut counter = 0; + for b in unix_byte_lines(&self.bytes) { + let s = match std::str::from_utf8(b) { + Ok(s) => s, + Err(_) => continue, // we can skip past things like invalid utf8 + }; + + let cmd: String = s.replace("<\\n>", "\n"); + + let offset = time::Duration::nanoseconds(counter); + counter += 1; + + let entry = History::import().timestamp(now - offset).command(cmd); + + h.push(entry.build().into()).await?; + } + + Ok(()) + } +} diff --git a/crates/atuin-client/src/import/nu_histdb.rs b/crates/atuin-client/src/import/nu_histdb.rs new file mode 100644 index 00000000..a13cb2b4 --- /dev/null +++ b/crates/atuin-client/src/import/nu_histdb.rs @@ -0,0 +1,113 @@ +// import old shell history! +// automatically hoover up all that we can find + +use std::path::PathBuf; + +use async_trait::async_trait; +use directories::BaseDirs; +use eyre::{Result, eyre}; +use sqlx::{Pool, sqlite::SqlitePool}; +use time::{Duration, OffsetDateTime}; + +use super::Importer; +use crate::history::History; +use crate::import::Loader; + +#[derive(sqlx::FromRow, Debug)] +pub struct HistDbEntry { + pub id: i64, + pub command_line: Vec<u8>, + pub start_timestamp: i64, + pub session_id: i64, + pub hostname: Vec<u8>, + pub cwd: Vec<u8>, + pub duration_ms: i64, + pub exit_status: i64, + pub more_info: Vec<u8>, +} + +impl From<HistDbEntry> for History { + fn from(histdb_item: HistDbEntry) -> Self { + let ts_secs = histdb_item.start_timestamp / 1000; + let ts_ns = (histdb_item.start_timestamp % 1000) * 1_000_000; + let imported = History::import() + .timestamp( + OffsetDateTime::from_unix_timestamp(ts_secs).unwrap() + + Duration::nanoseconds(ts_ns), + ) + .command(String::from_utf8(histdb_item.command_line).unwrap()) + .cwd(String::from_utf8(histdb_item.cwd).unwrap()) + .exit(histdb_item.exit_status) + .duration(histdb_item.duration_ms) + .session(format!("{:x}", histdb_item.session_id)) + .hostname(String::from_utf8(histdb_item.hostname).unwrap()); + + imported.build().into() + } +} + +#[derive(Debug)] +pub struct NuHistDb { + histdb: Vec<HistDbEntry>, +} + +/// Read db at given file, return vector of entries. +async fn hist_from_db(dbpath: PathBuf) -> Result<Vec<HistDbEntry>> { + let pool = SqlitePool::connect(dbpath.to_str().unwrap()).await?; + hist_from_db_conn(pool).await +} + +async fn hist_from_db_conn(pool: Pool<sqlx::Sqlite>) -> Result<Vec<HistDbEntry>> { + let query = r#" + SELECT + id, command_line, start_timestamp, session_id, hostname, cwd, duration_ms, exit_status, + more_info + FROM history + ORDER BY start_timestamp + "#; + let histdb_vec: Vec<HistDbEntry> = sqlx::query_as::<_, HistDbEntry>(query) + .fetch_all(&pool) + .await?; + Ok(histdb_vec) +} + +impl NuHistDb { + pub fn histpath() -> Result<PathBuf> { + let base = BaseDirs::new().ok_or_else(|| eyre!("could not determine data directory"))?; + let config_dir = base.config_dir().join("nushell"); + + let histdb_path = config_dir.join("history.sqlite3"); + if histdb_path.exists() { + Ok(histdb_path) + } else { + Err(eyre!("Could not find history file.")) + } + } +} + +#[async_trait] +impl Importer for NuHistDb { + // Not sure how this is used + const NAME: &'static str = "nu_histdb"; + + /// Creates a new NuHistDb and populates the history based on the pre-populated data + /// structure. + async fn new() -> Result<Self> { + let dbpath = NuHistDb::histpath()?; + let histdb_entry_vec = hist_from_db(dbpath).await?; + Ok(Self { + histdb: histdb_entry_vec, + }) + } + + async fn entries(&mut self) -> Result<usize> { + Ok(self.histdb.len()) + } + + async fn load(self, h: &mut impl Loader) -> Result<()> { + for i in self.histdb { + h.push(i.into()).await?; + } + Ok(()) + } +} diff --git a/crates/atuin-client/src/import/powershell.rs b/crates/atuin-client/src/import/powershell.rs new file mode 100644 index 00000000..86fd007d --- /dev/null +++ b/crates/atuin-client/src/import/powershell.rs @@ -0,0 +1,202 @@ +use async_trait::async_trait; +use directories::BaseDirs; +use eyre::{Result, eyre}; +use std::path::PathBuf; +use time::{Duration, OffsetDateTime}; + +use super::{Importer, Loader, count_lines, unix_byte_lines}; +use crate::history::History; +use crate::import::read_to_end; + +#[derive(Debug)] +pub struct PowerShell { + bytes: Vec<u8>, + line_count: Option<usize>, +} + +fn get_history_path() -> Result<PathBuf> { + let base = BaseDirs::new().ok_or_else(|| eyre!("could not determine data directory"))?; + + // The command line history in PowerShell is maintained by the PSReadLine module: + // https://learn.microsoft.com/en-us/powershell/module/psreadline/about/about_psreadline#command-history + // + // > PSReadLine maintains a history file containing all the commands and data you've entered from the command line. + // > The history files are a file named `$($Host.Name)_history.txt`. + // > On Windows systems the history file is stored at `$Env:APPDATA\Microsoft\Windows\PowerShell\PSReadLine`. + // > On non-Windows systems, the history files are stored at `$Env:XDG_DATA_HOME/powershell/PSReadLine` + // > or `$Env:HOME/.local/share/powershell/PSReadLine`. + + let dir = if cfg!(windows) { + base.data_dir() + .join("Microsoft") + .join("Windows") + .join("PowerShell") + .join("PSReadLine") + } else { + std::env::var("XDG_DATA_HOME") + .map_or_else( + |_| base.home_dir().join(".local").join("share"), + PathBuf::from, + ) + .join("powershell") + .join("PSReadLine") + }; + + // The history is stored in a file named `$($Host.Name)_history.txt`. + // For the default console host shipped by Microsoft,`$Host.Name` is `ConsoleHost`: + // https://learn.microsoft.com/en-us/dotnet/api/system.management.automation.host.pshost.name#remarks + + let file = dir.join("ConsoleHost_history.txt"); + + if file.is_file() { + Ok(file) + } else { + Err(eyre!("Could not find history file: {}", file.display())) + } +} + +#[async_trait] +impl Importer for PowerShell { + const NAME: &'static str = "PowerShell"; + + async fn new() -> Result<Self> { + let bytes = read_to_end(get_history_path()?)?; + Ok(Self { + bytes, + line_count: None, + }) + } + + async fn entries(&mut self) -> Result<usize> { + // Commands can be split over multiple lines, + // but this is only used for a progress bar, and multi-line commands + // should be quite rare, so this is not an issue in practice. + if self.line_count.is_none() { + self.line_count = Some(count_lines(&self.bytes)); + } + Ok(self.line_count.unwrap()) + } + + async fn load(mut self, h: &mut impl Loader) -> Result<()> { + let line_count = self.entries().await?; + let start = OffsetDateTime::now_utc() - Duration::milliseconds(line_count as i64); + + let mut counter = 0; + let mut iter = unix_byte_lines(&self.bytes); + + while let Some(s) = iter.next() { + let Ok(s) = read_line(s) else { + continue; // We can skip past things like invalid utf8 + }; + + let mut cmd = s.to_string(); + + // Multi-line commands end with a backtick, append the following lines. + while cmd.ends_with('`') { + cmd.pop(); + + let Some(next) = iter.next() else { + break; + }; + let Ok(next) = read_line(next) else { + break; + }; + + cmd.push('\n'); + cmd.push_str(next); + } + + if cmd.is_empty() { + continue; + } + + let offset = Duration::milliseconds(counter); + counter += 1; + + let entry = History::import().timestamp(start + offset).command(cmd); + h.push(entry.build().into()).await?; + } + + Ok(()) + } +} + +fn read_line(s: &[u8]) -> Result<&str> { + let s = str::from_utf8(s)?; + + // History is stored in CRLF on Windows, normalize the input to LF on all platforms. + let s = s.strip_suffix('\r').unwrap_or(s); + + Ok(s) +} + +#[cfg(test)] +mod test { + use super::*; + use crate::import::tests::TestLoader; + use itertools::assert_equal; + + const INPUT: &str = r#"cargo install atuin +cargo update +echo "first line` +second line` +` +last line" +echo foo + +echo bar +echo baz +"#; + + const EXPECTED: &[&str] = &[ + "cargo install atuin", + "cargo update", + "echo \"first line\nsecond line\n\nlast line\"", + "echo foo", + "echo bar", + "echo baz", + ]; + + #[tokio::test] + async fn test_import() { + let loader = import(INPUT).await; + + let actual = loader.buf.iter().map(|h| h.command.clone()); + let expected = EXPECTED.iter().map(|s| s.to_string()); + + assert_equal(actual, expected); + } + + #[tokio::test] + async fn test_crlf() { + let input = INPUT.replace("\n", "\r\n"); + let loader = import(input.as_str()).await; + + let actual = loader.buf.iter().map(|h| h.command.clone()); + let expected = EXPECTED.iter().map(|s| s.to_string()); + + assert_equal(actual, expected); + } + + #[tokio::test] + async fn test_timestamps() { + let loader = import(INPUT).await; + + let mut prev = loader.buf.first().unwrap().timestamp; + for current in loader.buf.iter().skip(1).map(|h| h.timestamp) { + assert!(current > prev); + prev = current; + } + } + + async fn import(input: &str) -> TestLoader { + let powershell = PowerShell { + bytes: input.as_bytes().to_vec(), + line_count: None, + }; + + let mut loader = TestLoader::default(); + powershell.load(&mut loader).await.unwrap(); + loader + } +} diff --git a/crates/atuin-client/src/import/replxx.rs b/crates/atuin-client/src/import/replxx.rs new file mode 100644 index 00000000..47d566cf --- /dev/null +++ b/crates/atuin-client/src/import/replxx.rs @@ -0,0 +1,137 @@ +use std::{path::PathBuf, str}; + +use async_trait::async_trait; +use directories::UserDirs; +use eyre::{Result, eyre}; +use time::{OffsetDateTime, PrimitiveDateTime, macros::format_description}; + +use super::{Importer, Loader, get_histfile_path, unix_byte_lines}; +use crate::history::History; +use crate::import::read_to_end; + +#[derive(Debug)] +pub struct Replxx { + bytes: Vec<u8>, +} + +fn default_histpath() -> Result<PathBuf> { + let user_dirs = UserDirs::new().ok_or_else(|| eyre!("could not find user directories"))?; + let home_dir = user_dirs.home_dir(); + + // There is no default histfile for replxx. + // Here we try a couple of common names. + let mut candidates = ["replxx_history.txt", ".histfile"].iter(); + loop { + match candidates.next() { + Some(candidate) => { + let histpath = home_dir.join(candidate); + if histpath.exists() { + break Ok(histpath); + } + } + None => { + break Err(eyre!( + "Could not find history file. Try setting and exporting $HISTFILE" + )); + } + } + } +} + +#[async_trait] +impl Importer for Replxx { + const NAME: &'static str = "replxx"; + + async fn new() -> Result<Self> { + let bytes = read_to_end(get_histfile_path(default_histpath)?)?; + Ok(Self { bytes }) + } + + async fn entries(&mut self) -> Result<usize> { + Ok(super::count_lines(&self.bytes) / 2) + } + + async fn load(self, h: &mut impl Loader) -> Result<()> { + let mut timestamp = OffsetDateTime::UNIX_EPOCH; + + for b in unix_byte_lines(&self.bytes) { + let s = std::str::from_utf8(b)?; + match try_parse_line_as_timestamp(s) { + Some(t) => timestamp = t, + None => { + // replxx uses ETB character (0x17) as line breaker + let cmd = s.replace('\u{0017}', "\n"); + let imported = History::import().timestamp(timestamp).command(cmd); + + h.push(imported.build().into()).await?; + } + } + } + + Ok(()) + } +} + +fn try_parse_line_as_timestamp(line: &str) -> Option<OffsetDateTime> { + // replxx history date time format: ### yyyy-mm-dd hh:mm:ss.xxx + let date_time_str = line.strip_prefix("### ")?; + let format = + format_description!("[year]-[month]-[day] [hour]:[minute]:[second].[subsecond digits:3]"); + + let primitive_date_time = PrimitiveDateTime::parse(date_time_str, format).ok()?; + // There is no safe way to get local time offset. + // For simplicity let's just assume UTC. + Some(primitive_date_time.assume_utc()) +} + +#[cfg(test)] +mod test { + + use crate::import::{Importer, tests::TestLoader}; + + use super::Replxx; + + #[tokio::test] + async fn parse_complex() { + let bytes = r#"### 2024-02-10 22:16:28.302 +select * from remote('127.0.0.1:20222', view(select 1)) +### 2024-02-10 22:16:36.919 +select * from numbers(10) +### 2024-02-10 22:16:41.710 +select * from system.numbers +### 2024-02-10 22:19:28.655 +select 1 +### 2024-02-22 11:15:33.046 +CREATE TABLE test( stamp DateTime('UTC'))ENGINE = MergeTreePARTITION BY toDate(stamp)order by tuple() as select toDateTime('2020-01-01')+number*60 from numbers(80000); +"# + .as_bytes() + .to_owned(); + + let replxx = Replxx { bytes }; + + let mut loader = TestLoader::default(); + replxx.load(&mut loader).await.unwrap(); + let mut history = loader.buf.into_iter(); + + // simple wrapper for replxx history entry + macro_rules! history { + ($timestamp:expr_2021, $command:expr_2021) => { + let h = history.next().expect("missing entry in history"); + assert_eq!(h.command.as_str(), $command); + assert_eq!(h.timestamp.unix_timestamp(), $timestamp); + }; + } + + history!( + 1707603388, + "select * from remote('127.0.0.1:20222', view(select 1))" + ); + history!(1707603396, "select * from numbers(10)"); + history!(1707603401, "select * from system.numbers"); + history!(1707603568, "select 1"); + history!( + 1708600533, + "CREATE TABLE test\n( stamp DateTime('UTC'))\nENGINE = MergeTree\nPARTITION BY toDate(stamp)\norder by tuple() as select toDateTime('2020-01-01')+number*60 from numbers(80000);" + ); + } +} diff --git a/crates/atuin-client/src/import/resh.rs b/crates/atuin-client/src/import/resh.rs new file mode 100644 index 00000000..de02d041 --- /dev/null +++ b/crates/atuin-client/src/import/resh.rs @@ -0,0 +1,140 @@ +use std::path::PathBuf; + +use async_trait::async_trait; +use directories::UserDirs; +use eyre::{Result, eyre}; +use serde::Deserialize; + +use atuin_common::utils::uuid_v7; +use time::OffsetDateTime; + +use super::{Importer, Loader, get_histfile_path, unix_byte_lines}; +use crate::history::History; +use crate::import::read_to_end; + +#[derive(Deserialize, Debug)] +#[serde(rename_all = "camelCase")] +pub struct ReshEntry { + pub cmd_line: String, + pub exit_code: i64, + pub shell: String, + pub uname: String, + pub session_id: String, + pub home: String, + pub lang: String, + pub lc_all: String, + pub login: String, + pub pwd: String, + pub pwd_after: String, + pub shell_env: String, + pub term: String, + pub real_pwd: String, + pub real_pwd_after: String, + pub pid: i64, + pub session_pid: i64, + pub host: String, + pub hosttype: String, + pub ostype: String, + pub machtype: String, + pub shlvl: i64, + pub timezone_before: String, + pub timezone_after: String, + pub realtime_before: f64, + pub realtime_after: f64, + pub realtime_before_local: f64, + pub realtime_after_local: f64, + pub realtime_duration: f64, + pub realtime_since_session_start: f64, + pub realtime_since_boot: f64, + pub git_dir: String, + pub git_real_dir: String, + pub git_origin_remote: String, + pub git_dir_after: String, + pub git_real_dir_after: String, + pub git_origin_remote_after: String, + pub machine_id: String, + pub os_release_id: String, + pub os_release_version_id: String, + pub os_release_id_like: String, + pub os_release_name: String, + pub os_release_pretty_name: String, + pub resh_uuid: String, + pub resh_version: String, + pub resh_revision: String, + pub parts_merged: bool, + pub recalled: bool, + pub recall_last_cmd_line: String, + pub cols: String, + pub lines: String, +} + +#[derive(Debug)] +pub struct Resh { + bytes: Vec<u8>, +} + +fn default_histpath() -> Result<PathBuf> { + let user_dirs = UserDirs::new().ok_or_else(|| eyre!("could not find user directories"))?; + let home_dir = user_dirs.home_dir(); + + Ok(home_dir.join(".resh_history.json")) +} + +#[async_trait] +impl Importer for Resh { + const NAME: &'static str = "resh"; + + async fn new() -> Result<Self> { + let bytes = read_to_end(get_histfile_path(default_histpath)?)?; + Ok(Self { bytes }) + } + + async fn entries(&mut self) -> Result<usize> { + Ok(super::count_lines(&self.bytes)) + } + + async fn load(self, h: &mut impl Loader) -> Result<()> { + for b in unix_byte_lines(&self.bytes) { + let s = match std::str::from_utf8(b) { + Ok(s) => s, + Err(_) => continue, // we can skip past things like invalid utf8 + }; + let entry = match serde_json::from_str::<ReshEntry>(s) { + Ok(e) => e, + Err(_) => continue, // skip invalid json :shrug: + }; + + #[allow(clippy::cast_possible_truncation)] + #[allow(clippy::cast_sign_loss)] + let timestamp = { + let secs = entry.realtime_before.floor() as i64; + let nanosecs = (entry.realtime_before.fract() * 1_000_000_000_f64).round() as i64; + OffsetDateTime::from_unix_timestamp(secs)? + time::Duration::nanoseconds(nanosecs) + }; + #[allow(clippy::cast_possible_truncation)] + #[allow(clippy::cast_sign_loss)] + let duration = { + let secs = entry.realtime_after.floor() as i64; + let nanosecs = (entry.realtime_after.fract() * 1_000_000_000_f64).round() as i64; + let base = OffsetDateTime::from_unix_timestamp(secs)? + + time::Duration::nanoseconds(nanosecs); + let difference = base - timestamp; + difference.whole_nanoseconds() as i64 + }; + + let imported = History::import() + .command(entry.cmd_line) + .timestamp(timestamp) + .duration(duration) + .exit(entry.exit_code) + .cwd(entry.pwd) + .hostname(entry.host) + // CHECK: should we add uuid here? It's not set in the other importers + .session(uuid_v7().as_simple().to_string()); + + h.push(imported.build().into()).await?; + } + + Ok(()) + } +} diff --git a/crates/atuin-client/src/import/xonsh.rs b/crates/atuin-client/src/import/xonsh.rs new file mode 100644 index 00000000..6f38de68 --- /dev/null +++ b/crates/atuin-client/src/import/xonsh.rs @@ -0,0 +1,234 @@ +use std::env; +use std::fs::{self, File}; +use std::path::{Path, PathBuf}; + +use async_trait::async_trait; +use directories::BaseDirs; +use eyre::{Result, eyre}; +use serde::Deserialize; +use time::OffsetDateTime; +use uuid::Uuid; +use uuid::timestamp::{Timestamp, context::NoContext}; + +use super::{Importer, Loader, get_histdir_path}; +use crate::history::History; +use crate::utils::get_host_user; + +// Note: both HistoryFile and HistoryData have other keys present in the JSON, we don't +// care about them so we leave them unspecified so as to avoid deserializing unnecessarily. +#[derive(Debug, Deserialize)] +struct HistoryFile { + data: HistoryData, +} + +#[derive(Debug, Deserialize)] +struct HistoryData { + sessionid: String, + cmds: Vec<HistoryCmd>, +} + +#[derive(Debug, Deserialize)] +struct HistoryCmd { + cwd: String, + inp: String, + rtn: Option<i64>, + ts: (f64, f64), +} + +#[derive(Debug)] +pub struct Xonsh { + // history is stored as a bunch of json files, one per session + sessions: Vec<HistoryData>, + hostname: String, +} + +fn xonsh_hist_dir(xonsh_data_dir: Option<String>) -> Result<PathBuf> { + // if running within xonsh, this will be available + if let Some(d) = xonsh_data_dir { + let mut path = PathBuf::from(d); + path.push("history_json"); + return Ok(path); + } + + // otherwise, fall back to default + let base = BaseDirs::new().ok_or_else(|| eyre!("Could not determine home directory"))?; + + let hist_dir = base.data_dir().join("xonsh/history_json"); + if hist_dir.exists() || cfg!(test) { + Ok(hist_dir) + } else { + Err(eyre!("Could not find xonsh history files")) + } +} + +fn load_sessions(hist_dir: &Path) -> Result<Vec<HistoryData>> { + let mut sessions = vec![]; + for entry in fs::read_dir(hist_dir)? { + let p = entry?.path(); + let ext = p.extension().and_then(|e| e.to_str()); + if p.is_file() + && ext == Some("json") + && let Some(data) = load_session(&p)? + { + sessions.push(data); + } + } + Ok(sessions) +} + +fn load_session(path: &Path) -> Result<Option<HistoryData>> { + let file = File::open(path)?; + // empty files are not valid json, so we can't deserialize them + if file.metadata()?.len() == 0 { + return Ok(None); + } + + let mut hist_file: HistoryFile = serde_json::from_reader(file)?; + + // if there are commands in this session, replace the existing UUIDv4 + // with a UUIDv7 generated from the timestamp of the first command + if let Some(cmd) = hist_file.data.cmds.first() { + let seconds = cmd.ts.0.trunc() as u64; + let nanos = (cmd.ts.0.fract() * 1_000_000_000_f64) as u32; + let ts = Timestamp::from_unix(NoContext, seconds, nanos); + hist_file.data.sessionid = Uuid::new_v7(ts).to_string(); + } + Ok(Some(hist_file.data)) +} + +#[async_trait] +impl Importer for Xonsh { + const NAME: &'static str = "xonsh"; + + async fn new() -> Result<Self> { + // wrap xonsh-specific path resolver in general one so that it respects $HISTPATH + let xonsh_data_dir = env::var("XONSH_DATA_DIR").ok(); + let hist_dir = get_histdir_path(|| xonsh_hist_dir(xonsh_data_dir))?; + let sessions = load_sessions(&hist_dir)?; + let hostname = get_host_user(); + Ok(Xonsh { sessions, hostname }) + } + + async fn entries(&mut self) -> Result<usize> { + let total = self.sessions.iter().map(|s| s.cmds.len()).sum(); + Ok(total) + } + + async fn load(self, loader: &mut impl Loader) -> Result<()> { + for session in self.sessions { + for cmd in session.cmds { + let (start, end) = cmd.ts; + let ts_nanos = (start * 1_000_000_000_f64) as i128; + let timestamp = OffsetDateTime::from_unix_timestamp_nanos(ts_nanos)?; + + let duration = (end - start) * 1_000_000_000_f64; + + match cmd.rtn { + Some(exit) => { + let entry = History::import() + .timestamp(timestamp) + .duration(duration.trunc() as i64) + .exit(exit) + .command(cmd.inp.trim()) + .cwd(cmd.cwd) + .session(session.sessionid.clone()) + .hostname(self.hostname.clone()); + loader.push(entry.build().into()).await?; + } + None => { + let entry = History::import() + .timestamp(timestamp) + .duration(duration.trunc() as i64) + .command(cmd.inp.trim()) + .cwd(cmd.cwd) + .session(session.sessionid.clone()) + .hostname(self.hostname.clone()); + loader.push(entry.build().into()).await?; + } + } + } + } + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use time::macros::datetime; + + use super::*; + + use crate::history::History; + use crate::import::tests::TestLoader; + + #[test] + fn test_hist_dir_xonsh() { + let hist_dir = xonsh_hist_dir(Some("/home/user/xonsh_data".to_string())).unwrap(); + assert_eq!( + hist_dir, + PathBuf::from("/home/user/xonsh_data/history_json") + ); + } + + #[tokio::test] + async fn test_import() { + let dir = PathBuf::from("tests/data/xonsh"); + let sessions = load_sessions(&dir).unwrap(); + let hostname = "box:user".to_string(); + let xonsh = Xonsh { sessions, hostname }; + + let mut loader = TestLoader::default(); + xonsh.load(&mut loader).await.unwrap(); + // order in buf will depend on filenames, so sort by timestamp for consistency + loader.buf.sort_by_key(|h| h.timestamp); + for (actual, expected) in loader.buf.iter().zip(expected_hist_entries().iter()) { + assert_eq!(actual.timestamp, expected.timestamp); + assert_eq!(actual.command, expected.command); + assert_eq!(actual.cwd, expected.cwd); + assert_eq!(actual.exit, expected.exit); + assert_eq!(actual.duration, expected.duration); + assert_eq!(actual.hostname, expected.hostname); + } + } + + fn expected_hist_entries() -> [History; 4] { + [ + History::import() + .timestamp(datetime!(2024-02-6 04:17:59.478272256 +00:00:00)) + .command("echo hello world!".to_string()) + .cwd("/home/user/Documents/code/atuin".to_string()) + .exit(0) + .duration(4651069) + .hostname("box:user".to_string()) + .build() + .into(), + History::import() + .timestamp(datetime!(2024-02-06 04:18:01.70632832 +00:00:00)) + .command("ls -l".to_string()) + .cwd("/home/user/Documents/code/atuin".to_string()) + .exit(0) + .duration(21288633) + .hostname("box:user".to_string()) + .build() + .into(), + History::import() + .timestamp(datetime!(2024-02-06 17:41:31.142515968 +00:00:00)) + .command("false".to_string()) + .cwd("/home/user/Documents/code/atuin/atuin-client".to_string()) + .exit(1) + .duration(10269403) + .hostname("box:user".to_string()) + .build() + .into(), + History::import() + .timestamp(datetime!(2024-02-06 17:41:32.271584 +00:00:00)) + .command("exit".to_string()) + .cwd("/home/user/Documents/code/atuin/atuin-client".to_string()) + .exit(0) + .duration(4259347) + .hostname("box:user".to_string()) + .build() + .into(), + ] + } +} diff --git a/crates/atuin-client/src/import/xonsh_sqlite.rs b/crates/atuin-client/src/import/xonsh_sqlite.rs new file mode 100644 index 00000000..7d50ac84 --- /dev/null +++ b/crates/atuin-client/src/import/xonsh_sqlite.rs @@ -0,0 +1,217 @@ +use std::env; +use std::path::PathBuf; + +use async_trait::async_trait; +use directories::BaseDirs; +use eyre::{Result, eyre}; +use futures::TryStreamExt; +use sqlx::{FromRow, Row, sqlite::SqlitePool}; +use time::OffsetDateTime; +use uuid::Uuid; +use uuid::timestamp::{Timestamp, context::NoContext}; + +use super::{Importer, Loader, get_histfile_path}; +use crate::history::History; +use crate::utils::get_host_user; + +#[derive(Debug, FromRow)] +struct HistDbEntry { + inp: String, + rtn: Option<i64>, + tsb: f64, + tse: f64, + cwd: String, + session_start: f64, +} + +impl HistDbEntry { + fn into_hist_with_hostname(self, hostname: String) -> History { + let ts_nanos = (self.tsb * 1_000_000_000_f64) as i128; + let timestamp = OffsetDateTime::from_unix_timestamp_nanos(ts_nanos).unwrap(); + + let session_ts_seconds = self.session_start.trunc() as u64; + let session_ts_nanos = (self.session_start.fract() * 1_000_000_000_f64) as u32; + let session_ts = Timestamp::from_unix(NoContext, session_ts_seconds, session_ts_nanos); + let session_id = Uuid::new_v7(session_ts).to_string(); + let duration = (self.tse - self.tsb) * 1_000_000_000_f64; + + if let Some(exit) = self.rtn { + let imported = History::import() + .timestamp(timestamp) + .duration(duration.trunc() as i64) + .exit(exit) + .command(self.inp) + .cwd(self.cwd) + .session(session_id) + .hostname(hostname); + imported.build().into() + } else { + let imported = History::import() + .timestamp(timestamp) + .duration(duration.trunc() as i64) + .command(self.inp) + .cwd(self.cwd) + .session(session_id) + .hostname(hostname); + imported.build().into() + } + } +} + +fn xonsh_db_path(xonsh_data_dir: Option<String>) -> Result<PathBuf> { + // if running within xonsh, this will be available + if let Some(d) = xonsh_data_dir { + let mut path = PathBuf::from(d); + path.push("xonsh-history.sqlite"); + return Ok(path); + } + + // otherwise, fall back to default + let base = BaseDirs::new().ok_or_else(|| eyre!("Could not determine home directory"))?; + + let hist_file = base.data_dir().join("xonsh/xonsh-history.sqlite"); + if hist_file.exists() || cfg!(test) { + Ok(hist_file) + } else { + Err(eyre!( + "Could not find xonsh history db at: {}", + hist_file.to_string_lossy() + )) + } +} + +#[derive(Debug)] +pub struct XonshSqlite { + pool: SqlitePool, + hostname: String, +} + +#[async_trait] +impl Importer for XonshSqlite { + const NAME: &'static str = "xonsh_sqlite"; + + async fn new() -> Result<Self> { + // wrap xonsh-specific path resolver in general one so that it respects $HISTPATH + let xonsh_data_dir = env::var("XONSH_DATA_DIR").ok(); + let db_path = get_histfile_path(|| xonsh_db_path(xonsh_data_dir))?; + let connection_str = db_path.to_str().ok_or_else(|| { + eyre!( + "Invalid path for SQLite database: {}", + db_path.to_string_lossy() + ) + })?; + + let pool = SqlitePool::connect(connection_str).await?; + let hostname = get_host_user(); + Ok(XonshSqlite { pool, hostname }) + } + + async fn entries(&mut self) -> Result<usize> { + let query = "SELECT COUNT(*) FROM xonsh_history"; + let row = sqlx::query(query).fetch_one(&self.pool).await?; + let count: u32 = row.get(0); + Ok(count as usize) + } + + async fn load(self, loader: &mut impl Loader) -> Result<()> { + let query = r#" + SELECT inp, rtn, tsb, tse, cwd, + MIN(tsb) OVER (PARTITION BY sessionid) AS session_start + FROM xonsh_history + ORDER BY rowid + "#; + + let mut entries = sqlx::query_as::<_, HistDbEntry>(query).fetch(&self.pool); + + let mut count = 0; + while let Some(entry) = entries.try_next().await? { + let hist = entry.into_hist_with_hostname(self.hostname.clone()); + loader.push(hist).await?; + count += 1; + } + + println!("Loaded: {count}"); + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use time::macros::datetime; + + use super::*; + + use crate::history::History; + use crate::import::tests::TestLoader; + + #[test] + fn test_db_path_xonsh() { + let db_path = xonsh_db_path(Some("/home/user/xonsh_data".to_string())).unwrap(); + assert_eq!( + db_path, + PathBuf::from("/home/user/xonsh_data/xonsh-history.sqlite") + ); + } + + #[tokio::test] + async fn test_import() { + let connection_str = "tests/data/xonsh-history.sqlite"; + let xonsh_sqlite = XonshSqlite { + pool: SqlitePool::connect(connection_str).await.unwrap(), + hostname: "box:user".to_string(), + }; + + let mut loader = TestLoader::default(); + xonsh_sqlite.load(&mut loader).await.unwrap(); + + for (actual, expected) in loader.buf.iter().zip(expected_hist_entries().iter()) { + assert_eq!(actual.timestamp, expected.timestamp); + assert_eq!(actual.command, expected.command); + assert_eq!(actual.cwd, expected.cwd); + assert_eq!(actual.exit, expected.exit); + assert_eq!(actual.duration, expected.duration); + assert_eq!(actual.hostname, expected.hostname); + } + } + + fn expected_hist_entries() -> [History; 4] { + [ + History::import() + .timestamp(datetime!(2024-02-6 17:56:21.130956288 +00:00:00)) + .command("echo hello world!".to_string()) + .cwd("/home/user/Documents/code/atuin".to_string()) + .exit(0) + .duration(2628564) + .hostname("box:user".to_string()) + .build() + .into(), + History::import() + .timestamp(datetime!(2024-02-06 17:56:28.190406144 +00:00:00)) + .command("ls -l".to_string()) + .cwd("/home/user/Documents/code/atuin".to_string()) + .exit(0) + .duration(9371519) + .hostname("box:user".to_string()) + .build() + .into(), + History::import() + .timestamp(datetime!(2024-02-06 17:56:46.989020928 +00:00:00)) + .command("false".to_string()) + .cwd("/home/user/Documents/code/atuin".to_string()) + .exit(1) + .duration(17337560) + .hostname("box:user".to_string()) + .build() + .into(), + History::import() + .timestamp(datetime!(2024-02-06 17:56:48.218384128 +00:00:00)) + .command("exit".to_string()) + .cwd("/home/user/Documents/code/atuin".to_string()) + .exit(0) + .duration(4599094) + .hostname("box:user".to_string()) + .build() + .into(), + ] + } +} diff --git a/crates/atuin-client/src/import/zsh.rs b/crates/atuin-client/src/import/zsh.rs new file mode 100644 index 00000000..11e2f371 --- /dev/null +++ b/crates/atuin-client/src/import/zsh.rs @@ -0,0 +1,230 @@ +// import old shell history! +// automatically hoover up all that we can find + +use std::borrow::Cow; +use std::path::PathBuf; + +use async_trait::async_trait; +use directories::UserDirs; +use eyre::{Result, eyre}; +use time::OffsetDateTime; + +use super::{Importer, Loader, get_histfile_path, unix_byte_lines}; +use crate::history::History; +use crate::import::read_to_end; + +#[derive(Debug)] +pub struct Zsh { + bytes: Vec<u8>, +} + +fn default_histpath() -> Result<PathBuf> { + // oh-my-zsh sets HISTFILE=~/.zhistory + // zsh has no default value for this var, but uses ~/.zhistory. + // zsh-newuser-install propose as default .histfile https://github.com/zsh-users/zsh/blob/master/Functions/Newuser/zsh-newuser-install#L794 + // we could maybe be smarter about this in the future :) + let user_dirs = UserDirs::new().ok_or_else(|| eyre!("could not find user directories"))?; + let home_dir = user_dirs.home_dir(); + + let mut candidates = [".zhistory", ".zsh_history", ".histfile"].iter(); + loop { + match candidates.next() { + Some(candidate) => { + let histpath = home_dir.join(candidate); + if histpath.exists() { + break Ok(histpath); + } + } + None => { + break Err(eyre!( + "Could not find history file. Try setting and exporting $HISTFILE" + )); + } + } + } +} + +#[async_trait] +impl Importer for Zsh { + const NAME: &'static str = "zsh"; + + async fn new() -> Result<Self> { + let bytes = read_to_end(get_histfile_path(default_histpath)?)?; + Ok(Self { bytes }) + } + + async fn entries(&mut self) -> Result<usize> { + Ok(super::count_lines(&self.bytes)) + } + + async fn load(self, h: &mut impl Loader) -> Result<()> { + let now = OffsetDateTime::now_utc(); + let mut line = String::new(); + + let mut counter = 0; + for b in unix_byte_lines(&self.bytes) { + let s = match unmetafy(b) { + Some(s) => s, + _ => continue, // we can skip past things like invalid utf8 + }; + + if let Some(s) = s.strip_suffix('\\') { + line.push_str(s); + line.push('\n'); + } else { + line.push_str(&s); + let command = std::mem::take(&mut line); + + if let Some(command) = command.strip_prefix(": ") { + counter += 1; + h.push(parse_extended(command, counter)).await?; + } else { + let offset = time::Duration::seconds(counter); + counter += 1; + + let imported = History::import() + // preserve ordering + .timestamp(now - offset) + .command(command.trim_end().to_string()); + + h.push(imported.build().into()).await?; + } + } + } + + Ok(()) + } +} + +fn parse_extended(line: &str, counter: i64) -> History { + let (time, duration) = line.split_once(':').unwrap(); + let (duration, command) = duration.split_once(';').unwrap(); + + let time = time + .parse::<i64>() + .ok() + .and_then(|t| OffsetDateTime::from_unix_timestamp(t).ok()) + .unwrap_or_else(OffsetDateTime::now_utc) + + time::Duration::milliseconds(counter); + + // use nanos, because why the hell not? we won't display them. + let duration = duration.parse::<i64>().map_or(-1, |t| t * 1_000_000_000); + + let imported = History::import() + .timestamp(time) + .command(command.trim_end().to_string()) + .duration(duration); + + imported.build().into() +} + +fn unmetafy(line: &[u8]) -> Option<Cow<'_, str>> { + if line.contains(&0x83) { + let mut s = Vec::with_capacity(line.len()); + let mut is_meta = false; + for ch in line { + if *ch == 0x83 { + is_meta = true; + } else if is_meta { + is_meta = false; + s.push(*ch ^ 32); + } else { + s.push(*ch) + } + } + String::from_utf8(s).ok().map(Cow::Owned) + } else { + std::str::from_utf8(line).ok().map(Cow::Borrowed) + } +} + +#[cfg(test)] +mod test { + use itertools::assert_equal; + + use crate::import::tests::TestLoader; + + use super::*; + + #[test] + fn test_parse_extended_simple() { + let parsed = parse_extended("1613322469:0;cargo install atuin", 0); + + assert_eq!(parsed.command, "cargo install atuin"); + assert_eq!(parsed.duration, 0); + assert_eq!( + parsed.timestamp, + OffsetDateTime::from_unix_timestamp(1_613_322_469).unwrap() + ); + + let parsed = parse_extended("1613322469:10;cargo install atuin;cargo update", 0); + + assert_eq!(parsed.command, "cargo install atuin;cargo update"); + assert_eq!(parsed.duration, 10_000_000_000); + assert_eq!( + parsed.timestamp, + OffsetDateTime::from_unix_timestamp(1_613_322_469).unwrap() + ); + + let parsed = parse_extended("1613322469:10;cargo :b̷i̶t̴r̵o̴t̴ ̵i̷s̴ ̷r̶e̵a̸l̷", 0); + + assert_eq!(parsed.command, "cargo :b̷i̶t̴r̵o̴t̴ ̵i̷s̴ ̷r̶e̵a̸l̷"); + assert_eq!(parsed.duration, 10_000_000_000); + assert_eq!( + parsed.timestamp, + OffsetDateTime::from_unix_timestamp(1_613_322_469).unwrap() + ); + + let parsed = parse_extended("1613322469:10;cargo install \\n atuin\n", 0); + + assert_eq!(parsed.command, "cargo install \\n atuin"); + assert_eq!(parsed.duration, 10_000_000_000); + assert_eq!( + parsed.timestamp, + OffsetDateTime::from_unix_timestamp(1_613_322_469).unwrap() + ); + } + + #[tokio::test] + async fn test_parse_file() { + let bytes = r": 1613322469:0;cargo install atuin +: 1613322469:10;cargo install atuin; \\ +cargo update +: 1613322469:10;cargo :b̷i̶t̴r̵o̴t̴ ̵i̷s̴ ̷r̶e̵a̸l̷ +" + .as_bytes() + .to_owned(); + + let mut zsh = Zsh { bytes }; + assert_eq!(zsh.entries().await.unwrap(), 4); + + let mut loader = TestLoader::default(); + zsh.load(&mut loader).await.unwrap(); + + assert_equal( + loader.buf.iter().map(|h| h.command.as_str()), + [ + "cargo install atuin", + "cargo install atuin; \\\ncargo update", + "cargo :b̷i̶t̴r̵o̴t̴ ̵i̷s̴ ̷r̶e̵a̸l̷", + ], + ); + } + + #[tokio::test] + async fn test_parse_metafied() { + let bytes = + b"echo \xe4\xbd\x83\x80\xe5\xa5\xbd\nls ~/\xe9\x83\xbf\xb3\xe4\xb9\x83\xb0\n".to_vec(); + + let mut zsh = Zsh { bytes }; + assert_eq!(zsh.entries().await.unwrap(), 2); + + let mut loader = TestLoader::default(); + zsh.load(&mut loader).await.unwrap(); + + assert_equal( + loader.buf.iter().map(|h| h.command.as_str()), + ["echo 你好", "ls ~/音乐"], + ); + } +} diff --git a/crates/atuin-client/src/import/zsh_histdb.rs b/crates/atuin-client/src/import/zsh_histdb.rs new file mode 100644 index 00000000..cac85566 --- /dev/null +++ b/crates/atuin-client/src/import/zsh_histdb.rs @@ -0,0 +1,249 @@ +// import old shell history from zsh-histdb! +// automatically hoover up all that we can find + +// As far as i can tell there are no version numbers in the histdb sqlite DB, so we're going based +// on the schema from 2022-05-01 +// +// I have run into some histories that will not import b/c of non UTF-8 characters. +// + +// +// An Example sqlite query for hsitdb data: +// +//id|session|command_id|place_id|exit_status|start_time|duration|id|argv|id|host|dir +// +// +// select +// history.id, +// history.start_time, +// places.host, +// places.dir, +// commands.argv +// from history +// left join commands on history.command_id = commands.id +// left join places on history.place_id = places.id ; +// +// CREATE TABLE history (id integer primary key autoincrement, +// session int, +// command_id int references commands (id), +// place_id int references places (id), +// exit_status int, +// start_time int, +// duration int); +// + +use std::collections::HashMap; +use std::path::{Path, PathBuf}; + +use async_trait::async_trait; +use atuin_common::utils::uuid_v7; +use directories::UserDirs; +use eyre::{Result, eyre}; +use sqlx::{Pool, sqlite::SqlitePool}; +use time::PrimitiveDateTime; + +use super::Importer; +use crate::history::History; +use crate::import::Loader; +use crate::utils::{get_hostname, get_username}; + +#[derive(sqlx::FromRow, Debug)] +pub struct HistDbEntryCount { + pub count: usize, +} + +#[derive(sqlx::FromRow, Debug)] +pub struct HistDbEntry { + pub id: i64, + pub start_time: PrimitiveDateTime, + pub host: Vec<u8>, + pub dir: Vec<u8>, + pub argv: Vec<u8>, + pub duration: i64, + pub exit_status: i64, + pub session: i64, +} + +#[derive(Debug)] +pub struct ZshHistDb { + histdb: Vec<HistDbEntry>, + username: String, +} + +/// Read db at given file, return vector of entries. +async fn hist_from_db(dbpath: PathBuf) -> Result<Vec<HistDbEntry>> { + let pool = SqlitePool::connect(dbpath.to_str().unwrap()).await?; + hist_from_db_conn(pool).await +} + +async fn hist_from_db_conn(pool: Pool<sqlx::Sqlite>) -> Result<Vec<HistDbEntry>> { + let query = r#" + SELECT + history.id, history.start_time, history.duration, places.host, places.dir, + commands.argv, history.exit_status, history.session + FROM history + LEFT JOIN commands ON history.command_id = commands.id + LEFT JOIN places ON history.place_id = places.id + ORDER BY history.start_time + "#; + let histdb_vec: Vec<HistDbEntry> = sqlx::query_as::<_, HistDbEntry>(query) + .fetch_all(&pool) + .await?; + Ok(histdb_vec) +} + +impl ZshHistDb { + pub fn histpath_candidate() -> PathBuf { + // By default histdb database is `${HOME}/.histdb/zsh-history.db` + // This can be modified by ${HISTDB_FILE} + // + // if [[ -z ${HISTDB_FILE} ]]; then + // typeset -g HISTDB_FILE="${HOME}/.histdb/zsh-history.db" + let user_dirs = UserDirs::new().unwrap(); // should catch error here? + let home_dir = user_dirs.home_dir(); + std::env::var("HISTDB_FILE") + .as_ref() + .map(|x| Path::new(x).to_path_buf()) + .unwrap_or_else(|_err| home_dir.join(".histdb/zsh-history.db")) + } + pub fn histpath() -> Result<PathBuf> { + let histdb_path = ZshHistDb::histpath_candidate(); + if histdb_path.exists() { + Ok(histdb_path) + } else { + Err(eyre!( + "Could not find history file. Try setting $HISTDB_FILE" + )) + } + } +} + +#[async_trait] +impl Importer for ZshHistDb { + // Not sure how this is used + const NAME: &'static str = "zsh_histdb"; + + /// Creates a new ZshHistDb and populates the history based on the pre-populated data + /// structure. + async fn new() -> Result<Self> { + let dbpath = ZshHistDb::histpath()?; + let histdb_entry_vec = hist_from_db(dbpath).await?; + Ok(Self { + histdb: histdb_entry_vec, + username: get_username(), + }) + } + + async fn entries(&mut self) -> Result<usize> { + Ok(self.histdb.len()) + } + + async fn load(self, h: &mut impl Loader) -> Result<()> { + let mut session_map = HashMap::new(); + for entry in self.histdb { + let command = match std::str::from_utf8(&entry.argv) { + Ok(s) => s.trim_end(), + Err(_) => continue, // we can skip past things like invalid utf8 + }; + let cwd = match std::str::from_utf8(&entry.dir) { + Ok(s) => s.trim_end(), + Err(_) => continue, // we can skip past things like invalid utf8 + }; + let hostname = format!( + "{}:{}", + String::from_utf8(entry.host).unwrap_or_else(|_e| get_hostname()), + self.username + ); + let session = session_map.entry(entry.session).or_insert_with(uuid_v7); + + let imported = History::import() + .timestamp(entry.start_time.assume_utc()) + .command(command) + .cwd(cwd) + .duration(entry.duration * 1_000_000_000) + .exit(entry.exit_status) + .session(session.as_simple().to_string()) + .hostname(hostname) + .build(); + h.push(imported.into()).await?; + } + Ok(()) + } +} + +#[cfg(test)] +mod test { + + use super::*; + use sqlx::sqlite::SqlitePoolOptions; + use std::env; + #[tokio::test(flavor = "multi_thread")] + #[allow(unsafe_code)] + async fn test_env_vars() { + let test_env_db = "nonstd-zsh-history.db"; + let key = "HISTDB_FILE"; + // TODO: Audit that the environment access only happens in single-threaded code. + unsafe { env::set_var(key, test_env_db) }; + + // test the env got set + assert_eq!(env::var(key).unwrap(), test_env_db.to_string()); + + // test histdb returns the proper db from previous step + let histdb_path = ZshHistDb::histpath_candidate(); + assert_eq!(histdb_path.to_str().unwrap(), test_env_db); + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_import() { + let pool: SqlitePool = SqlitePoolOptions::new() + .min_connections(2) + .connect(":memory:") + .await + .unwrap(); + + // sql dump directly from a test database. + let db_sql = r#" + PRAGMA foreign_keys=OFF; + BEGIN TRANSACTION; + CREATE TABLE commands (id integer primary key autoincrement, argv text, unique(argv) on conflict ignore); + INSERT INTO commands VALUES(1,'pwd'); + INSERT INTO commands VALUES(2,'curl google.com'); + INSERT INTO commands VALUES(3,'bash'); + CREATE TABLE places (id integer primary key autoincrement, host text, dir text, unique(host, dir) on conflict ignore); + INSERT INTO places VALUES(1,'mbp16.local','/home/noyez'); + CREATE TABLE history (id integer primary key autoincrement, + session int, + command_id int references commands (id), + place_id int references places (id), + exit_status int, + start_time int, + duration int); + INSERT INTO history VALUES(1,0,1,1,0,1651497918,1); + INSERT INTO history VALUES(2,0,2,1,0,1651497923,1); + INSERT INTO history VALUES(3,0,3,1,NULL,1651497930,NULL); + DELETE FROM sqlite_sequence; + INSERT INTO sqlite_sequence VALUES('commands',3); + INSERT INTO sqlite_sequence VALUES('places',3); + INSERT INTO sqlite_sequence VALUES('history',3); + CREATE INDEX hist_time on history(start_time); + CREATE INDEX place_dir on places(dir); + CREATE INDEX place_host on places(host); + CREATE INDEX history_command_place on history(command_id, place_id); + COMMIT; "#; + + sqlx::query(db_sql).execute(&pool).await.unwrap(); + + // test histdb iterator + let histdb_vec = hist_from_db_conn(pool).await.unwrap(); + let histdb = ZshHistDb { + histdb: histdb_vec, + username: get_username(), + }; + + println!("h: {:#?}", histdb.histdb); + println!("counter: {:?}", histdb.histdb.len()); + for i in histdb.histdb { + println!("{i:?}"); + } + } +} diff --git a/crates/atuin-client/src/lib.rs b/crates/atuin-client/src/lib.rs new file mode 100644 index 00000000..352e5746 --- /dev/null +++ b/crates/atuin-client/src/lib.rs @@ -0,0 +1,30 @@ +#![deny(unsafe_code)] + +#[macro_use] +extern crate log; + +#[cfg(feature = "sync")] +pub mod api_client; +#[cfg(feature = "hub")] +pub mod hub; +#[cfg(feature = "sync")] +pub mod login; +#[cfg(feature = "sync")] +pub mod register; +#[cfg(feature = "sync")] +pub mod sync; + +pub mod database; +pub mod encryption; +pub mod history; +pub mod import; +pub mod logout; +pub mod meta; +pub mod ordering; +pub mod plugin; +pub mod record; +pub mod secrets; +pub mod settings; +pub mod theme; + +mod utils; diff --git a/crates/atuin-client/src/login.rs b/crates/atuin-client/src/login.rs new file mode 100644 index 00000000..d72d1c09 --- /dev/null +++ b/crates/atuin-client/src/login.rs @@ -0,0 +1,87 @@ +use std::path::PathBuf; + +use atuin_common::api::LoginRequest; +use eyre::{Context, Result, bail}; +use tokio::fs::File; +use tokio::io::AsyncWriteExt; + +use crate::{ + api_client, + encryption::{Key, decode_key, encode_key, load_key}, + record::{sqlite_store::SqliteStore, store::Store}, + settings::Settings, +}; + +pub async fn login( + settings: &Settings, + store: &SqliteStore, + username: String, + password: String, + key: String, +) -> Result<String> { + // try parse the key as a mnemonic... + let key = match bip39::Mnemonic::from_phrase(&key, bip39::Language::English) { + Ok(mnemonic) => encode_key(Key::from_slice(mnemonic.entropy()))?, + Err(err) => { + match err { + // assume they copied in the base64 key + bip39::ErrorKind::InvalidWord(_) => key, + bip39::ErrorKind::InvalidChecksum => { + bail!("key mnemonic was not valid") + } + bip39::ErrorKind::InvalidKeysize(_) + | bip39::ErrorKind::InvalidWordLength(_) + | bip39::ErrorKind::InvalidEntropyLength(_, _) => { + bail!("key was not the correct length") + } + } + } + }; + + let key_path = settings.key_path.as_str(); + let key_path = PathBuf::from(key_path); + + if !key_path.exists() { + if decode_key(key.clone()).is_err() { + bail!("the specified key was invalid"); + } + + let mut file = File::create(&key_path).await?; + file.write_all(key.as_bytes()).await?; + } else { + // we now know that the user has logged in specifying a key, AND that the key path + // exists + + // 1. check if the saved key and the provided key match. if so, nothing to do. + // 2. if not, re-encrypt the local history and overwrite the key + let current_key: [u8; 32] = load_key(settings)?.into(); + + let encoded = key.clone(); // gonna want to save it in a bit + let new_key: [u8; 32] = decode_key(key) + .context("could not decode provided key - is not valid base64")? + .into(); + + if new_key != current_key { + println!("\nRe-encrypting local store with new key"); + + store.re_encrypt(¤t_key, &new_key).await?; + + println!("Writing new key"); + let mut file = File::create(&key_path).await?; + file.write_all(encoded.as_bytes()).await?; + } + } + + let session = api_client::login( + settings.sync_address.as_str(), + LoginRequest { username, password }, + ) + .await?; + + Settings::meta_store() + .await? + .save_session(&session.session) + .await?; + + Ok(session.session) +} diff --git a/crates/atuin-client/src/logout.rs b/crates/atuin-client/src/logout.rs new file mode 100644 index 00000000..80f0ad73 --- /dev/null +++ b/crates/atuin-client/src/logout.rs @@ -0,0 +1,17 @@ +use eyre::Result; + +use crate::settings::Settings; + +pub async fn logout() -> Result<()> { + let meta = Settings::meta_store().await?; + + if meta.logged_in().await? { + meta.delete_session().await?; + meta.delete_hub_session().await?; + println!("You have logged out!"); + } else { + println!("You are not logged in"); + } + + Ok(()) +} diff --git a/crates/atuin-client/src/meta.rs b/crates/atuin-client/src/meta.rs new file mode 100644 index 00000000..eb6dd8cf --- /dev/null +++ b/crates/atuin-client/src/meta.rs @@ -0,0 +1,384 @@ +use std::path::Path; +use std::str::FromStr; +use std::time::Duration; + +use atuin_common::record::HostId; +use eyre::{Result, eyre}; +use sqlx::sqlite::{SqliteConnectOptions, SqliteJournalMode, SqlitePool, SqlitePoolOptions}; +use time::{OffsetDateTime, format_description::well_known::Rfc3339}; +use tokio::sync::OnceCell; +use uuid::Uuid; + +// Filenames for the legacy plain-text files that we migrate from. +const LEGACY_HOST_ID_FILENAME: &str = "host_id"; +const LEGACY_LAST_SYNC_FILENAME: &str = "last_sync_time"; +const LEGACY_LAST_VERSION_CHECK_FILENAME: &str = "last_version_check_time"; +const LEGACY_LATEST_VERSION_FILENAME: &str = "latest_version"; +const LEGACY_SESSION_FILENAME: &str = "session"; + +const KEY_HOST_ID: &str = "host_id"; +const KEY_LAST_SYNC: &str = "last_sync_time"; +const KEY_LAST_VERSION_CHECK: &str = "last_version_check_time"; +const KEY_LATEST_VERSION: &str = "latest_version"; +const KEY_SESSION: &str = "session"; +const KEY_HUB_SESSION: &str = "hub_session"; +const KEY_FILES_MIGRATED: &str = "files_migrated"; + +pub struct MetaStore { + pool: SqlitePool, + cached_host_id: OnceCell<HostId>, +} + +impl MetaStore { + pub async fn new(path: impl AsRef<Path>, timeout: f64) -> Result<Self> { + let path = path.as_ref(); + let path_str = path + .as_os_str() + .to_str() + .ok_or_else(|| eyre!("meta database path is not valid UTF-8: {path:?}"))?; + debug!("opening meta sqlite database at {path:?}"); + + let is_memory = path_str.contains(":memory:"); + + if !is_memory + && !path.exists() + && let Some(dir) = path.parent() + { + fs_err::create_dir_all(dir)?; + } + + // Use DELETE journal mode instead of WAL. This is a small, infrequently- + // written KV store — WAL's concurrency benefits aren't needed, and DELETE + // mode avoids creating auxiliary -wal/-shm files that complicate + // permission handling. + let opts = SqliteConnectOptions::from_str(path_str)? + .journal_mode(SqliteJournalMode::Delete) + .optimize_on_close(true, None) + .create_if_missing(true); + + let pool = SqlitePoolOptions::new() + .acquire_timeout(Duration::from_secs_f64(timeout)) + .connect_with(opts) + .await?; + + sqlx::migrate!("./meta-migrations").run(&pool).await?; + + // Session tokens are stored in this database, so restrict permissions. + #[cfg(unix)] + if !is_memory { + use std::os::unix::fs::PermissionsExt; + std::fs::set_permissions(path, std::fs::Permissions::from_mode(0o600))?; + } + + let store = Self { + pool, + cached_host_id: OnceCell::const_new(), + }; + + if !is_memory { + store.migrate_files().await?; + } + + Ok(store) + } + + // Generic key-value operations + + pub async fn get(&self, key: &str) -> Result<Option<String>> { + let row: Option<(String,)> = sqlx::query_as("SELECT value FROM meta WHERE key = ?1") + .bind(key) + .fetch_optional(&self.pool) + .await?; + + Ok(row.map(|r| r.0)) + } + + pub async fn set(&self, key: &str, value: &str) -> Result<()> { + sqlx::query( + "INSERT INTO meta (key, value, updated_at) VALUES (?1, ?2, strftime('%s', 'now')) + ON CONFLICT(key) DO UPDATE SET value = ?2, updated_at = strftime('%s', 'now')", + ) + .bind(key) + .bind(value) + .execute(&self.pool) + .await?; + + Ok(()) + } + + pub async fn delete(&self, key: &str) -> Result<()> { + sqlx::query("DELETE FROM meta WHERE key = ?1") + .bind(key) + .execute(&self.pool) + .await?; + + Ok(()) + } + + // Typed accessors + + pub async fn host_id(&self) -> Result<HostId> { + self.cached_host_id + .get_or_try_init(|| async { + if let Some(id) = self.get(KEY_HOST_ID).await? { + let parsed = Uuid::from_str(id.as_str()) + .map_err(|e| eyre!("failed to parse host ID: {e}"))?; + return Ok(HostId(parsed)); + } + + let uuid = atuin_common::utils::uuid_v7(); + self.set(KEY_HOST_ID, uuid.as_simple().to_string().as_ref()) + .await?; + + Ok(HostId(uuid)) + }) + .await + .copied() + } + + pub async fn last_sync(&self) -> Result<OffsetDateTime> { + match self.get(KEY_LAST_SYNC).await? { + Some(v) => Ok(OffsetDateTime::parse(v.as_str(), &Rfc3339)?), + None => Ok(OffsetDateTime::UNIX_EPOCH), + } + } + + pub async fn save_sync_time(&self) -> Result<()> { + self.set( + KEY_LAST_SYNC, + OffsetDateTime::now_utc().format(&Rfc3339)?.as_str(), + ) + .await + } + + pub async fn last_version_check(&self) -> Result<OffsetDateTime> { + match self.get(KEY_LAST_VERSION_CHECK).await? { + Some(v) => Ok(OffsetDateTime::parse(v.as_str(), &Rfc3339)?), + None => Ok(OffsetDateTime::UNIX_EPOCH), + } + } + + pub async fn save_version_check_time(&self) -> Result<()> { + self.set( + KEY_LAST_VERSION_CHECK, + OffsetDateTime::now_utc().format(&Rfc3339)?.as_str(), + ) + .await + } + + pub async fn latest_version(&self) -> Result<Option<String>> { + self.get(KEY_LATEST_VERSION).await + } + + pub async fn save_latest_version(&self, version: &str) -> Result<()> { + self.set(KEY_LATEST_VERSION, version).await + } + + pub async fn session_token(&self) -> Result<Option<String>> { + self.get(KEY_SESSION).await + } + + pub async fn save_session(&self, token: &str) -> Result<()> { + self.set(KEY_SESSION, token).await + } + + pub async fn delete_session(&self) -> Result<()> { + self.delete(KEY_SESSION).await + } + + pub async fn logged_in(&self) -> Result<bool> { + Ok(self.session_token().await?.is_some() || self.hub_session_token().await?.is_some()) + } + + // Hub session methods (separate from sync session, used for Hub-specific features like AI) + + pub async fn hub_session_token(&self) -> Result<Option<String>> { + self.get(KEY_HUB_SESSION).await + } + + pub async fn save_hub_session(&self, token: &str) -> Result<()> { + self.set(KEY_HUB_SESSION, token).await + } + + pub async fn delete_hub_session(&self) -> Result<()> { + self.delete(KEY_HUB_SESSION).await + } + + pub async fn hub_logged_in(&self) -> Result<bool> { + Ok(self.hub_session_token().await?.is_some()) + } + + // File migration: on first open, migrate old plain-text files into the database. + // Old files are left in place for safe downgrades. + + async fn migrate_files(&self) -> Result<()> { + if self.get(KEY_FILES_MIGRATED).await?.is_some() { + return Ok(()); + } + + let data_dir = crate::settings::Settings::effective_data_dir(); + + // host_id — validate as UUID + let host_id_path = data_dir.join(LEGACY_HOST_ID_FILENAME); + if host_id_path.exists() + && let Ok(value) = fs_err::read_to_string(&host_id_path) + { + let value = value.trim(); + if !value.is_empty() { + if Uuid::from_str(value).is_ok() { + self.set(KEY_HOST_ID, value).await?; + } else { + warn!("skipping migration of host_id: invalid UUID {value:?}"); + } + } + } + + // last_sync_time — validate as RFC3339 + let sync_path = data_dir.join(LEGACY_LAST_SYNC_FILENAME); + if sync_path.exists() + && let Ok(value) = fs_err::read_to_string(&sync_path) + { + let value = value.trim(); + if !value.is_empty() { + if OffsetDateTime::parse(value, &Rfc3339).is_ok() { + self.set(KEY_LAST_SYNC, value).await?; + } else { + warn!("skipping migration of last_sync_time: invalid RFC3339 {value:?}"); + } + } + } + + // last_version_check_time — validate as RFC3339 + let version_check_path = data_dir.join(LEGACY_LAST_VERSION_CHECK_FILENAME); + if version_check_path.exists() + && let Ok(value) = fs_err::read_to_string(&version_check_path) + { + let value = value.trim(); + if !value.is_empty() { + if OffsetDateTime::parse(value, &Rfc3339).is_ok() { + self.set(KEY_LAST_VERSION_CHECK, value).await?; + } else { + warn!( + "skipping migration of last_version_check_time: invalid RFC3339 {value:?}" + ); + } + } + } + + // latest_version — no strict validation, just non-empty + let latest_version_path = data_dir.join(LEGACY_LATEST_VERSION_FILENAME); + if latest_version_path.exists() + && let Ok(value) = fs_err::read_to_string(&latest_version_path) + { + let value = value.trim(); + if !value.is_empty() { + self.set(KEY_LATEST_VERSION, value).await?; + } + } + + // session token — no strict validation, just non-empty + let session_path = data_dir.join(LEGACY_SESSION_FILENAME); + if session_path.exists() + && let Ok(value) = fs_err::read_to_string(&session_path) + { + let value = value.trim(); + if !value.is_empty() { + self.set(KEY_SESSION, value).await?; + } + } + + self.set(KEY_FILES_MIGRATED, "true").await?; + + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + async fn new_test_store() -> MetaStore { + MetaStore::new("sqlite::memory:", 2.0).await.unwrap() + } + + #[tokio::test] + async fn test_get_set_delete() { + let store = new_test_store().await; + + assert_eq!(store.get("foo").await.unwrap(), None); + + store.set("foo", "bar").await.unwrap(); + assert_eq!(store.get("foo").await.unwrap(), Some("bar".to_string())); + + store.set("foo", "baz").await.unwrap(); + assert_eq!(store.get("foo").await.unwrap(), Some("baz".to_string())); + + store.delete("foo").await.unwrap(); + assert_eq!(store.get("foo").await.unwrap(), None); + } + + #[tokio::test] + async fn test_host_id_generation_and_stability() { + let store = new_test_store().await; + + let id1 = store.host_id().await.unwrap(); + let id2 = store.host_id().await.unwrap(); + + assert_eq!(id1, id2, "host_id should be stable across calls"); + } + + #[tokio::test] + async fn test_sync_time() { + let store = new_test_store().await; + + let t = store.last_sync().await.unwrap(); + assert_eq!(t, OffsetDateTime::UNIX_EPOCH); + + store.save_sync_time().await.unwrap(); + let t = store.last_sync().await.unwrap(); + assert!(t > OffsetDateTime::UNIX_EPOCH); + } + + #[tokio::test] + async fn test_version_check_time() { + let store = new_test_store().await; + + let t = store.last_version_check().await.unwrap(); + assert_eq!(t, OffsetDateTime::UNIX_EPOCH); + + store.save_version_check_time().await.unwrap(); + let t = store.last_version_check().await.unwrap(); + assert!(t > OffsetDateTime::UNIX_EPOCH); + } + + #[tokio::test] + async fn test_session_crud() { + let store = new_test_store().await; + + assert!(!store.logged_in().await.unwrap()); + assert_eq!(store.session_token().await.unwrap(), None); + + store.save_session("tok123").await.unwrap(); + assert!(store.logged_in().await.unwrap()); + assert_eq!( + store.session_token().await.unwrap(), + Some("tok123".to_string()) + ); + + store.delete_session().await.unwrap(); + assert!(!store.logged_in().await.unwrap()); + } + + #[tokio::test] + async fn test_latest_version() { + let store = new_test_store().await; + + assert_eq!(store.latest_version().await.unwrap(), None); + + store.save_latest_version("1.2.3").await.unwrap(); + assert_eq!( + store.latest_version().await.unwrap(), + Some("1.2.3".to_string()) + ); + } +} diff --git a/crates/atuin-client/src/ordering.rs b/crates/atuin-client/src/ordering.rs new file mode 100644 index 00000000..4e5ec84c --- /dev/null +++ b/crates/atuin-client/src/ordering.rs @@ -0,0 +1,32 @@ +use minspan::minspan; + +use super::{history::History, settings::SearchMode}; + +pub fn reorder_fuzzy(mode: SearchMode, query: &str, res: Vec<History>) -> Vec<History> { + match mode { + SearchMode::Fuzzy => reorder(query, |x| &x.command, res), + _ => res, + } +} + +fn reorder<F, A>(query: &str, f: F, res: Vec<A>) -> Vec<A> +where + F: Fn(&A) -> &String, + A: Clone, +{ + let mut r = res.clone(); + let qvec = &query.chars().collect(); + r.sort_by_cached_key(|h| { + // TODO for fzf search we should sum up scores for each matched term + let (from, to) = match minspan::span(qvec, &(f(h).chars().collect())) { + Some(x) => x, + // this is a little unfortunate: when we are asked to match a query that is found nowhere, + // we don't want to return a None, as the comparison behaviour would put the worst matches + // at the front. therefore, we'll return a set of indices that are one larger than the longest + // possible legitimate match. This is meaningless except as a comparison. + None => (0, res.len()), + }; + 1 + to - from + }); + r +} diff --git a/crates/atuin-client/src/plugin.rs b/crates/atuin-client/src/plugin.rs new file mode 100644 index 00000000..21a2bcef --- /dev/null +++ b/crates/atuin-client/src/plugin.rs @@ -0,0 +1,97 @@ +use std::collections::HashMap; + +#[derive(Debug, Clone)] +pub struct OfficialPlugin { + pub name: String, + pub description: String, + pub install_message: String, +} + +impl OfficialPlugin { + pub fn new(name: &str, description: &str, install_message: &str) -> Self { + Self { + name: name.to_string(), + description: description.to_string(), + install_message: install_message.to_string(), + } + } +} + +pub struct OfficialPluginRegistry { + plugins: HashMap<String, OfficialPlugin>, +} + +impl OfficialPluginRegistry { + pub fn new() -> Self { + let mut registry = Self { + plugins: HashMap::new(), + }; + + // Register official plugins + registry.register_official_plugins(); + + registry + } + + fn register_official_plugins(&mut self) { + // atuin-update plugin + self.plugins.insert( + "update".to_string(), + OfficialPlugin::new( + "update", + "Update atuin to the latest version", + "The 'atuin update' command is provided by the atuin-update plugin.\n\ + It is only installed if you used the install script\n \ + If you used a package manager (brew, apt, etc), please continue to use it for updates" + ), + ); + } + + pub fn get_plugin(&self, name: &str) -> Option<&OfficialPlugin> { + self.plugins.get(name) + } + + pub fn is_official_plugin(&self, name: &str) -> bool { + self.plugins.contains_key(name) + } + + pub fn get_install_message(&self, name: &str) -> Option<&str> { + self.plugins + .get(name) + .map(|plugin| plugin.install_message.as_str()) + } +} + +impl Default for OfficialPluginRegistry { + fn default() -> Self { + Self::new() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_registry_creation() { + let registry = OfficialPluginRegistry::new(); + assert!(registry.is_official_plugin("update")); + assert!(!registry.is_official_plugin("nonexistent")); + } + + #[test] + fn test_get_plugin() { + let registry = OfficialPluginRegistry::new(); + let plugin = registry.get_plugin("update"); + assert!(plugin.is_some()); + assert_eq!(plugin.unwrap().name, "update"); + } + + #[test] + fn test_get_install_message() { + let registry = OfficialPluginRegistry::new(); + let message = registry.get_install_message("update"); + assert!(message.is_some()); + assert!(message.unwrap().contains("atuin-update")); + } +} diff --git a/crates/atuin-client/src/record/encryption.rs b/crates/atuin-client/src/record/encryption.rs new file mode 100644 index 00000000..176d75f5 --- /dev/null +++ b/crates/atuin-client/src/record/encryption.rs @@ -0,0 +1,373 @@ +use atuin_common::record::{ + AdditionalData, DecryptedData, EncryptedData, Encryption, HostId, RecordId, RecordIdx, +}; +use base64::{Engine, engine::general_purpose}; +use eyre::{Context, Result, ensure}; +use rusty_paserk::{Key, KeyId, Local, PieWrappedKey}; +use rusty_paseto::core::{ + ImplicitAssertion, Key as DataKey, Local as LocalPurpose, Paseto, PasetoNonce, Payload, V4, +}; +use serde::{Deserialize, Serialize}; + +/// Use PASETO V4 Local encryption using the additional data as an implicit assertion. +#[allow(non_camel_case_types)] +pub struct PASETO_V4; + +/* +Why do we use a random content-encryption key? +Originally I was planning on using a derived key for encryption based on additional data. +This would be a lot more secure than using the master key directly. + +However, there's an established norm of using a random key. This scheme might be otherwise known as +- client-side encryption +- envelope encryption +- key wrapping + +A HSM (Hardware Security Module) provider, eg: AWS, Azure, GCP, or even a physical device like a YubiKey +will have some keys that they keep to themselves. These keys never leave their physical hardware. +If they never leave the hardware, then encrypting large amounts of data means giving them the data and waiting. +This is not a practical solution. Instead, generate a unique key for your data, encrypt that using your HSM +and then store that with your data. + +See + - <https://docs.aws.amazon.com/wellarchitected/latest/financial-services-industry-lens/use-envelope-encryption-with-customer-master-keys.html> + - <https://cloud.google.com/kms/docs/envelope-encryption> + - <https://learn.microsoft.com/en-us/azure/storage/blobs/client-side-encryption?tabs=dotnet#encryption-and-decryption-via-the-envelope-technique> + - <https://www.yubico.com/gb/product/yubihsm-2-fips/> + - <https://cheatsheetseries.owasp.org/cheatsheets/Cryptographic_Storage_Cheat_Sheet.html#encrypting-stored-keys> + +Why would we care? In the past we have received some requests for company solutions. If in future we can configure a +KMS service with little effort, then that would solve a lot of issues for their security team. + +Even for personal use, if a user is not comfortable with sharing keys between hosts, +GCP HSM costs $1/month and $0.03 per 10,000 key operations. Assuming an active user runs +1000 atuin records a day, that would only cost them $1 and 10 cent a month. + +Additionally, key rotations are much simpler using this scheme. Rotating a key is as simple as re-encrypting the CEK, and not the message contents. +This makes it very fast to rotate a key in bulk. + +For future reference, with asymmetric encryption, you can encrypt the CEK without the HSM's involvement, but decrypting +will need the HSM. This allows the encryption path to still be extremely fast (no network calls) but downloads/decryption +that happens in the background can make the network calls to the HSM +*/ + +impl Encryption for PASETO_V4 { + fn re_encrypt( + mut data: EncryptedData, + _ad: AdditionalData, + old_key: &[u8; 32], + new_key: &[u8; 32], + ) -> Result<EncryptedData> { + let cek = Self::decrypt_cek(data.content_encryption_key, old_key)?; + data.content_encryption_key = Self::encrypt_cek(cek, new_key); + Ok(data) + } + + fn encrypt(data: DecryptedData, ad: AdditionalData, key: &[u8; 32]) -> EncryptedData { + // generate a random key for this entry + // aka content-encryption-key (CEK) + let random_key = Key::<V4, Local>::new_os_random(); + + // encode the implicit assertions + let assertions = Assertions::from(ad).encode(); + + // build the payload and encrypt the token + let payload = serde_json::to_string(&AtuinPayload { + data: general_purpose::URL_SAFE_NO_PAD.encode(data.0), + }) + .expect("json encoding can't fail"); + let nonce = DataKey::<32>::try_new_random().expect("could not source from random"); + let nonce = PasetoNonce::<V4, LocalPurpose>::from(&nonce); + + let token = Paseto::<V4, LocalPurpose>::builder() + .set_payload(Payload::from(payload.as_str())) + .set_implicit_assertion(ImplicitAssertion::from(assertions.as_str())) + .try_encrypt(&random_key.into(), &nonce) + .expect("error encrypting atuin data"); + + EncryptedData { + data: token, + content_encryption_key: Self::encrypt_cek(random_key, key), + } + } + + fn decrypt(data: EncryptedData, ad: AdditionalData, key: &[u8; 32]) -> Result<DecryptedData> { + let token = data.data; + let cek = Self::decrypt_cek(data.content_encryption_key, key)?; + + // encode the implicit assertions + let assertions = Assertions::from(ad).encode(); + + // decrypt the payload with the footer and implicit assertions + let payload = Paseto::<V4, LocalPurpose>::try_decrypt( + &token, + &cek.into(), + None, + ImplicitAssertion::from(&*assertions), + ) + .context("could not decrypt entry")?; + + let payload: AtuinPayload = serde_json::from_str(&payload)?; + let data = general_purpose::URL_SAFE_NO_PAD.decode(payload.data)?; + Ok(DecryptedData(data)) + } +} + +impl PASETO_V4 { + fn decrypt_cek(wrapped_cek: String, key: &[u8; 32]) -> Result<Key<V4, Local>> { + let wrapping_key = Key::<V4, Local>::from_bytes(*key); + + // let wrapping_key = PasetoSymmetricKey::from(Key::from(key)); + + let AtuinFooter { kid, wpk } = serde_json::from_str(&wrapped_cek) + .context("wrapped cek did not contain the correct contents")?; + + // check that the wrapping key matches the required key to decrypt. + // In future, we could support multiple keys and use this key to + // look up the key rather than only allow one key. + // For now though we will only support the one key and key rotation will + // have to be a hard reset + let current_kid = wrapping_key.to_id(); + + ensure!( + current_kid == kid, + "attempting to decrypt with incorrect key. currently using {current_kid}, expecting {kid}" + ); + + // decrypt the random key + Ok(wpk.unwrap_key(&wrapping_key)?) + } + + fn encrypt_cek(cek: Key<V4, Local>, key: &[u8; 32]) -> String { + // aka key-encryption-key (KEK) + let wrapping_key = Key::<V4, Local>::from_bytes(*key); + + // wrap the random key so we can decrypt it later + let wrapped_cek = AtuinFooter { + wpk: cek.wrap_pie(&wrapping_key), + kid: wrapping_key.to_id(), + }; + serde_json::to_string(&wrapped_cek).expect("could not serialize wrapped cek") + } +} + +#[derive(Serialize, Deserialize)] +struct AtuinPayload { + data: String, +} + +#[derive(Serialize, Deserialize)] +/// Well-known footer claims for decrypting. This is not encrypted but is stored in the record. +/// <https://github.com/paseto-standard/paseto-spec/blob/master/docs/02-Implementation-Guide/04-Claims.md#optional-footer-claims> +struct AtuinFooter { + /// Wrapped key + wpk: PieWrappedKey<V4, Local>, + /// ID of the key which was used to wrap + kid: KeyId<V4, Local>, +} + +/// Used in the implicit assertions. This is not encrypted and not stored in the data blob. +// This cannot be changed, otherwise it breaks the authenticated encryption. +#[derive(Debug, Copy, Clone, Serialize)] +struct Assertions<'a> { + id: &'a RecordId, + idx: &'a RecordIdx, + version: &'a str, + tag: &'a str, + host: &'a HostId, +} + +impl<'a> From<AdditionalData<'a>> for Assertions<'a> { + fn from(ad: AdditionalData<'a>) -> Self { + Self { + id: ad.id, + version: ad.version, + tag: ad.tag, + host: ad.host, + idx: ad.idx, + } + } +} + +impl Assertions<'_> { + fn encode(&self) -> String { + serde_json::to_string(self).expect("could not serialize implicit assertions") + } +} + +#[cfg(test)] +mod tests { + use atuin_common::{ + record::{Host, Record}, + utils::uuid_v7, + }; + + use super::*; + + #[test] + fn round_trip() { + let key = Key::<V4, Local>::new_os_random(); + + let ad = AdditionalData { + id: &RecordId(uuid_v7()), + version: "v0", + tag: "kv", + host: &HostId(uuid_v7()), + idx: &0, + }; + + let data = DecryptedData(vec![1, 2, 3, 4]); + + let encrypted = PASETO_V4::encrypt(data.clone(), ad, &key.to_bytes()); + let decrypted = PASETO_V4::decrypt(encrypted, ad, &key.to_bytes()).unwrap(); + assert_eq!(decrypted, data); + } + + #[test] + fn same_entry_different_output() { + let key = Key::<V4, Local>::new_os_random(); + + let ad = AdditionalData { + id: &RecordId(uuid_v7()), + version: "v0", + tag: "kv", + host: &HostId(uuid_v7()), + idx: &0, + }; + + let data = DecryptedData(vec![1, 2, 3, 4]); + + let encrypted = PASETO_V4::encrypt(data.clone(), ad, &key.to_bytes()); + let encrypted2 = PASETO_V4::encrypt(data, ad, &key.to_bytes()); + + assert_ne!( + encrypted.data, encrypted2.data, + "re-encrypting the same contents should have different output due to key randomization" + ); + } + + #[test] + fn cannot_decrypt_different_key() { + let key = Key::<V4, Local>::new_os_random(); + let fake_key = Key::<V4, Local>::new_os_random(); + + let ad = AdditionalData { + id: &RecordId(uuid_v7()), + version: "v0", + tag: "kv", + host: &HostId(uuid_v7()), + idx: &0, + }; + + let data = DecryptedData(vec![1, 2, 3, 4]); + + let encrypted = PASETO_V4::encrypt(data, ad, &key.to_bytes()); + let _ = PASETO_V4::decrypt(encrypted, ad, &fake_key.to_bytes()).unwrap_err(); + } + + #[test] + fn cannot_decrypt_different_id() { + let key = Key::<V4, Local>::new_os_random(); + + let ad = AdditionalData { + id: &RecordId(uuid_v7()), + version: "v0", + tag: "kv", + host: &HostId(uuid_v7()), + idx: &0, + }; + + let data = DecryptedData(vec![1, 2, 3, 4]); + + let encrypted = PASETO_V4::encrypt(data, ad, &key.to_bytes()); + + let ad = AdditionalData { + id: &RecordId(uuid_v7()), + ..ad + }; + let _ = PASETO_V4::decrypt(encrypted, ad, &key.to_bytes()).unwrap_err(); + } + + #[test] + fn re_encrypt_round_trip() { + let key1 = Key::<V4, Local>::new_os_random(); + let key2 = Key::<V4, Local>::new_os_random(); + + let ad = AdditionalData { + id: &RecordId(uuid_v7()), + version: "v0", + tag: "kv", + host: &HostId(uuid_v7()), + idx: &0, + }; + + let data = DecryptedData(vec![1, 2, 3, 4]); + + let encrypted1 = PASETO_V4::encrypt(data.clone(), ad, &key1.to_bytes()); + let encrypted2 = + PASETO_V4::re_encrypt(encrypted1.clone(), ad, &key1.to_bytes(), &key2.to_bytes()) + .unwrap(); + + // we only re-encrypt the content keys + assert_eq!(encrypted1.data, encrypted2.data); + assert_ne!( + encrypted1.content_encryption_key, + encrypted2.content_encryption_key + ); + + let decrypted = PASETO_V4::decrypt(encrypted2, ad, &key2.to_bytes()).unwrap(); + + assert_eq!(decrypted, data); + } + + #[test] + fn full_record_round_trip() { + let key = [0x55; 32]; + let record = Record::builder() + .id(RecordId(uuid_v7())) + .version("v0".to_owned()) + .tag("kv".to_owned()) + .host(Host::new(HostId(uuid_v7()))) + .timestamp(1687244806000000) + .data(DecryptedData(vec![1, 2, 3, 4])) + .idx(0) + .build(); + + let encrypted = record.encrypt::<PASETO_V4>(&key); + + assert!(!encrypted.data.data.is_empty()); + assert!(!encrypted.data.content_encryption_key.is_empty()); + + let decrypted = encrypted.decrypt::<PASETO_V4>(&key).unwrap(); + + assert_eq!(decrypted.data.0, [1, 2, 3, 4]); + } + + #[test] + fn full_record_round_trip_fail() { + let key = [0x55; 32]; + let record = Record::builder() + .id(RecordId(uuid_v7())) + .version("v0".to_owned()) + .tag("kv".to_owned()) + .host(Host::new(HostId(uuid_v7()))) + .timestamp(1687244806000000) + .data(DecryptedData(vec![1, 2, 3, 4])) + .idx(0) + .build(); + + let encrypted = record.encrypt::<PASETO_V4>(&key); + + let mut enc1 = encrypted.clone(); + enc1.host = Host::new(HostId(uuid_v7())); + let _ = enc1 + .decrypt::<PASETO_V4>(&key) + .expect_err("tampering with the host should result in auth failure"); + + let mut enc2 = encrypted; + enc2.id = RecordId(uuid_v7()); + let _ = enc2 + .decrypt::<PASETO_V4>(&key) + .expect_err("tampering with the id should result in auth failure"); + } +} diff --git a/crates/atuin-client/src/record/mod.rs b/crates/atuin-client/src/record/mod.rs new file mode 100644 index 00000000..c40fd395 --- /dev/null +++ b/crates/atuin-client/src/record/mod.rs @@ -0,0 +1,6 @@ +pub mod encryption; +pub mod sqlite_store; +pub mod store; + +#[cfg(feature = "sync")] +pub mod sync; diff --git a/crates/atuin-client/src/record/sqlite_store.rs b/crates/atuin-client/src/record/sqlite_store.rs new file mode 100644 index 00000000..ed51f3fd --- /dev/null +++ b/crates/atuin-client/src/record/sqlite_store.rs @@ -0,0 +1,642 @@ +// Here we are using sqlite as a pretty dumb store, and will not be running any complex queries. +// Multiple stores of multiple types are all stored in one chonky table (for now), and we just index +// by tag/host + +use std::str::FromStr; +use std::{path::Path, time::Duration}; + +use async_trait::async_trait; +use eyre::{Result, eyre}; +use fs_err as fs; + +use sqlx::{ + Row, + sqlite::{SqliteConnectOptions, SqliteJournalMode, SqlitePool, SqlitePoolOptions, SqliteRow}, +}; + +use atuin_common::record::{ + EncryptedData, Host, HostId, Record, RecordId, RecordIdx, RecordStatus, +}; +use atuin_common::utils; +use uuid::Uuid; + +use super::encryption::PASETO_V4; +use super::store::Store; + +#[derive(Debug, Clone)] +pub struct SqliteStore { + pool: SqlitePool, +} + +impl SqliteStore { + pub async fn new(path: impl AsRef<Path>, timeout: f64) -> Result<Self> { + let path = path.as_ref(); + + debug!("opening sqlite database at {path:?}"); + + if utils::broken_symlink(path) { + eprintln!( + "Atuin: Sqlite db path ({path:?}) is a broken symlink. Unable to read or create replacement." + ); + std::process::exit(1); + } + + if !path.exists() + && let Some(dir) = path.parent() + { + fs::create_dir_all(dir)?; + } + + let opts = SqliteConnectOptions::from_str(path.as_os_str().to_str().unwrap())? + .journal_mode(SqliteJournalMode::Wal) + .foreign_keys(true) + .create_if_missing(true); + + let pool = SqlitePoolOptions::new() + .acquire_timeout(Duration::from_secs_f64(timeout)) + .connect_with(opts) + .await?; + + Self::setup_db(&pool).await?; + + Ok(Self { pool }) + } + + async fn setup_db(pool: &SqlitePool) -> Result<()> { + debug!("running sqlite database setup"); + + sqlx::migrate!("./record-migrations").run(pool).await?; + + Ok(()) + } + + async fn save_raw( + tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, + r: &Record<EncryptedData>, + ) -> Result<()> { + // In sqlite, we are "limited" to i64. But that is still fine, until 2262. + sqlx::query( + "insert or ignore into store(id, idx, host, tag, timestamp, version, data, cek) + values(?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)", + ) + .bind(r.id.0.as_hyphenated().to_string()) + .bind(r.idx as i64) + .bind(r.host.id.0.as_hyphenated().to_string()) + .bind(r.tag.as_str()) + .bind(r.timestamp as i64) + .bind(r.version.as_str()) + .bind(r.data.data.as_str()) + .bind(r.data.content_encryption_key.as_str()) + .execute(&mut **tx) + .await?; + + Ok(()) + } + + fn query_row(row: SqliteRow) -> Record<EncryptedData> { + let idx: i64 = row.get("idx"); + let timestamp: i64 = row.get("timestamp"); + + // tbh at this point things are pretty fucked so just panic + let id = Uuid::from_str(row.get("id")).expect("invalid id UUID format in sqlite DB"); + let host = Uuid::from_str(row.get("host")).expect("invalid host UUID format in sqlite DB"); + + Record { + id: RecordId(id), + idx: idx as u64, + host: Host::new(HostId(host)), + timestamp: timestamp as u64, + tag: row.get("tag"), + version: row.get("version"), + data: EncryptedData { + data: row.get("data"), + content_encryption_key: row.get("cek"), + }, + } + } + + async fn load_all(&self) -> Result<Vec<Record<EncryptedData>>> { + let res = sqlx::query("select * from store ") + .map(Self::query_row) + .fetch_all(&self.pool) + .await?; + + Ok(res) + } +} + +#[async_trait] +impl Store for SqliteStore { + async fn push_batch( + &self, + records: impl Iterator<Item = &Record<EncryptedData>> + Send + Sync, + ) -> Result<()> { + let mut tx = self.pool.begin().await?; + + for record in records { + Self::save_raw(&mut tx, record).await?; + } + + tx.commit().await?; + + Ok(()) + } + + async fn get(&self, id: RecordId) -> Result<Record<EncryptedData>> { + let res = sqlx::query("select * from store where store.id = ?1") + .bind(id.0.as_hyphenated().to_string()) + .map(Self::query_row) + .fetch_one(&self.pool) + .await?; + + Ok(res) + } + + async fn delete(&self, id: RecordId) -> Result<()> { + sqlx::query("delete from store where id = ?1") + .bind(id.0.as_hyphenated().to_string()) + .execute(&self.pool) + .await?; + + Ok(()) + } + + async fn delete_all(&self) -> Result<()> { + sqlx::query("delete from store").execute(&self.pool).await?; + + Ok(()) + } + + async fn last(&self, host: HostId, tag: &str) -> Result<Option<Record<EncryptedData>>> { + let res = + sqlx::query("select * from store where host=?1 and tag=?2 order by idx desc limit 1") + .bind(host.0.as_hyphenated().to_string()) + .bind(tag) + .map(Self::query_row) + .fetch_one(&self.pool) + .await; + + match res { + Err(sqlx::Error::RowNotFound) => Ok(None), + Err(e) => Err(eyre!("an error occurred: {}", e)), + Ok(record) => Ok(Some(record)), + } + } + + async fn first(&self, host: HostId, tag: &str) -> Result<Option<Record<EncryptedData>>> { + self.idx(host, tag, 0).await + } + + async fn len_all(&self) -> Result<u64> { + let res: Result<(i64,), sqlx::Error> = sqlx::query_as("select count(*) from store") + .fetch_one(&self.pool) + .await; + match res { + Err(e) => Err(eyre!("failed to fetch local store len: {}", e)), + Ok(v) => Ok(v.0 as u64), + } + } + + async fn len_tag(&self, tag: &str) -> Result<u64> { + let res: Result<(i64,), sqlx::Error> = + sqlx::query_as("select count(*) from store where tag=?1") + .bind(tag) + .fetch_one(&self.pool) + .await; + match res { + Err(e) => Err(eyre!("failed to fetch local store len: {}", e)), + Ok(v) => Ok(v.0 as u64), + } + } + + async fn len(&self, host: HostId, tag: &str) -> Result<u64> { + let last = self.last(host, tag).await?; + + if let Some(last) = last { + return Ok(last.idx + 1); + } + + return Ok(0); + } + + async fn next( + &self, + host: HostId, + tag: &str, + idx: RecordIdx, + limit: u64, + ) -> Result<Vec<Record<EncryptedData>>> { + let res = sqlx::query( + "select * from store where idx >= ?1 and host = ?2 and tag = ?3 order by idx asc limit ?4", + ) + .bind(idx as i64) + .bind(host.0.as_hyphenated().to_string()) + .bind(tag) + .bind(limit as i64) + .map(Self::query_row) + .fetch_all(&self.pool) + .await?; + + Ok(res) + } + + async fn idx( + &self, + host: HostId, + tag: &str, + idx: RecordIdx, + ) -> Result<Option<Record<EncryptedData>>> { + let res = sqlx::query("select * from store where idx = ?1 and host = ?2 and tag = ?3") + .bind(idx as i64) + .bind(host.0.as_hyphenated().to_string()) + .bind(tag) + .map(Self::query_row) + .fetch_one(&self.pool) + .await; + + match res { + Err(sqlx::Error::RowNotFound) => Ok(None), + Err(e) => Err(eyre!("an error occurred: {}", e)), + Ok(v) => Ok(Some(v)), + } + } + + async fn status(&self) -> Result<RecordStatus> { + let mut status = RecordStatus::new(); + + let res: Result<Vec<(String, String, i64)>, sqlx::Error> = + sqlx::query_as("select host, tag, max(idx) from store group by host, tag") + .fetch_all(&self.pool) + .await; + + let res = match res { + Err(e) => return Err(eyre!("failed to fetch local store status: {}", e)), + Ok(v) => v, + }; + + for i in res { + let host = HostId( + Uuid::from_str(i.0.as_str()).expect("failed to parse uuid for local store status"), + ); + + status.set_raw(host, i.1, i.2 as u64); + } + + Ok(status) + } + + async fn all_tagged(&self, tag: &str) -> Result<Vec<Record<EncryptedData>>> { + let res = sqlx::query("select * from store where tag = ?1 order by timestamp asc") + .bind(tag) + .map(Self::query_row) + .fetch_all(&self.pool) + .await?; + + Ok(res) + } + + /// Reencrypt every single item in this store with a new key + /// Be careful - this may mess with sync. + async fn re_encrypt(&self, old_key: &[u8; 32], new_key: &[u8; 32]) -> Result<()> { + // Load all the records + // In memory like some of the other code here + // This will never be called in a hot loop, and only under the following circumstances + // 1. The user has logged into a new account, with a new key. They are unlikely to have a + // lot of data + // 2. The user has encountered some sort of issue, and runs a maintenance command that + // invokes this + let all = self.load_all().await?; + + let re_encrypted = all + .into_iter() + .map(|record| record.re_encrypt::<PASETO_V4>(old_key, new_key)) + .collect::<Result<Vec<_>>>()?; + + // next up, we delete all the old data and reinsert the new stuff + // do it in one transaction, so if anything fails we rollback OK + + let mut tx = self.pool.begin().await?; + + let res = sqlx::query("delete from store").execute(&mut *tx).await?; + + let rows = res.rows_affected(); + debug!("deleted {rows} rows"); + + // don't call push_batch, as it will start its own transaction + // call the underlying save_raw + + for record in re_encrypted { + Self::save_raw(&mut tx, &record).await?; + } + + tx.commit().await?; + + Ok(()) + } + + /// Verify that every record in this store can be decrypted with the current key + /// Someday maybe also check each tag/record can be deserialized, but not for now. + async fn verify(&self, key: &[u8; 32]) -> Result<()> { + let all = self.load_all().await?; + + all.into_iter() + .map(|record| record.decrypt::<PASETO_V4>(key)) + .collect::<Result<Vec<_>>>()?; + + Ok(()) + } + + /// Verify that every record in this store can be decrypted with the current key + /// Someday maybe also check each tag/record can be deserialized, but not for now. + async fn purge(&self, key: &[u8; 32]) -> Result<()> { + let all = self.load_all().await?; + + for record in all.iter() { + match record.clone().decrypt::<PASETO_V4>(key) { + Ok(_) => continue, + Err(_) => { + println!( + "Failed to decrypt {}, deleting", + record.id.0.as_hyphenated() + ); + + self.delete(record.id).await?; + } + } + } + + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use atuin_common::{ + record::{DecryptedData, EncryptedData, Host, HostId, Record}, + utils::uuid_v7, + }; + + use crate::{ + encryption::generate_encoded_key, + record::{encryption::PASETO_V4, store::Store}, + settings::test_local_timeout, + }; + + use super::SqliteStore; + + fn test_record() -> Record<EncryptedData> { + Record::builder() + .host(Host::new(HostId(atuin_common::utils::uuid_v7()))) + .version("v1".into()) + .tag(atuin_common::utils::uuid_v7().simple().to_string()) + .data(EncryptedData { + data: "1234".into(), + content_encryption_key: "1234".into(), + }) + .idx(0) + .build() + } + + #[tokio::test] + async fn create_db() { + let db = SqliteStore::new(":memory:", test_local_timeout()).await; + + assert!( + db.is_ok(), + "db could not be created, {:?}", + db.err().unwrap() + ); + } + + #[tokio::test] + async fn push_record() { + let db = SqliteStore::new(":memory:", test_local_timeout()) + .await + .unwrap(); + let record = test_record(); + + db.push(&record).await.expect("failed to insert record"); + } + + #[tokio::test] + async fn get_record() { + let db = SqliteStore::new(":memory:", test_local_timeout()) + .await + .unwrap(); + let record = test_record(); + db.push(&record).await.unwrap(); + + let new_record = db.get(record.id).await.expect("failed to fetch record"); + + assert_eq!(record, new_record, "records are not equal"); + } + + #[tokio::test] + async fn last() { + let db = SqliteStore::new(":memory:", test_local_timeout()) + .await + .unwrap(); + let record = test_record(); + db.push(&record).await.unwrap(); + + let last = db + .last(record.host.id, record.tag.as_str()) + .await + .expect("failed to get store len"); + + assert_eq!( + last.unwrap().id, + record.id, + "expected to get back the same record that was inserted" + ); + } + + #[tokio::test] + async fn first() { + let db = SqliteStore::new(":memory:", test_local_timeout()) + .await + .unwrap(); + let record = test_record(); + db.push(&record).await.unwrap(); + + let first = db + .first(record.host.id, record.tag.as_str()) + .await + .expect("failed to get store len"); + + assert_eq!( + first.unwrap().id, + record.id, + "expected to get back the same record that was inserted" + ); + } + + #[tokio::test] + async fn len() { + let db = SqliteStore::new(":memory:", test_local_timeout()) + .await + .unwrap(); + let record = test_record(); + db.push(&record).await.unwrap(); + + let len = db + .len(record.host.id, record.tag.as_str()) + .await + .expect("failed to get store len"); + + assert_eq!(len, 1, "expected length of 1 after insert"); + } + + #[tokio::test] + async fn len_tag() { + let db = SqliteStore::new(":memory:", test_local_timeout()) + .await + .unwrap(); + let record = test_record(); + db.push(&record).await.unwrap(); + + let len = db + .len_tag(record.tag.as_str()) + .await + .expect("failed to get store len"); + + assert_eq!(len, 1, "expected length of 1 after insert"); + } + + #[tokio::test] + async fn len_different_tags() { + let db = SqliteStore::new(":memory:", test_local_timeout()) + .await + .unwrap(); + + // these have different tags, so the len should be the same + // we model multiple stores within one database + // new store = new tag = independent length + let first = test_record(); + let second = test_record(); + + db.push(&first).await.unwrap(); + db.push(&second).await.unwrap(); + + let first_len = db.len(first.host.id, first.tag.as_str()).await.unwrap(); + let second_len = db.len(second.host.id, second.tag.as_str()).await.unwrap(); + + assert_eq!(first_len, 1, "expected length of 1 after insert"); + assert_eq!(second_len, 1, "expected length of 1 after insert"); + } + + #[tokio::test] + async fn append_a_bunch() { + let db = SqliteStore::new(":memory:", test_local_timeout()) + .await + .unwrap(); + + let mut tail = test_record(); + db.push(&tail).await.expect("failed to push record"); + + for _ in 1..100 { + tail = tail.append(vec![1, 2, 3, 4]).encrypt::<PASETO_V4>(&[0; 32]); + db.push(&tail).await.unwrap(); + } + + assert_eq!( + db.len(tail.host.id, tail.tag.as_str()).await.unwrap(), + 100, + "failed to insert 100 records" + ); + + assert_eq!( + db.len_tag(tail.tag.as_str()).await.unwrap(), + 100, + "failed to insert 100 records" + ); + } + + #[tokio::test] + async fn append_a_big_bunch() { + let db = SqliteStore::new(":memory:", test_local_timeout()) + .await + .unwrap(); + + let mut records: Vec<Record<EncryptedData>> = Vec::with_capacity(10000); + + let mut tail = test_record(); + records.push(tail.clone()); + + for _ in 1..10000 { + tail = tail.append(vec![1, 2, 3]).encrypt::<PASETO_V4>(&[0; 32]); + records.push(tail.clone()); + } + + db.push_batch(records.iter()).await.unwrap(); + + assert_eq!( + db.len(tail.host.id, tail.tag.as_str()).await.unwrap(), + 10000, + "failed to insert 10k records" + ); + } + + #[tokio::test] + async fn re_encrypt() { + let store = SqliteStore::new(":memory:", test_local_timeout()) + .await + .unwrap(); + let (key, _) = generate_encoded_key().unwrap(); + let data = vec![0u8, 1u8, 2u8, 3u8]; + let host_id = HostId(uuid_v7()); + + for i in 0..10 { + let record = Record::builder() + .host(Host::new(host_id)) + .version(String::from("test")) + .tag(String::from("test")) + .idx(i) + .data(DecryptedData(data.clone())) + .build(); + + let record = record.encrypt::<PASETO_V4>(&key.into()); + store + .push(&record) + .await + .expect("failed to push encrypted record"); + } + + // first, check that we can decrypt the data with the current key + let all = store.all_tagged("test").await.unwrap(); + + assert_eq!(all.len(), 10, "failed to fetch all records"); + + for record in all { + let decrypted = record.decrypt::<PASETO_V4>(&key.into()).unwrap(); + assert_eq!(decrypted.data.0, data); + } + + // reencrypt the store, then check if + // 1) it cannot be decrypted with the old key + // 2) it can be decrypted with the new key + + let (new_key, _) = generate_encoded_key().unwrap(); + store + .re_encrypt(&key.into(), &new_key.into()) + .await + .expect("failed to re-encrypt store"); + + let all = store.all_tagged("test").await.unwrap(); + + for record in all.iter() { + let decrypted = record.clone().decrypt::<PASETO_V4>(&key.into()); + assert!( + decrypted.is_err(), + "did not get error decrypting with old key after re-encrypt" + ) + } + + for record in all { + let decrypted = record.decrypt::<PASETO_V4>(&new_key.into()).unwrap(); + assert_eq!(decrypted.data.0, data); + } + + assert_eq!(store.len(host_id, "test").await.unwrap(), 10); + } +} diff --git a/crates/atuin-client/src/record/store.rs b/crates/atuin-client/src/record/store.rs new file mode 100644 index 00000000..49ca4968 --- /dev/null +++ b/crates/atuin-client/src/record/store.rs @@ -0,0 +1,60 @@ +use async_trait::async_trait; +use eyre::Result; + +use atuin_common::record::{EncryptedData, HostId, Record, RecordId, RecordIdx, RecordStatus}; + +/// A record store stores records +/// In more detail - we tend to need to process this into _another_ format to actually query it. +/// As is, the record store is intended as the source of truth for arbitrary data, which could +/// be shell history, kvs, etc. +#[async_trait] +pub trait Store { + // Push a record + async fn push(&self, record: &Record<EncryptedData>) -> Result<()> { + self.push_batch(std::iter::once(record)).await + } + + // Push a batch of records, all in one transaction + async fn push_batch( + &self, + records: impl Iterator<Item = &Record<EncryptedData>> + Send + Sync, + ) -> Result<()>; + + async fn get(&self, id: RecordId) -> Result<Record<EncryptedData>>; + + async fn delete(&self, id: RecordId) -> Result<()>; + async fn delete_all(&self) -> Result<()>; + + async fn len_all(&self) -> Result<u64>; + async fn len(&self, host: HostId, tag: &str) -> Result<u64>; + async fn len_tag(&self, tag: &str) -> Result<u64>; + + async fn last(&self, host: HostId, tag: &str) -> Result<Option<Record<EncryptedData>>>; + async fn first(&self, host: HostId, tag: &str) -> Result<Option<Record<EncryptedData>>>; + + async fn re_encrypt(&self, old_key: &[u8; 32], new_key: &[u8; 32]) -> Result<()>; + async fn verify(&self, key: &[u8; 32]) -> Result<()>; + async fn purge(&self, key: &[u8; 32]) -> Result<()>; + + /// Get the next `limit` records, after and including the given index + async fn next( + &self, + host: HostId, + tag: &str, + idx: RecordIdx, + limit: u64, + ) -> Result<Vec<Record<EncryptedData>>>; + + /// Get the first record for a given host and tag + async fn idx( + &self, + host: HostId, + tag: &str, + idx: RecordIdx, + ) -> Result<Option<Record<EncryptedData>>>; + + async fn status(&self) -> Result<RecordStatus>; + + /// Get all records for a given tag + async fn all_tagged(&self, tag: &str) -> Result<Vec<Record<EncryptedData>>>; +} diff --git a/crates/atuin-client/src/record/sync.rs b/crates/atuin-client/src/record/sync.rs new file mode 100644 index 00000000..37840b75 --- /dev/null +++ b/crates/atuin-client/src/record/sync.rs @@ -0,0 +1,628 @@ +// do a sync :O +use std::{cmp::Ordering, fmt::Write}; + +use eyre::Result; +use thiserror::Error; + +use super::store::Store; +use crate::{api_client::Client, settings::Settings}; + +use atuin_common::record::{Diff, HostId, RecordId, RecordIdx, RecordStatus}; +use indicatif::{ProgressBar, ProgressState, ProgressStyle}; + +#[derive(Error, Debug)] +pub enum SyncError { + #[error("the local store is ahead of the remote, but for another host. has remote lost data?")] + LocalAheadOtherHost, + + #[error("an issue with the local database occurred: {msg:?}")] + LocalStoreError { msg: String }, + + #[error("something has gone wrong with the sync logic: {msg:?}")] + SyncLogicError { msg: String }, + + #[error("operational error: {msg:?}")] + OperationalError { msg: String }, + + #[error("a request to the sync server failed: {msg:?}")] + RemoteRequestError { msg: String }, +} + +#[derive(Debug, Eq, PartialEq)] +pub enum Operation { + // Either upload or download until the states matches the below + Upload { + local: RecordIdx, + remote: Option<RecordIdx>, + host: HostId, + tag: String, + }, + Download { + local: Option<RecordIdx>, + remote: RecordIdx, + host: HostId, + tag: String, + }, + Noop { + host: HostId, + tag: String, + }, +} + +pub async fn diff( + settings: &Settings, + store: &impl Store, +) -> Result<(Vec<Diff>, RecordStatus), SyncError> { + let client = Client::new( + &settings.sync_address, + settings + .sync_auth_token() + .await + .map_err(|e| SyncError::RemoteRequestError { msg: e.to_string() })?, + settings.network_connect_timeout, + settings.network_timeout, + ) + .map_err(|e| SyncError::OperationalError { msg: e.to_string() })?; + + let local_index = store + .status() + .await + .map_err(|e| SyncError::LocalStoreError { msg: e.to_string() })?; + + let remote_index = client + .record_status() + .await + .map_err(|e| SyncError::RemoteRequestError { msg: e.to_string() })?; + + let diff = local_index.diff(&remote_index); + + Ok((diff, remote_index)) +} + +// Take a diff, along with a local store, and resolve it into a set of operations. +// With the store as context, we can determine if a tail exists locally or not and therefore if it needs uploading or download. +// In theory this could be done as a part of the diffing stage, but it's easier to reason +// about and test this way +pub async fn operations( + diffs: Vec<Diff>, + _store: &impl Store, +) -> Result<Vec<Operation>, SyncError> { + let mut operations = Vec::with_capacity(diffs.len()); + + for diff in diffs { + let op = match (diff.local, diff.remote) { + // We both have it! Could be either. Compare. + (Some(local), Some(remote)) => match local.cmp(&remote) { + Ordering::Equal => Operation::Noop { + host: diff.host, + tag: diff.tag, + }, + Ordering::Greater => Operation::Upload { + local, + remote: Some(remote), + host: diff.host, + tag: diff.tag, + }, + Ordering::Less => Operation::Download { + local: Some(local), + remote, + host: diff.host, + tag: diff.tag, + }, + }, + + // Remote has it, we don't. Gotta be download + (None, Some(remote)) => Operation::Download { + local: None, + remote, + host: diff.host, + tag: diff.tag, + }, + + // We have it, remote doesn't. Gotta be upload. + (Some(local), None) => Operation::Upload { + local, + remote: None, + host: diff.host, + tag: diff.tag, + }, + + // something is pretty fucked. + (None, None) => { + return Err(SyncError::SyncLogicError { + msg: String::from( + "diff has nothing for local or remote - (host, tag) does not exist", + ), + }); + } + }; + + operations.push(op); + } + + // sort them - purely so we have a stable testing order, and can rely on + // same input = same output + // We can sort by ID so long as we continue to use UUIDv7 or something + // with the same properties + + operations.sort_by_key(|op| match op { + Operation::Noop { host, tag } => (0, *host, tag.clone()), + + Operation::Upload { host, tag, .. } => (1, *host, tag.clone()), + + Operation::Download { host, tag, .. } => (2, *host, tag.clone()), + }); + + Ok(operations) +} + +async fn sync_upload( + store: &impl Store, + client: &Client<'_>, + host: HostId, + tag: String, + local: RecordIdx, + remote: Option<RecordIdx>, + page_size: u64, +) -> Result<i64, SyncError> { + let remote = remote.unwrap_or(0); + let expected = local - remote; + let mut progress = 0; + + let pb = ProgressBar::new(expected); + pb.set_style(ProgressStyle::with_template("{spinner:.green} [{elapsed_precise}] [{wide_bar:.cyan/blue}] {human_pos}/{human_len} ({eta})") + .unwrap() + .with_key("eta", |state: &ProgressState, w: &mut dyn Write| write!(w, "{:.1}s", state.eta().as_secs_f64()).unwrap()) + .progress_chars("#>-")); + + println!( + "Uploading {} records to {}/{}", + expected, + host.0.as_simple(), + tag + ); + + loop { + let page = store + .next(host, tag.as_str(), remote + progress, page_size) + .await + .map_err(|e| { + error!("failed to read upload page: {e:?}"); + + SyncError::LocalStoreError { msg: e.to_string() } + })?; + + if page.is_empty() { + break; + } + + client.post_records(&page).await.map_err(|e| { + error!("failed to post records: {e:?}"); + + SyncError::RemoteRequestError { msg: e.to_string() } + })?; + + progress += page.len() as u64; + pb.set_position(progress); + + if progress >= expected { + break; + } + } + + pb.finish_with_message("Uploaded records"); + + Ok(progress as i64) +} + +async fn sync_download( + store: &impl Store, + client: &Client<'_>, + host: HostId, + tag: String, + local: Option<RecordIdx>, + remote: RecordIdx, + page_size: u64, +) -> Result<Vec<RecordId>, SyncError> { + let local = local.unwrap_or(0); + let expected = remote - local; + let mut progress = 0; + let mut ret = Vec::new(); + + println!( + "Downloading {} records from {}/{}", + expected, + host.0.as_simple(), + tag + ); + + let pb = ProgressBar::new(expected); + pb.set_style(ProgressStyle::with_template("{spinner:.green} [{elapsed_precise}] [{wide_bar:.cyan/blue}] {human_pos}/{human_len} ({eta})") + .unwrap() + .with_key("eta", |state: &ProgressState, w: &mut dyn Write| write!(w, "{:.1}s", state.eta().as_secs_f64()).unwrap()) + .progress_chars("#>-")); + + loop { + let page = client + .next_records(host, tag.clone(), local + progress, page_size) + .await + .map_err(|e| SyncError::RemoteRequestError { msg: e.to_string() })?; + + if page.is_empty() { + break; + } + + store + .push_batch(page.iter()) + .await + .map_err(|e| SyncError::LocalStoreError { msg: e.to_string() })?; + + ret.extend(page.iter().map(|f| f.id)); + + progress += page.len() as u64; + pb.set_position(progress); + + if progress >= expected { + break; + } + } + + pb.finish_with_message("Downloaded records"); + + Ok(ret) +} + +pub async fn sync_remote( + operations: Vec<Operation>, + local_store: &impl Store, + settings: &Settings, + page_size: u64, +) -> Result<(i64, Vec<RecordId>), SyncError> { + let client = Client::new( + &settings.sync_address, + settings + .sync_auth_token() + .await + .map_err(|e| SyncError::RemoteRequestError { msg: e.to_string() })?, + settings.network_connect_timeout, + settings.network_timeout, + ) + .expect("failed to create client"); + + let mut uploaded = 0; + let mut downloaded = Vec::new(); + + // this can totally run in parallel, but lets get it working first + for i in operations { + match i { + Operation::Upload { + host, + tag, + local, + remote, + } => { + uploaded += + sync_upload(local_store, &client, host, tag, local, remote, page_size).await? + } + + Operation::Download { + host, + tag, + local, + remote, + } => { + let mut d = + sync_download(local_store, &client, host, tag, local, remote, page_size) + .await?; + downloaded.append(&mut d) + } + + Operation::Noop { .. } => continue, + } + } + + Ok((uploaded, downloaded)) +} + +pub async fn sync( + settings: &Settings, + store: &impl Store, +) -> Result<(i64, Vec<RecordId>), SyncError> { + let (diff, _) = diff(settings, store).await?; + let operations = operations(diff, store).await?; + let (uploaded, downloaded) = sync_remote(operations, store, settings, 100).await?; + + Ok((uploaded, downloaded)) +} + +#[cfg(test)] +mod tests { + use atuin_common::record::{Diff, EncryptedData, HostId, Record}; + use pretty_assertions::assert_eq; + + use crate::{ + record::{ + encryption::PASETO_V4, + sqlite_store::SqliteStore, + store::Store, + sync::{self, Operation}, + }, + settings::test_local_timeout, + }; + + fn test_record() -> Record<EncryptedData> { + Record::builder() + .host(atuin_common::record::Host::new(HostId( + atuin_common::utils::uuid_v7(), + ))) + .version("v1".into()) + .tag(atuin_common::utils::uuid_v7().simple().to_string()) + .data(EncryptedData { + data: String::new(), + content_encryption_key: String::new(), + }) + .idx(0) + .build() + } + + // Take a list of local records, and a list of remote records. + // Return the local database, and a diff of local/remote, ready to build + // ops + async fn build_test_diff( + local_records: Vec<Record<EncryptedData>>, + remote_records: Vec<Record<EncryptedData>>, + ) -> (SqliteStore, Vec<Diff>) { + let local_store = SqliteStore::new(":memory:", test_local_timeout()) + .await + .expect("failed to open in memory sqlite"); + let remote_store = SqliteStore::new(":memory:", test_local_timeout()) + .await + .expect("failed to open in memory sqlite"); // "remote" + + for i in local_records { + local_store.push(&i).await.unwrap(); + } + + for i in remote_records { + remote_store.push(&i).await.unwrap(); + } + + let local_index = local_store.status().await.unwrap(); + let remote_index = remote_store.status().await.unwrap(); + + let diff = local_index.diff(&remote_index); + + (local_store, diff) + } + + #[tokio::test] + async fn test_basic_diff() { + // a diff where local is ahead of remote. nothing else. + + let record = test_record(); + let (store, diff) = build_test_diff(vec![record.clone()], vec![]).await; + + assert_eq!(diff.len(), 1); + + let operations = sync::operations(diff, &store).await.unwrap(); + + assert_eq!(operations.len(), 1); + + assert_eq!( + operations[0], + Operation::Upload { + host: record.host.id, + tag: record.tag, + local: record.idx, + remote: None, + } + ); + } + + #[tokio::test] + async fn build_two_way_diff() { + // a diff where local is ahead of remote for one, and remote for + // another. One upload, one download + + let shared_record = test_record(); + let remote_ahead = test_record(); + + let local_ahead = shared_record + .append(vec![1, 2, 3]) + .encrypt::<PASETO_V4>(&[0; 32]); + + assert_eq!(local_ahead.idx, 1); + + let local = vec![shared_record.clone(), local_ahead.clone()]; // local knows about the already synced, and something newer in the same store + let remote = vec![shared_record.clone(), remote_ahead.clone()]; // remote knows about the already-synced, and one new record in a new store + + let (store, diff) = build_test_diff(local, remote).await; + let operations = sync::operations(diff, &store).await.unwrap(); + + assert_eq!(operations.len(), 2); + + assert_eq!( + operations, + vec![ + // Or in otherwords, local is ahead by one + Operation::Upload { + host: local_ahead.host.id, + tag: local_ahead.tag, + local: 1, + remote: Some(0), + }, + // Or in other words, remote knows of a record in an entirely new store (tag) + Operation::Download { + host: remote_ahead.host.id, + tag: remote_ahead.tag, + local: None, + remote: 0, + }, + ] + ); + } + + #[tokio::test] + async fn build_complex_diff() { + // One shared, ahead but known only by remote + // One known only by local + // One known only by remote + + let shared_record = test_record(); + let local_only = test_record(); + + let local_only_20 = test_record(); + let local_only_21 = local_only_20 + .append(vec![1, 2, 3]) + .encrypt::<PASETO_V4>(&[0; 32]); + let local_only_22 = local_only_21 + .append(vec![1, 2, 3]) + .encrypt::<PASETO_V4>(&[0; 32]); + let local_only_23 = local_only_22 + .append(vec![1, 2, 3]) + .encrypt::<PASETO_V4>(&[0; 32]); + + let remote_only = test_record(); + + let remote_only_20 = test_record(); + let remote_only_21 = remote_only_20 + .append(vec![2, 3, 2]) + .encrypt::<PASETO_V4>(&[0; 32]); + let remote_only_22 = remote_only_21 + .append(vec![2, 3, 2]) + .encrypt::<PASETO_V4>(&[0; 32]); + let remote_only_23 = remote_only_22 + .append(vec![2, 3, 2]) + .encrypt::<PASETO_V4>(&[0; 32]); + let remote_only_24 = remote_only_23 + .append(vec![2, 3, 2]) + .encrypt::<PASETO_V4>(&[0; 32]); + + let second_shared = test_record(); + let second_shared_remote_ahead = second_shared + .append(vec![1, 2, 3]) + .encrypt::<PASETO_V4>(&[0; 32]); + let second_shared_remote_ahead2 = second_shared_remote_ahead + .append(vec![1, 2, 3]) + .encrypt::<PASETO_V4>(&[0; 32]); + + let third_shared = test_record(); + let third_shared_local_ahead = third_shared + .append(vec![1, 2, 3]) + .encrypt::<PASETO_V4>(&[0; 32]); + let third_shared_local_ahead2 = third_shared_local_ahead + .append(vec![1, 2, 3]) + .encrypt::<PASETO_V4>(&[0; 32]); + + let fourth_shared = test_record(); + let fourth_shared_remote_ahead = fourth_shared + .append(vec![1, 2, 3]) + .encrypt::<PASETO_V4>(&[0; 32]); + let fourth_shared_remote_ahead2 = fourth_shared_remote_ahead + .append(vec![1, 2, 3]) + .encrypt::<PASETO_V4>(&[0; 32]); + + let local = vec![ + shared_record.clone(), + second_shared.clone(), + third_shared.clone(), + fourth_shared.clone(), + fourth_shared_remote_ahead.clone(), + // single store, only local has it + local_only.clone(), + // bigger store, also only known by local + local_only_20.clone(), + local_only_21.clone(), + local_only_22.clone(), + local_only_23.clone(), + // another shared store, but local is ahead on this one + third_shared_local_ahead.clone(), + third_shared_local_ahead2.clone(), + ]; + + let remote = vec![ + remote_only.clone(), + remote_only_20.clone(), + remote_only_21.clone(), + remote_only_22.clone(), + remote_only_23.clone(), + remote_only_24.clone(), + shared_record.clone(), + second_shared.clone(), + third_shared.clone(), + second_shared_remote_ahead.clone(), + second_shared_remote_ahead2.clone(), + fourth_shared.clone(), + fourth_shared_remote_ahead.clone(), + fourth_shared_remote_ahead2.clone(), + ]; // remote knows about the already-synced, and one new record in a new store + + let (store, diff) = build_test_diff(local, remote).await; + let operations = sync::operations(diff, &store).await.unwrap(); + + assert_eq!(operations.len(), 7); + + let mut result_ops = vec![ + // We started with a shared record, but the remote knows of two newer records in the + // same store + Operation::Download { + local: Some(0), + remote: 2, + host: second_shared_remote_ahead.host.id, + tag: second_shared_remote_ahead.tag, + }, + // We have a shared record, local knows of the first two but not the last + Operation::Download { + local: Some(1), + remote: 2, + host: fourth_shared_remote_ahead2.host.id, + tag: fourth_shared_remote_ahead2.tag, + }, + // Remote knows of a store with a single record that local does not have + Operation::Download { + local: None, + remote: 0, + host: remote_only.host.id, + tag: remote_only.tag, + }, + // Remote knows of a store with a bunch of records that local does not have + Operation::Download { + local: None, + remote: 4, + host: remote_only_20.host.id, + tag: remote_only_20.tag, + }, + // Local knows of a record in a store that remote does not have + Operation::Upload { + local: 0, + remote: None, + host: local_only.host.id, + tag: local_only.tag, + }, + // Local knows of 4 records in a store that remote does not have + Operation::Upload { + local: 3, + remote: None, + host: local_only_20.host.id, + tag: local_only_20.tag, + }, + // Local knows of 2 more records in a shared store that remote only has one of + Operation::Upload { + local: 2, + remote: Some(0), + host: third_shared.host.id, + tag: third_shared.tag, + }, + ]; + + result_ops.sort_by_key(|op| match op { + Operation::Noop { host, tag } => (0, *host, tag.clone()), + + Operation::Upload { host, tag, .. } => (1, *host, tag.clone()), + + Operation::Download { host, tag, .. } => (2, *host, tag.clone()), + }); + + assert_eq!(result_ops, operations); + } +} diff --git a/crates/atuin-client/src/register.rs b/crates/atuin-client/src/register.rs new file mode 100644 index 00000000..ad077dd1 --- /dev/null +++ b/crates/atuin-client/src/register.rs @@ -0,0 +1,20 @@ +use eyre::Result; + +use crate::{api_client, settings::Settings}; + +pub async fn register_classic( + settings: &Settings, + username: String, + email: String, + password: String, +) -> Result<String> { + let session = + api_client::register(settings.sync_address.as_str(), &username, &email, &password).await?; + + let meta = Settings::meta_store().await?; + meta.save_session(&session.session).await?; + + let _key = crate::encryption::load_key(settings)?; + + Ok(session.session) +} diff --git a/crates/atuin-client/src/secrets.rs b/crates/atuin-client/src/secrets.rs new file mode 100644 index 00000000..459e6238 --- /dev/null +++ b/crates/atuin-client/src/secrets.rs @@ -0,0 +1,167 @@ +// This file will probably trigger a lot of scanners. Sorry. + +use regex::RegexSet; +use std::sync::LazyLock; + +pub enum TestValue<'a> { + Single(&'a str), + Multiple(&'a [&'a str]), +} + +/// A list of `(name, regex, test)`, where `test` should match against `regex`. +pub static SECRET_PATTERNS: &[(&str, &str, TestValue)] = &[ + ( + "AWS Access Key ID", + "A[KS]IA[0-9A-Z]{16}", + TestValue::Single("AKIAIOSFODNN7EXAMPLE"), + ), + ( + "AWS Secret Access Key env var", + "AWS_SECRET_ACCESS_KEY", + TestValue::Single("AWS_SECRET_ACCESS_KEY=KEYDATA"), + ), + ( + "AWS Session Token env var", + "AWS_SESSION_TOKEN", + TestValue::Single("AWS_SESSION_TOKEN=KEYDATA"), + ), + ( + "Microsoft Azure secret access key env var", + "AZURE_.*_KEY", + TestValue::Single("export AZURE_STORAGE_ACCOUNT_KEY=KEYDATA"), + ), + ( + "Google cloud platform key env var", + "GOOGLE_SERVICE_ACCOUNT_KEY", + TestValue::Single("export GOOGLE_SERVICE_ACCOUNT_KEY=KEYDATA"), + ), + ( + "Atuin login", + r"atuin\s+login", + TestValue::Single( + "atuin login -u mycoolusername -p mycoolpassword -k \"lots of random words\"", + ), + ), + ( + "GitHub PAT (old)", + "ghp_[a-zA-Z0-9]{36}", + TestValue::Single("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"), // legit, I expired it + ), + ( + "GitHub PAT (new)", + "gh1_[A-Za-z0-9]{21}_[A-Za-z0-9]{59}|github_pat_[0-9][A-Za-z0-9]{21}_[A-Za-z0-9]{59}", + TestValue::Multiple(&[ + "gh1_1234567890abcdefghijk_1234567890abcdefghijklmnopqrstuvwxyz1234567890abcdefghijklm", + "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", // also legit, also expired + ]), + ), + ( + "GitHub OAuth Access Token", + "gho_[A-Za-z0-9]{36}", + TestValue::Single("gho_1234567890abcdefghijklmnopqrstuvwx000"), // not a real token + ), + ( + "GitHub OAuth Access Token (user)", + "ghu_[A-Za-z0-9]{36}", + TestValue::Single("ghu_1234567890abcdefghijklmnopqrstuvwx000"), // not a real token + ), + ( + "GitHub App Installation Access Token", + "ghs_[A-Za-z0-9]{36}", + TestValue::Single("ghs_1234567890abcdefghijklmnopqrstuvwx000"), // not a real token + ), + ( + "GitHub Refresh Token", + "ghr_[A-Za-z0-9]{76}", + TestValue::Single( + "ghr_1234567890abcdefghijklmnopqrstuvwx1234567890abcdefghijklmnopqrstuvwx1234567890abcdefghijklmnopqrstuvwx", + ), // not a real token + ), + ( + "GitHub App Installation Access Token v1", + "v1\\.[0-9A-Fa-f]{40}", + TestValue::Single("v1.1234567890abcdef1234567890abcdef12345678"), // not a real token + ), + ( + "GitLab PAT", + "glpat-[a-zA-Z0-9_]{20}", + TestValue::Single("glpat-RkE_BG5p_bbjML21WSfy"), + ), + ( + "Slack OAuth v2 bot", + "xoxb-[0-9]{11}-[0-9]{11}-[0-9a-zA-Z]{24}", + TestValue::Single("xoxb-17653672481-19874698323-pdFZKVeTuE8sk7oOcBrzbqgy"), + ), + ( + "Slack OAuth v2 user token", + "xoxp-[0-9]{11}-[0-9]{11}-[0-9a-zA-Z]{24}", + TestValue::Single("xoxp-17653672481-19874698323-pdFZKVeTuE8sk7oOcBrzbqgy"), + ), + ( + "Slack webhook", + "T[a-zA-Z0-9_]{8}/B[a-zA-Z0-9_]{8}/[a-zA-Z0-9_]{24}", + TestValue::Single( + "https://hooks.slack.com/services/T00000000/B00000000/XXXXXXXXXXXXXXXXXXXXXXXX", + ), + ), + ( + "Stripe test key", + "sk_test_[0-9a-zA-Z]{24}", + TestValue::Single("sk_test_1234567890abcdefghijklmnop"), + ), + ( + "Stripe live key", + "sk_live_[0-9a-zA-Z]{24}", + TestValue::Single("sk_live_1234567890abcdefghijklmnop"), + ), + ( + "Netlify authentication token", + "nf[pcoub]_[0-9a-zA-Z]{36}", + TestValue::Single("nfp_nBh7BdJxUwyaBBwFzpyD29MMFT6pZ9wq5634"), + ), + ( + "npm token", + "npm_[A-Za-z0-9]{36}", + TestValue::Single("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"), + ), + ( + "Pulumi personal access token", + "pul-[0-9a-f]{40}", + TestValue::Single("pul-683c2770662c51d960d72ec27613be7653c5cb26"), + ), +]; + +/// The `regex` expressions from [`SECRET_PATTERNS`] compiled into a `RegexSet`. +pub static SECRET_PATTERNS_RE: LazyLock<RegexSet> = LazyLock::new(|| { + let exprs = SECRET_PATTERNS.iter().map(|f| f.1); + RegexSet::new(exprs).expect("Failed to build secrets regex") +}); + +#[cfg(test)] +mod tests { + use regex::Regex; + + use crate::secrets::{SECRET_PATTERNS, TestValue}; + + #[test] + fn test_secrets() { + for (name, regex, test) in SECRET_PATTERNS { + let re = + Regex::new(regex).unwrap_or_else(|_| panic!("Failed to compile regex for {name}")); + + match test { + TestValue::Single(test) => { + assert!(re.is_match(test), "{name} test failed!"); + } + TestValue::Multiple(tests) => { + for test_str in tests.iter() { + assert!( + re.is_match(test_str), + "{name} test with value \"{test_str}\" failed!" + ); + } + } + } + } + } +} diff --git a/crates/atuin-client/src/settings.rs b/crates/atuin-client/src/settings.rs new file mode 100644 index 00000000..2a96a2b3 --- /dev/null +++ b/crates/atuin-client/src/settings.rs @@ -0,0 +1,1844 @@ +use std::{collections::HashMap, fmt, io::prelude::*, path::PathBuf, str::FromStr, sync::OnceLock}; +use tokio::sync::OnceCell; + +use atuin_common::record::HostId; +use atuin_common::utils; +use clap::ValueEnum; +use config::{ + Config, ConfigBuilder, Environment, File as ConfigFile, FileFormat, builder::DefaultState, +}; +use eyre::{Context, Error, Result, bail, eyre}; +use fs_err::{File, create_dir_all}; +use humantime::parse_duration; +use regex::RegexSet; +use semver::Version; +use serde::{Deserialize, Serialize}; +use serde_with::DeserializeFromStr; +use time::{OffsetDateTime, UtcOffset, format_description::FormatItem, macros::format_description}; + +pub const HISTORY_PAGE_SIZE: i64 = 100; +static EXAMPLE_CONFIG: &str = include_str!("../config.toml"); + +static DATA_DIR: OnceLock<PathBuf> = OnceLock::new(); +static META_CONFIG: OnceLock<(String, f64)> = OnceLock::new(); +static META_STORE: OnceCell<crate::meta::MetaStore> = OnceCell::const_new(); + +mod dotfiles; +mod kv; +pub(crate) mod meta; +mod scripts; +pub mod watcher; + +#[derive(Clone, Debug, Deserialize, Copy, ValueEnum, PartialEq, Serialize)] +pub enum SearchMode { + #[serde(rename = "prefix")] + Prefix, + + #[serde(rename = "fulltext")] + #[clap(aliases = &["fulltext"])] + FullText, + + #[serde(rename = "fuzzy")] + Fuzzy, + + #[serde(rename = "skim")] + Skim, + + #[serde(rename = "daemon-fuzzy")] + #[clap(aliases = &["daemon-fuzzy"])] + DaemonFuzzy, +} + +impl SearchMode { + pub fn as_str(&self) -> &'static str { + match self { + SearchMode::Prefix => "PREFIX", + SearchMode::FullText => "FULLTXT", + SearchMode::Fuzzy => "FUZZY", + SearchMode::Skim => "SKIM", + SearchMode::DaemonFuzzy => "DAEMON", + } + } + pub fn next(&self, settings: &Settings) -> Self { + match self { + SearchMode::Prefix => SearchMode::FullText, + // if the user is using skim, we go to skim + SearchMode::FullText if settings.search_mode == SearchMode::Skim => SearchMode::Skim, + // if the user is using daemon-fuzzy, we go to daemon-fuzzy + SearchMode::FullText if settings.search_mode == SearchMode::DaemonFuzzy => { + SearchMode::DaemonFuzzy + } + // otherwise fuzzy. + SearchMode::FullText => SearchMode::Fuzzy, + SearchMode::Fuzzy | SearchMode::Skim | SearchMode::DaemonFuzzy => SearchMode::Prefix, + } + } +} + +#[derive(Clone, Debug, Deserialize, Copy, PartialEq, Eq, ValueEnum, Serialize)] +pub enum FilterMode { + #[serde(rename = "global")] + Global = 0, + + #[serde(rename = "host")] + Host = 1, + + #[serde(rename = "session")] + Session = 2, + + #[serde(rename = "directory")] + Directory = 3, + + #[serde(rename = "workspace")] + Workspace = 4, + + #[serde(rename = "session-preload")] + SessionPreload = 5, +} + +impl FilterMode { + pub fn as_str(&self) -> &'static str { + match self { + FilterMode::Global => "GLOBAL", + FilterMode::Host => "HOST", + FilterMode::Session => "SESSION", + FilterMode::Directory => "DIRECTORY", + FilterMode::Workspace => "WORKSPACE", + FilterMode::SessionPreload => "SESSION+", + } + } +} + +#[derive(Clone, Debug, Deserialize, Copy, Serialize)] +pub enum ExitMode { + #[serde(rename = "return-original")] + ReturnOriginal, + + #[serde(rename = "return-query")] + ReturnQuery, +} + +// FIXME: Can use upstream Dialect enum if https://github.com/stevedonovan/chrono-english/pull/16 is merged +// FIXME: Above PR was merged, but dependency was changed to interim (fork of chrono-english) in the ... interim +#[derive(Clone, Debug, Deserialize, Copy, Serialize)] +pub enum Dialect { + #[serde(rename = "us")] + Us, + + #[serde(rename = "uk")] + Uk, +} + +impl From<Dialect> for interim::Dialect { + fn from(d: Dialect) -> interim::Dialect { + match d { + Dialect::Uk => interim::Dialect::Uk, + Dialect::Us => interim::Dialect::Us, + } + } +} + +/// Type wrapper around `time::UtcOffset` to support a wider variety of timezone formats. +/// +/// Note that the parsing of this struct needs to be done before starting any +/// multithreaded runtime, otherwise it will fail on most Unix systems. +/// +/// See: <https://github.com/atuinsh/atuin/pull/1517#discussion_r1447516426> +#[derive(Clone, Copy, Debug, Eq, PartialEq, DeserializeFromStr, Serialize)] +pub struct Timezone(pub UtcOffset); +impl fmt::Display for Timezone { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + self.0.fmt(f) + } +} +/// format: <+|-><hour>[:<minute>[:<second>]] +static OFFSET_FMT: &[FormatItem<'_>] = format_description!( + "[offset_hour sign:mandatory padding:none][optional [:[offset_minute padding:none][optional [:[offset_second padding:none]]]]]" +); +impl FromStr for Timezone { + type Err = Error; + + fn from_str(s: &str) -> Result<Self> { + // local timezone + if matches!(s.to_lowercase().as_str(), "l" | "local") { + // There have been some timezone issues, related to errors fetching it on some + // platforms + // Rather than fail to start, fallback to UTC. The user should still be able to specify + // their timezone manually in the config file. + let offset = UtcOffset::current_local_offset().unwrap_or(UtcOffset::UTC); + return Ok(Self(offset)); + } + + if matches!(s.to_lowercase().as_str(), "0" | "utc") { + let offset = UtcOffset::UTC; + return Ok(Self(offset)); + } + + // offset from UTC + if let Ok(offset) = UtcOffset::parse(s, OFFSET_FMT) { + return Ok(Self(offset)); + } + + // IDEA: Currently named timezones are not supported, because the well-known crate + // for this is `chrono_tz`, which is not really interoperable with the datetime crate + // that we currently use - `time`. If ever we migrate to using `chrono`, this would + // be a good feature to add. + + bail!(r#""{s}" is not a valid timezone spec"#) + } +} + +#[derive(Clone, Debug, Deserialize, Copy, Serialize)] +pub enum Style { + #[serde(rename = "auto")] + Auto, + + #[serde(rename = "full")] + Full, + + #[serde(rename = "compact")] + Compact, +} + +#[derive(Clone, Debug, Deserialize, Copy, Serialize)] +pub enum WordJumpMode { + #[serde(rename = "emacs")] + Emacs, + + #[serde(rename = "subl")] + Subl, +} + +#[derive(Clone, Debug, Deserialize, Copy, PartialEq, Eq, ValueEnum, Serialize)] +pub enum KeymapMode { + #[serde(rename = "emacs")] + Emacs, + + #[serde(rename = "vim-normal")] + VimNormal, + + #[serde(rename = "vim-insert")] + VimInsert, + + #[serde(rename = "auto")] + Auto, +} + +impl KeymapMode { + pub fn as_str(&self) -> &'static str { + match self { + KeymapMode::Emacs => "EMACS", + KeymapMode::VimNormal => "VIMNORMAL", + KeymapMode::VimInsert => "VIMINSERT", + KeymapMode::Auto => "AUTO", + } + } +} + +// We want to translate the config to crossterm::cursor::SetCursorStyle, but +// the original type does not implement trait serde::Deserialize unfortunately. +// It seems impossible to implement Deserialize for external types when it is +// used in HashMap (https://stackoverflow.com/questions/67142663). We instead +// define an adapter type. +#[derive(Clone, Debug, Deserialize, Copy, PartialEq, Eq, ValueEnum, Serialize)] +pub enum CursorStyle { + #[serde(rename = "default")] + DefaultUserShape, + + #[serde(rename = "blink-block")] + BlinkingBlock, + + #[serde(rename = "steady-block")] + SteadyBlock, + + #[serde(rename = "blink-underline")] + BlinkingUnderScore, + + #[serde(rename = "steady-underline")] + SteadyUnderScore, + + #[serde(rename = "blink-bar")] + BlinkingBar, + + #[serde(rename = "steady-bar")] + SteadyBar, +} + +impl CursorStyle { + pub fn as_str(&self) -> &'static str { + match self { + CursorStyle::DefaultUserShape => "DEFAULT", + CursorStyle::BlinkingBlock => "BLINKBLOCK", + CursorStyle::SteadyBlock => "STEADYBLOCK", + CursorStyle::BlinkingUnderScore => "BLINKUNDERLINE", + CursorStyle::SteadyUnderScore => "STEADYUNDERLINE", + CursorStyle::BlinkingBar => "BLINKBAR", + CursorStyle::SteadyBar => "STEADYBAR", + } + } +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct Stats { + #[serde(default = "Stats::common_prefix_default")] + pub common_prefix: Vec<String>, // sudo, etc. commands we want to strip off + #[serde(default = "Stats::common_subcommands_default")] + pub common_subcommands: Vec<String>, // kubectl, commands we should consider subcommands for + #[serde(default = "Stats::ignored_commands_default")] + pub ignored_commands: Vec<String>, // cd, ls, etc. commands we want to completely hide from stats +} + +impl Stats { + fn common_prefix_default() -> Vec<String> { + vec!["sudo", "doas"].into_iter().map(String::from).collect() + } + + fn common_subcommands_default() -> Vec<String> { + vec![ + "apt", + "cargo", + "composer", + "dnf", + "docker", + "dotnet", + "git", + "go", + "ip", + "jj", + "kubectl", + "nix", + "nmcli", + "npm", + "pecl", + "pnpm", + "podman", + "port", + "systemctl", + "tmux", + "yarn", + ] + .into_iter() + .map(String::from) + .collect() + } + + fn ignored_commands_default() -> Vec<String> { + vec![] + } +} + +impl Default for Stats { + fn default() -> Self { + Self { + common_prefix: Self::common_prefix_default(), + common_subcommands: Self::common_subcommands_default(), + ignored_commands: Self::ignored_commands_default(), + } + } +} + +#[derive(Clone, Debug, Deserialize, Default, Serialize)] +pub struct Sync { + pub records: bool, +} + +/// Sync protocol type for authentication. +/// +/// This setting is primarily for development/testing. When not explicitly set, +/// the protocol is inferred from the sync_address: +/// - Default sync address (api.atuin.sh) → Hub protocol +/// - Custom sync address → Legacy protocol +/// +/// Set explicitly to "hub" to use Hub authentication with a custom sync_address +/// (useful for local development against a Hub instance). +#[derive(Clone, Copy, Debug, Deserialize, PartialEq, Eq, Serialize, Default)] +#[serde(rename_all = "lowercase")] +pub enum SyncProtocol { + /// Use Hub authentication (Bearer token from Hub OAuth flow) + Hub, + /// Use legacy CLI authentication (Token from CLI register/login) + Legacy, + /// Infer from sync_address (default behavior) + #[default] + Auto, +} + +#[derive(Clone, Debug, Deserialize, Default, Serialize)] +pub struct Keys { + pub scroll_exits: bool, + pub exit_past_line_start: bool, + pub accept_past_line_end: bool, + pub accept_past_line_start: bool, + pub accept_with_backspace: bool, + pub prefix: String, +} + +impl Keys { + /// The standard default values for all `[keys]` options. + /// These match the config defaults set in `builder_with_data_dir()`. + pub fn standard_defaults() -> Self { + Keys { + scroll_exits: true, + exit_past_line_start: true, + accept_past_line_end: true, + accept_past_line_start: false, + accept_with_backspace: false, + prefix: "a".to_string(), + } + } + + /// Returns true if any value differs from the standard defaults. + pub fn has_non_default_values(&self) -> bool { + let d = Self::standard_defaults(); + self.scroll_exits != d.scroll_exits + || self.exit_past_line_start != d.exit_past_line_start + || self.accept_past_line_end != d.accept_past_line_end + || self.accept_past_line_start != d.accept_past_line_start + || self.accept_with_backspace != d.accept_with_backspace + || self.prefix != d.prefix + } +} + +/// A single rule within a conditional keybinding config. +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct KeyRuleConfig { + /// Optional condition expression (e.g. "cursor-at-start", "input-empty && no-results"). + /// If absent, the rule always matches. + #[serde(default)] + pub when: Option<String>, + /// The action to perform (e.g. "exit", "cursor-left", "accept"). + pub action: String, +} + +/// A keybinding config value: either a simple action string or an ordered list of conditional rules. +#[derive(Clone, Debug, Deserialize, Serialize)] +#[serde(untagged)] +pub enum KeyBindingConfig { + /// Simple unconditional binding: `"ctrl-c" = "return-original"` + Simple(String), + /// Conditional binding: `"left" = [{ when = "cursor-at-start", action = "exit" }, { action = "cursor-left" }]` + Rules(Vec<KeyRuleConfig>), +} + +/// User-facing keymap configuration. Each mode maps key strings to bindings. +/// Keys present here override the defaults for that key; unmentioned keys keep defaults. +#[derive(Clone, Debug, Deserialize, Serialize, Default)] +pub struct KeymapConfig { + #[serde(default)] + pub emacs: HashMap<String, KeyBindingConfig>, + #[serde(default, rename = "vim-normal")] + pub vim_normal: HashMap<String, KeyBindingConfig>, + #[serde(default, rename = "vim-insert")] + pub vim_insert: HashMap<String, KeyBindingConfig>, + #[serde(default)] + pub inspector: HashMap<String, KeyBindingConfig>, + #[serde(default)] + pub prefix: HashMap<String, KeyBindingConfig>, +} + +impl KeymapConfig { + /// Returns true if no keybinding overrides are configured in any mode. + pub fn is_empty(&self) -> bool { + self.emacs.is_empty() + && self.vim_normal.is_empty() + && self.vim_insert.is_empty() + && self.inspector.is_empty() + && self.prefix.is_empty() + } +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct Preview { + pub strategy: PreviewStrategy, +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct Theme { + /// Name of desired theme ("default" for base) + pub name: String, + + /// Whether any available additional theme debug should be shown + pub debug: Option<bool>, + + /// How many levels of parenthood will be traversed if needed + pub max_depth: Option<u8>, +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct Daemon { + /// Use the daemon to sync + /// If enabled, history hooks are routed through the daemon. + #[serde(alias = "enable")] + pub enabled: bool, + + /// Automatically start and manage a local daemon when needed. + pub autostart: bool, + + /// The daemon will handle sync on an interval. How often to sync, in seconds. + pub sync_frequency: u64, + + /// The path to the unix socket used by the daemon + pub socket_path: String, + + /// Path to the daemon pidfile used for process coordination. + pub pidfile_path: String, + + /// Use a socket passed via systemd's socket activation protocol, instead of the path + pub systemd_socket: bool, + + /// The port that should be used for TCP on non unix systems + pub tcp_port: u64, +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct Search { + /// The list of enabled filter modes, in order of priority. + pub filters: Vec<FilterMode>, + + /// The recency score multiplier for the search index (default: 1.0). + /// Values < 1.0 reduce weight, > 1.0 increase weight, 0.0 disables. + pub recency_score_multiplier: f64, + + /// The frequency score multiplier for the search index (default: 1.0). + /// Values < 1.0 reduce weight, > 1.0 increase weight, 0.0 disables. + pub frequency_score_multiplier: f64, + + /// The overall frecency score multiplier for the search index (default: 1.0). + /// Applied after combining recency and frequency scores. + pub frecency_score_multiplier: f64, +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct Tmux { + /// Enable using atuin with tmux popup (tmux >= 3.2) + pub enabled: bool, + + /// Width of the tmux popup (percentage) + pub width: String, + + /// Height of the tmux popup (percentage) + pub height: String, +} + +/// Log level for file logging. Maps to tracing's LevelFilter. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Deserialize, Serialize)] +#[serde(rename_all = "lowercase")] +pub enum LogLevel { + Trace, + Debug, + #[default] + Info, + Warn, + Error, +} + +impl LogLevel { + /// Convert to a tracing directive string for use with EnvFilter. + pub fn as_directive(&self) -> &'static str { + match self { + LogLevel::Trace => "trace", + LogLevel::Debug => "debug", + LogLevel::Info => "info", + LogLevel::Warn => "warn", + LogLevel::Error => "error", + } + } +} + +/// Configuration for a specific log type (search or daemon). +#[derive(Clone, Debug, Default, Deserialize, Serialize)] +pub struct LogConfig { + /// Log file name (relative to dir) or absolute path. + pub file: String, + + /// Override global enabled setting for this log type. + pub enabled: Option<bool>, + + /// Override global level setting for this log type. + pub level: Option<LogLevel>, + + /// Override global retention days setting for this log type. + pub retention: Option<u64>, +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct Logs { + /// Enable file logging globally. Defaults to true. + #[serde(default = "Logs::default_enabled")] + pub enabled: bool, + + /// Directory for log files. Defaults to ~/.atuin/logs + pub dir: String, + + /// Default log level for file logging. Defaults to "info". + /// Note: ATUIN_LOG environment variable overrides this. + #[serde(default)] + pub level: LogLevel, + + /// Default retention days for log files. Defaults to 4. + #[serde(default = "Logs::default_retention")] + pub retention: u64, + + /// Search log settings + #[serde(default)] + pub search: LogConfig, + + /// Daemon log settings + #[serde(default)] + pub daemon: LogConfig, + + /// AI log settings + #[serde(default)] + pub ai: LogConfig, +} + +#[derive(Default, Clone, Debug, Deserialize, Serialize)] +pub struct Ai { + /// Whether or not the AI features are enabled. + pub enabled: bool, + + /// The address of the Atuin AI endpoint. Used for AI features like command generation. + /// Only necessary for custom AI endpoints. + pub endpoint: Option<String>, + + /// The API token for the Atuin AI endpoint. Used for AI features like command generation. + /// Only necessary for custom AI endpoints. + pub api_token: Option<String>, + + /// Whether or not to send the current working directory to the AI endpoint. + pub send_cwd: bool, +} + +impl Default for Preview { + fn default() -> Self { + Self { + strategy: PreviewStrategy::Auto, + } + } +} + +impl Default for Theme { + fn default() -> Self { + Self { + name: "".to_string(), + debug: None::<bool>, + max_depth: Some(10), + } + } +} + +impl Default for Daemon { + fn default() -> Self { + Self { + enabled: false, + autostart: false, + sync_frequency: 300, + socket_path: "".to_string(), + pidfile_path: "".to_string(), + systemd_socket: false, + tcp_port: 8889, + } + } +} + +impl Default for Logs { + fn default() -> Self { + Self { + enabled: true, + dir: "".to_string(), + level: LogLevel::default(), + retention: Self::default_retention(), + search: LogConfig { + file: "search.log".to_string(), + ..Default::default() + }, + daemon: LogConfig { + file: "daemon.log".to_string(), + ..Default::default() + }, + ai: LogConfig { + file: "ai.log".to_string(), + ..Default::default() + }, + } + } +} + +impl Logs { + fn default_enabled() -> bool { + true + } + + fn default_retention() -> u64 { + 4 + } + + /// Returns whether search logging is enabled. + /// Uses search-specific setting if set, otherwise falls back to global. + pub fn search_enabled(&self) -> bool { + self.search.enabled.unwrap_or(self.enabled) + } + + /// Returns whether daemon logging is enabled. + /// Uses daemon-specific setting if set, otherwise falls back to global. + pub fn daemon_enabled(&self) -> bool { + self.daemon.enabled.unwrap_or(self.enabled) + } + + /// Returns whether AI logging is enabled. + /// Uses AI-specific setting if set, otherwise falls back to global. + pub fn ai_enabled(&self) -> bool { + self.ai.enabled.unwrap_or(self.enabled) + } + + /// Returns the log level for search logging. + /// Uses search-specific setting if set, otherwise falls back to global. + pub fn search_level(&self) -> LogLevel { + self.search.level.unwrap_or(self.level) + } + + /// Returns the log level for daemon logging. + /// Uses daemon-specific setting if set, otherwise falls back to global. + pub fn daemon_level(&self) -> LogLevel { + self.daemon.level.unwrap_or(self.level) + } + + /// Returns the log level for AI logging. + /// Uses AI-specific setting if set, otherwise falls back to global. + pub fn ai_level(&self) -> LogLevel { + self.ai.level.unwrap_or(self.level) + } + + /// Returns the retention days for search logging. + /// Uses search-specific setting if set, otherwise falls back to global. + pub fn search_retention(&self) -> u64 { + self.search.retention.unwrap_or(self.retention) + } + + /// Returns the retention days for daemon logging. + /// Uses daemon-specific setting if set, otherwise falls back to global. + pub fn daemon_retention(&self) -> u64 { + self.daemon.retention.unwrap_or(self.retention) + } + + /// Returns the retention days for AI logging. + /// Uses AI-specific setting if set, otherwise falls back to global. + pub fn ai_retention(&self) -> u64 { + self.ai.retention.unwrap_or(self.retention) + } + + /// Returns the full path for the search log file. + pub fn search_path(&self) -> PathBuf { + let path = PathBuf::from(&self.search.file); + PathBuf::from(&self.dir).join(path) + } + + /// Returns the full path for the daemon log file. + pub fn daemon_path(&self) -> PathBuf { + let path = PathBuf::from(&self.daemon.file); + PathBuf::from(&self.dir).join(path) + } + + /// Returns the full path for the AI log file. + pub fn ai_path(&self) -> PathBuf { + let path = PathBuf::from(&self.ai.file); + PathBuf::from(&self.dir).join(path) + } +} + +impl Default for Search { + fn default() -> Self { + Self { + filters: vec![ + FilterMode::Global, + FilterMode::Host, + FilterMode::Session, + FilterMode::SessionPreload, + FilterMode::Workspace, + FilterMode::Directory, + ], + + recency_score_multiplier: 1.0, + frequency_score_multiplier: 1.0, + frecency_score_multiplier: 1.0, + } + } +} + +impl Default for Tmux { + fn default() -> Self { + Self { + enabled: false, + width: "80%".to_string(), + height: "60%".to_string(), + } + } +} + +// The preview height strategy also takes max_preview_height into account. +#[derive(Clone, Debug, Deserialize, Copy, PartialEq, Eq, ValueEnum, Serialize)] +pub enum PreviewStrategy { + // Preview height is calculated for the length of the selected command. + #[serde(rename = "auto")] + Auto, + + // Preview height is calculated for the length of the longest command stored in the history. + #[serde(rename = "static")] + Static, + + // max_preview_height is used as fixed height. + #[serde(rename = "fixed")] + Fixed, +} + +/// Column types available for the interactive search UI. +#[derive(Clone, Copy, Debug, Deserialize, PartialEq, Eq, Serialize)] +#[serde(rename_all = "lowercase")] +pub enum UiColumnType { + /// Command execution duration (e.g., "123ms") + Duration, + /// Relative time since execution (e.g., "59s ago") + Time, + /// Absolute timestamp (e.g., "2025-01-22 14:35") + Datetime, + /// Working directory + Directory, + /// Hostname + Host, + /// Username + User, + /// Exit code + Exit, + /// The command itself (should be last, expands to fill) + Command, +} + +impl UiColumnType { + /// Returns the default width for this column type (in characters). + /// The Command column returns 0 as it expands to fill remaining space. + pub fn default_width(&self) -> u16 { + match self { + UiColumnType::Duration => 5, // "814ms" + UiColumnType::Time => 9, // "459ms ago" + UiColumnType::Datetime => 16, // "2025-01-22 14:35" + UiColumnType::Directory => 20, + UiColumnType::Host => 15, + UiColumnType::User => 10, + UiColumnType::Exit => { + if cfg!(windows) { + 11 // 32-bit integer on Windows: "-1978335212" + } else { + 3 // Usually a byte on Unix + } + } + UiColumnType::Command => 0, // Expands to fill + } + } +} + +/// A column configuration with type and optional custom width. +/// Can be specified as just a string (uses default width) or as an object with type and width. +#[derive(Clone, Debug, Serialize)] +pub struct UiColumn { + pub column_type: UiColumnType, + pub width: u16, + /// If true, this column expands to fill remaining space. Only one column should expand. + pub expand: bool, +} + +impl UiColumn { + pub fn new(column_type: UiColumnType) -> Self { + Self { + width: column_type.default_width(), + expand: column_type == UiColumnType::Command, + column_type, + } + } + + pub fn with_width(column_type: UiColumnType, width: u16) -> Self { + Self { + column_type, + width, + expand: column_type == UiColumnType::Command, + } + } +} + +// Custom deserialize to handle both string and object formats: +// "duration" or { type = "duration", width = 8, expand = true } +impl<'de> serde::Deserialize<'de> for UiColumn { + fn deserialize<D>(deserializer: D) -> Result<Self, D::Error> + where + D: serde::Deserializer<'de>, + { + use serde::de::{self, MapAccess, Visitor}; + + struct UiColumnVisitor; + + impl<'de> Visitor<'de> for UiColumnVisitor { + type Value = UiColumn; + + fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { + formatter.write_str( + "a column type string or an object with 'type' and optional 'width'/'expand'", + ) + } + + fn visit_str<E>(self, value: &str) -> Result<UiColumn, E> + where + E: de::Error, + { + let column_type: UiColumnType = + serde::Deserialize::deserialize(serde::de::value::StrDeserializer::new(value))?; + Ok(UiColumn::new(column_type)) + } + + fn visit_map<M>(self, mut map: M) -> Result<UiColumn, M::Error> + where + M: MapAccess<'de>, + { + let mut column_type: Option<UiColumnType> = None; + let mut width: Option<u16> = None; + let mut expand: Option<bool> = None; + + while let Some(key) = map.next_key::<String>()? { + match key.as_str() { + "type" => { + column_type = Some(map.next_value()?); + } + "width" => { + width = Some(map.next_value()?); + } + "expand" => { + expand = Some(map.next_value()?); + } + _ => { + let _: serde::de::IgnoredAny = map.next_value()?; + } + } + } + + let column_type = column_type.ok_or_else(|| de::Error::missing_field("type"))?; + let width = width.unwrap_or_else(|| column_type.default_width()); + let expand = expand.unwrap_or(column_type == UiColumnType::Command); + Ok(UiColumn { + column_type, + width, + expand, + }) + } + } + + deserializer.deserialize_any(UiColumnVisitor) + } +} + +/// UI-specific settings for the interactive search. +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct Ui { + /// Columns to display in interactive search, from left to right. + /// The indicator column (" > ") is always shown first implicitly. + /// The "command" column should be last as it expands to fill remaining space. + /// Can be simple strings or objects with type and width. + #[serde(default = "Ui::default_columns")] + pub columns: Vec<UiColumn>, +} + +impl Ui { + fn default_columns() -> Vec<UiColumn> { + vec![ + UiColumn::new(UiColumnType::Duration), + UiColumn::new(UiColumnType::Time), + UiColumn::new(UiColumnType::Command), + ] + } + + /// Validate the UI configuration. + /// Returns an error if more than one column has expand = true. + pub fn validate(&self) -> Result<()> { + let expand_count = self.columns.iter().filter(|c| c.expand).count(); + if expand_count > 1 { + bail!( + "Only one column can have expand = true, but {} columns are set to expand", + expand_count + ); + } + Ok(()) + } +} + +impl Default for Ui { + fn default() -> Self { + Self { + columns: Self::default_columns(), + } + } +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct Settings { + pub data_dir: Option<String>, + pub dialect: Dialect, + pub timezone: Timezone, + pub style: Style, + pub auto_sync: bool, + pub update_check: bool, + + /// The sync address for atuin. + pub sync_address: String, + + /// Sync protocol for authentication. When set to "auto" (default), the protocol + /// is inferred from sync_address. Set to "hub" to force Hub auth with a custom + /// sync_address (useful for local development). + #[serde(default)] + pub sync_protocol: SyncProtocol, + + pub sync_frequency: String, + pub db_path: String, + pub record_store_path: String, + pub key_path: String, + pub search_mode: SearchMode, + pub filter_mode: Option<FilterMode>, + pub filter_mode_shell_up_key_binding: Option<FilterMode>, + pub search_mode_shell_up_key_binding: Option<SearchMode>, + pub shell_up_key_binding: bool, + pub inline_height: u16, + pub inline_height_shell_up_key_binding: Option<u16>, + pub invert: bool, + pub show_preview: bool, + pub max_preview_height: u16, + pub show_help: bool, + pub show_tabs: bool, + pub show_numeric_shortcuts: bool, + pub auto_hide_height: u16, + pub exit_mode: ExitMode, + pub keymap_mode: KeymapMode, + pub keymap_mode_shell: KeymapMode, + pub keymap_cursor: HashMap<String, CursorStyle>, + pub word_jump_mode: WordJumpMode, + pub word_chars: String, + pub scroll_context_lines: usize, + pub history_format: String, + pub prefers_reduced_motion: bool, + pub store_failed: bool, + + #[serde(with = "serde_regex", default = "RegexSet::empty", skip_serializing)] + pub history_filter: RegexSet, + + #[serde(with = "serde_regex", default = "RegexSet::empty", skip_serializing)] + pub cwd_filter: RegexSet, + + pub secrets_filter: bool, + pub workspaces: bool, + pub ctrl_n_shortcuts: bool, + + pub network_connect_timeout: u64, + pub network_timeout: u64, + pub local_timeout: f64, + pub enter_accept: bool, + pub smart_sort: bool, + pub command_chaining: bool, + + #[serde(default)] + pub stats: Stats, + + #[serde(default)] + pub sync: Sync, + + #[serde(default)] + pub keys: Keys, + + #[serde(default)] + pub keymap: KeymapConfig, + + #[serde(default)] + pub preview: Preview, + + #[serde(default)] + pub dotfiles: dotfiles::Settings, + + #[serde(default)] + pub daemon: Daemon, + + #[serde(default)] + pub search: Search, + + #[serde(default)] + pub theme: Theme, + + #[serde(default)] + pub ui: Ui, + + #[serde(default)] + pub scripts: scripts::Settings, + + #[serde(default)] + pub kv: kv::Settings, + + #[serde(default)] + pub tmux: Tmux, + + #[serde(default)] + pub logs: Logs, + + #[serde(default)] + pub meta: meta::Settings, + + #[serde(default)] + pub ai: Ai, +} + +impl Settings { + pub fn utc() -> Self { + Self::builder() + .expect("Could not build default") + .set_override("timezone", "0") + .expect("failed to override timezone with UTC") + .build() + .expect("Could not build config") + .try_deserialize() + .expect("Could not deserialize config") + } + + pub(crate) fn effective_data_dir() -> PathBuf { + DATA_DIR + .get() + .cloned() + .unwrap_or_else(atuin_common::utils::data_dir) + } + + // -- Meta store: lazily initialized on first access -- + + pub async fn meta_store() -> Result<&'static crate::meta::MetaStore> { + META_STORE + .get_or_try_init(|| async { + let (db_path, timeout) = META_CONFIG.get().ok_or_else(|| { + eyre!("meta store config not set — Settings::new() has not been called") + })?; + crate::meta::MetaStore::new(db_path, *timeout).await + }) + .await + } + + pub async fn host_id() -> Result<HostId> { + Self::meta_store().await?.host_id().await + } + + pub async fn last_sync() -> Result<OffsetDateTime> { + Self::meta_store().await?.last_sync().await + } + + pub async fn save_sync_time() -> Result<()> { + Self::meta_store().await?.save_sync_time().await + } + + pub async fn last_version_check() -> Result<OffsetDateTime> { + Self::meta_store().await?.last_version_check().await + } + + pub async fn save_version_check_time() -> Result<()> { + Self::meta_store().await?.save_version_check_time().await + } + + pub async fn should_sync(&self) -> Result<bool> { + if !self.auto_sync || !Self::meta_store().await?.logged_in().await? { + return Ok(false); + } + + if self.sync_frequency == "0" { + return Ok(true); + } + + match parse_duration(self.sync_frequency.as_str()) { + Ok(d) => { + let d = time::Duration::try_from(d)?; + Ok(OffsetDateTime::now_utc() - Settings::last_sync().await? >= d) + } + Err(e) => Err(eyre!("failed to check sync: {}", e)), + } + } + + pub async fn logged_in(&self) -> Result<bool> { + Self::meta_store().await?.logged_in().await + } + + pub async fn session_token(&self) -> Result<String> { + match Self::meta_store().await?.session_token().await? { + Some(token) => Ok(token), + None => Err(eyre!("Tried to load session; not logged in")), + } + } + + pub async fn hub_session_token(&self) -> Result<String> { + match Self::meta_store().await?.hub_session_token().await? { + Some(token) => Ok(token), + None => Err(eyre!("Tried to load hub session; not logged in")), + } + } + + /// Default sync address for Atuin's hosted service + pub const DEFAULT_SYNC_ADDRESS: &'static str = "https://api.atuin.sh"; + + /// Default Hub web/API endpoint for Atuin's hosted service + pub const DEFAULT_HUB_ENDPOINT: &'static str = "https://hub.atuin.sh"; + + /// Normalize a URL for comparison by trimming trailing slashes + fn normalize_url(url: &str) -> &str { + url.trim_end_matches('/') + } + + /// Check if a URL matches one of Atuin's official hosted addresses + fn is_official_address(url: &str) -> bool { + let normalized = Self::normalize_url(url); + normalized == Self::normalize_url(Self::DEFAULT_SYNC_ADDRESS) + || normalized == Self::normalize_url(Self::DEFAULT_HUB_ENDPOINT) + } + + /// Returns whether this configuration uses Hub-style sync. + /// + /// Hub sync uses Bearer token authentication and is the default for + /// Atuin's hosted service. This returns true when: + /// - `sync_protocol` is explicitly set to `Hub`, OR + /// - `sync_protocol` is `Auto` and `sync_address` is an official Atuin address + pub fn is_hub_sync(&self) -> bool { + match self.sync_protocol { + SyncProtocol::Hub => true, + SyncProtocol::Legacy => false, + SyncProtocol::Auto => Self::is_official_address(&self.sync_address), + } + } + + /// Returns the base URL for the Hub endpoint. + /// + /// For Atuin's official hosted service, this always returns `https://hub.atuin.sh` + /// regardless of whether `sync_address` is `api.atuin.sh` or `hub.atuin.sh`. + /// For self-hosted instances, returns the configured `sync_address`. + pub fn active_hub_endpoint(&self) -> Option<String> { + if self.is_hub_sync() { + if Self::is_official_address(&self.sync_address) { + Some(Self::DEFAULT_HUB_ENDPOINT.to_string()) + } else { + Some(self.sync_address.clone()) + } + } else { + None + } + } + + /// Returns the best available auth token for sync operations. + /// + /// Token priority when using Hub sync: + /// 1. Hub token (Bearer) - enables unified Hub auth + /// 2. CLI session token (Token) - fallback if Hub token revoked + /// + /// For legacy/self-hosted sync, only CLI session token is used. + /// + /// Hub tokens are preferred when available because they provide unified + /// authentication across CLI and Hub features, and users can manage them + /// via the Hub web interface. + #[cfg(feature = "sync")] + pub async fn sync_auth_token(&self) -> Result<crate::api_client::AuthToken> { + use crate::api_client::AuthToken; + + let meta = Self::meta_store().await?; + + // Try Hub token first if we're using Hub sync + if self.is_hub_sync() + && let Some(hub_token) = meta.hub_session_token().await? + { + return Ok(AuthToken::Bearer(hub_token)); + } + + // Fall back to CLI session token + match meta.session_token().await? { + Some(token) => Ok(AuthToken::Token(token)), + None => Err(eyre!( + "Not logged in - no Hub session or CLI session found. \ + Run 'atuin login' or 'atuin register' to authenticate." + )), + } + } + + #[cfg(feature = "check-update")] + async fn needs_update_check(&self) -> Result<bool> { + let last_check = Settings::last_version_check().await?; + let diff = OffsetDateTime::now_utc() - last_check; + + // Check a max of once per hour + Ok(diff.whole_hours() >= 1) + } + + #[cfg(feature = "check-update")] + async fn latest_version(&self) -> Result<Version> { + // Default to the current version, and if that doesn't parse, a version so high it's unlikely to ever + // suggest upgrading. + let current = + Version::parse(env!("CARGO_PKG_VERSION")).unwrap_or(Version::new(100000, 0, 0)); + + if !self.needs_update_check().await? { + let meta = Self::meta_store().await?; + let version = match meta.latest_version().await? { + Some(v) => Version::parse(&v).unwrap_or(current), + None => current, + }; + + return Ok(version); + } + + #[cfg(feature = "sync")] + let latest = crate::api_client::latest_version().await.unwrap_or(current); + + #[cfg(not(feature = "sync"))] + let latest = current; + + let meta = Self::meta_store().await?; + Settings::save_version_check_time().await?; + meta.save_latest_version(&latest.to_string()).await?; + + Ok(latest) + } + + // Return Some(latest version) if an update is needed. Otherwise, none. + #[cfg(feature = "check-update")] + pub async fn needs_update(&self) -> Option<Version> { + if !self.update_check { + return None; + } + + let current = + Version::parse(env!("CARGO_PKG_VERSION")).unwrap_or(Version::new(100000, 0, 0)); + + let latest = self.latest_version().await; + + if latest.is_err() { + return None; + } + + let latest = latest.unwrap(); + + if latest > current { + return Some(latest); + } + + None + } + + pub fn default_filter_mode(&self, git_root: bool) -> FilterMode { + self.filter_mode + .filter(|x| self.search.filters.contains(x)) + .or_else(|| { + self.search + .filters + .iter() + .find(|x| match (x, git_root, self.workspaces) { + (FilterMode::Workspace, true, true) => true, + (FilterMode::Workspace, _, _) => false, + (_, _, _) => true, + }) + .copied() + }) + .unwrap_or(FilterMode::Global) + } + + #[cfg(not(feature = "check-update"))] + pub async fn needs_update(&self) -> Option<Version> { + None + } + + pub fn builder() -> Result<ConfigBuilder<DefaultState>> { + Self::builder_with_data_dir(&atuin_common::utils::data_dir()) + } + + fn builder_with_data_dir(data_dir: &std::path::Path) -> Result<ConfigBuilder<DefaultState>> { + let db_path = data_dir.join("history.db"); + let record_store_path = data_dir.join("records.db"); + let kv_path = data_dir.join("kv.db"); + let scripts_path = data_dir.join("scripts.db"); + let socket_path = atuin_common::utils::runtime_dir().join("atuin.sock"); + let pidfile_path = data_dir.join("atuin-daemon.pid"); + let logs_dir = atuin_common::utils::logs_dir(); + + let key_path = data_dir.join("key"); + let meta_path = data_dir.join("meta.db"); + + Ok(Config::builder() + .set_default("history_format", "{time}\t{command}\t{duration}")? + .set_default("db_path", db_path.to_str())? + .set_default("record_store_path", record_store_path.to_str())? + .set_default("key_path", key_path.to_str())? + .set_default("dialect", "us")? + .set_default("timezone", "local")? + .set_default("auto_sync", true)? + .set_default("update_check", cfg!(feature = "check-update"))? + .set_default("sync_address", "https://api.atuin.sh")? + .set_default("sync_frequency", "5m")? + .set_default("search_mode", "fuzzy")? + .set_default("filter_mode", None::<String>)? + .set_default("style", "compact")? + .set_default("inline_height", 40)? + .set_default("show_preview", true)? + .set_default("preview.strategy", "auto")? + .set_default("max_preview_height", 4)? + .set_default("show_help", true)? + .set_default("show_tabs", true)? + .set_default("show_numeric_shortcuts", true)? + .set_default("auto_hide_height", 8)? + .set_default("invert", false)? + .set_default("exit_mode", "return-original")? + .set_default("word_jump_mode", "emacs")? + .set_default( + "word_chars", + "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789", + )? + .set_default("scroll_context_lines", 1)? + .set_default("shell_up_key_binding", false)? + .set_default("workspaces", false)? + .set_default("ctrl_n_shortcuts", false)? + .set_default("secrets_filter", true)? + .set_default("network_connect_timeout", 5)? + .set_default("network_timeout", 30)? + .set_default("local_timeout", 2.0)? + // enter_accept defaults to false here, but true in the default config file. The dissonance is + // intentional! + // Existing users will get the default "False", so we don't mess with any potential + // muscle memory. + // New users will get the new default, that is more similar to what they are used to. + .set_default("enter_accept", false)? + .set_default("sync.records", true)? + .set_default("keys.scroll_exits", true)? + .set_default("keys.accept_past_line_end", true)? + .set_default("keys.exit_past_line_start", true)? + .set_default("keys.accept_past_line_start", false)? + .set_default("keys.accept_with_backspace", false)? + .set_default("keys.prefix", "a")? + .set_default("keymap_mode", "emacs")? + .set_default("keymap_mode_shell", "auto")? + .set_default("keymap_cursor", HashMap::<String, String>::new())? + .set_default("smart_sort", false)? + .set_default("command_chaining", false)? + .set_default("store_failed", true)? + .set_default("daemon.sync_frequency", 300)? + .set_default("daemon.enabled", false)? + .set_default("daemon.autostart", false)? + .set_default("daemon.socket_path", socket_path.to_str())? + .set_default("daemon.pidfile_path", pidfile_path.to_str())? + .set_default("daemon.systemd_socket", false)? + .set_default("daemon.tcp_port", 8889)? + .set_default("logs.enabled", true)? + .set_default("logs.dir", logs_dir.to_str())? + .set_default("logs.level", "info")? + .set_default("logs.search.file", "search.log")? + .set_default("logs.daemon.file", "daemon.log")? + .set_default("logs.ai.file", "ai.log")? + .set_default("kv.db_path", kv_path.to_str())? + .set_default("scripts.db_path", scripts_path.to_str())? + .set_default("search.recency_score_multiplier", 1.0)? + .set_default("search.frequency_score_multiplier", 1.0)? + .set_default("search.frecency_score_multiplier", 1.0)? + .set_default("meta.db_path", meta_path.to_str())? + .set_default("ai.enabled", false)? + .set_default("ai.send_cwd", false)? + .set_default( + "search.filters", + vec![ + "global", + "host", + "session", + "workspace", + "directory", + "session-preload", + ], + )? + .set_default("theme.name", "default")? + .set_default("theme.debug", None::<bool>)? + .set_default("tmux.enabled", false)? + .set_default("tmux.width", "80%")? + .set_default("tmux.height", "60%")? + .set_default( + "prefers_reduced_motion", + std::env::var("NO_MOTION") + .ok() + .map(|_| config::Value::new(None, config::ValueKind::Boolean(true))) + .unwrap_or_else(|| config::Value::new(None, config::ValueKind::Boolean(false))), + )? + .add_source( + Environment::with_prefix("atuin") + .prefix_separator("_") + .separator("__"), + )) + } + + pub fn get_config_path() -> Result<PathBuf> { + let config_dir = atuin_common::utils::config_dir(); + + create_dir_all(&config_dir) + .wrap_err_with(|| format!("could not create dir {config_dir:?}"))?; + + let mut config_file = if let Ok(p) = std::env::var("ATUIN_CONFIG_DIR") { + PathBuf::from(p) + } else { + let mut config_file = PathBuf::new(); + config_file.push(config_dir); + config_file + }; + + config_file.push("config.toml"); + + Ok(config_file) + } + + pub fn new() -> Result<Self> { + let config_file = Self::get_config_path()?; + + // extract data_dir first so we can use it as the base for other path defaults + let effective_data_dir = if config_file.exists() { + #[derive(Deserialize, Default)] + struct DataDirOnly { + data_dir: Option<String>, + } + + let config_file_str = config_file + .to_str() + .ok_or_else(|| eyre!("config file path is not valid UTF-8"))?; + + let partial_config = Config::builder() + .add_source(ConfigFile::new(config_file_str, FileFormat::Toml)) + .add_source( + Environment::with_prefix("atuin") + .prefix_separator("_") + .separator("__"), + ) + .build() + .ok(); + + let custom_data_dir = partial_config + .and_then(|c| c.try_deserialize::<DataDirOnly>().ok()) + .and_then(|d| d.data_dir); + + match custom_data_dir { + Some(dir) => { + let expanded = shellexpand::full(&dir) + .map_err(|e| eyre!("failed to expand data_dir path: {}", e))?; + PathBuf::from(expanded.as_ref()) + } + None => atuin_common::utils::data_dir(), + } + } else { + atuin_common::utils::data_dir() + }; + + DATA_DIR.set(effective_data_dir.clone()).ok(); + + create_dir_all(&effective_data_dir) + .wrap_err_with(|| format!("could not create dir {effective_data_dir:?}"))?; + + let mut config_builder = Self::builder_with_data_dir(&effective_data_dir)?; + + config_builder = if config_file.exists() { + let config_file_str = config_file + .to_str() + .ok_or_else(|| eyre!("config file path is not valid UTF-8"))?; + config_builder.add_source(ConfigFile::new(config_file_str, FileFormat::Toml)) + } else { + let mut file = File::create(config_file).wrap_err("could not create config file")?; + file.write_all(EXAMPLE_CONFIG.as_bytes()) + .wrap_err("could not write default config file")?; + + config_builder + }; + + let config = config_builder.build()?; + let mut settings: Settings = config + .try_deserialize() + .map_err(|e| eyre!("failed to deserialize: {}", e))?; + + // all paths should be expanded + settings.db_path = Self::expand_path(settings.db_path)?; + settings.record_store_path = Self::expand_path(settings.record_store_path)?; + settings.key_path = Self::expand_path(settings.key_path)?; + settings.daemon.socket_path = Self::expand_path(settings.daemon.socket_path)?; + settings.daemon.pidfile_path = Self::expand_path(settings.daemon.pidfile_path)?; + settings.logs.dir = Self::expand_path(settings.logs.dir)?; + settings.logs.search.file = Self::expand_path(settings.logs.search.file)?; + settings.logs.daemon.file = Self::expand_path(settings.logs.daemon.file)?; + + // Validate UI settings + settings.ui.validate()?; + + // Register meta store config for lazy initialization on first access + META_CONFIG + .set((settings.meta.db_path.clone(), settings.local_timeout)) + .ok(); + + Ok(settings) + } + + fn expand_path(path: String) -> Result<String> { + shellexpand::full(&path) + .map(|p| p.to_string()) + .map_err(|e| eyre!("failed to expand path: {}", e)) + } + + pub fn example_config() -> &'static str { + EXAMPLE_CONFIG + } + + pub fn paths_ok(&self) -> bool { + let paths = [ + &self.db_path, + &self.record_store_path, + &self.key_path, + &self.meta.db_path, + ]; + paths.iter().all(|p| !utils::broken_symlink(p)) + } +} + +impl Default for Settings { + fn default() -> Self { + // if this panics something is very wrong, as the default config + // does not build or deserialize into the settings struct + Self::builder() + .expect("Could not build default") + .build() + .expect("Could not build config") + .try_deserialize() + .expect("Could not deserialize config") + } +} + +/// Initialize the meta store configuration for testing. +/// +/// This should only be used in tests. It allows tests to bypass the normal +/// Settings::new() flow while still being able to use Settings::host_id() +/// and other meta store dependent functions. +/// +/// # Safety +/// This function is not thread-safe with concurrent calls to Settings::new() +/// or other meta store initialization. Only call from tests. +#[doc(hidden)] +pub fn init_meta_config_for_testing(meta_db_path: impl Into<String>, local_timeout: f64) { + META_CONFIG.set((meta_db_path.into(), local_timeout)).ok(); +} + +#[cfg(test)] +pub(crate) fn test_local_timeout() -> f64 { + std::env::var("ATUIN_TEST_LOCAL_TIMEOUT") + .ok() + .and_then(|x| x.parse().ok()) + // this hardcoded value should be replaced by a simple way to get the + // default local_timeout of Settings if possible + .unwrap_or(2.0) +} + +#[cfg(test)] +mod tests { + use std::str::FromStr; + + use eyre::Result; + + use super::Timezone; + + #[test] + fn can_parse_offset_timezone_spec() -> Result<()> { + assert_eq!(Timezone::from_str("+02")?.0.as_hms(), (2, 0, 0)); + assert_eq!(Timezone::from_str("-04")?.0.as_hms(), (-4, 0, 0)); + assert_eq!(Timezone::from_str("+05:30")?.0.as_hms(), (5, 30, 0)); + assert_eq!(Timezone::from_str("-09:30")?.0.as_hms(), (-9, -30, 0)); + + // single digit hours are allowed + assert_eq!(Timezone::from_str("+2")?.0.as_hms(), (2, 0, 0)); + assert_eq!(Timezone::from_str("-4")?.0.as_hms(), (-4, 0, 0)); + assert_eq!(Timezone::from_str("+5:30")?.0.as_hms(), (5, 30, 0)); + assert_eq!(Timezone::from_str("-9:30")?.0.as_hms(), (-9, -30, 0)); + + // fully qualified form + assert_eq!(Timezone::from_str("+09:30:00")?.0.as_hms(), (9, 30, 0)); + assert_eq!(Timezone::from_str("-09:30:00")?.0.as_hms(), (-9, -30, 0)); + + // these offsets don't really exist but are supported anyway + assert_eq!(Timezone::from_str("+0:5")?.0.as_hms(), (0, 5, 0)); + assert_eq!(Timezone::from_str("-0:5")?.0.as_hms(), (0, -5, 0)); + assert_eq!(Timezone::from_str("+01:23:45")?.0.as_hms(), (1, 23, 45)); + assert_eq!(Timezone::from_str("-01:23:45")?.0.as_hms(), (-1, -23, -45)); + + // require a leading sign for clarity + assert!(Timezone::from_str("5").is_err()); + assert!(Timezone::from_str("10:30").is_err()); + + Ok(()) + } + + #[test] + fn can_choose_workspace_filters_when_in_git_context() -> Result<()> { + let mut settings = super::Settings::default(); + settings.search.filters = vec![ + super::FilterMode::Workspace, + super::FilterMode::Host, + super::FilterMode::Directory, + super::FilterMode::Session, + super::FilterMode::Global, + ]; + settings.workspaces = true; + + assert_eq!( + settings.default_filter_mode(true), + super::FilterMode::Workspace, + ); + + Ok(()) + } + + #[test] + fn wont_choose_workspace_filters_when_not_in_git_context() -> Result<()> { + let mut settings = super::Settings::default(); + settings.search.filters = vec![ + super::FilterMode::Workspace, + super::FilterMode::Host, + super::FilterMode::Directory, + super::FilterMode::Session, + super::FilterMode::Global, + ]; + settings.workspaces = true; + + assert_eq!(settings.default_filter_mode(false), super::FilterMode::Host,); + + Ok(()) + } + + #[test] + fn wont_choose_workspace_filters_when_workspaces_disabled() -> Result<()> { + let mut settings = super::Settings::default(); + settings.search.filters = vec![ + super::FilterMode::Workspace, + super::FilterMode::Host, + super::FilterMode::Directory, + super::FilterMode::Session, + super::FilterMode::Global, + ]; + settings.workspaces = false; + + assert_eq!(settings.default_filter_mode(true), super::FilterMode::Host,); + + Ok(()) + } + + #[test] + fn builder_with_data_dir_uses_custom_paths() -> Result<()> { + use std::path::PathBuf; + + let custom_dir = PathBuf::from("/custom/data/dir"); + let builder = super::Settings::builder_with_data_dir(&custom_dir)?; + let config = builder.build()?; + + let db_path: String = config.get("db_path")?; + let key_path: String = config.get("key_path")?; + let record_store_path: String = config.get("record_store_path")?; + let kv_db_path: String = config.get("kv.db_path")?; + let scripts_db_path: String = config.get("scripts.db_path")?; + let meta_db_path: String = config.get("meta.db_path")?; + let daemon_socket_path: String = config.get("daemon.socket_path")?; + let daemon_pidfile_path: String = config.get("daemon.pidfile_path")?; + let daemon_autostart: bool = config.get("daemon.autostart")?; + + assert_eq!(db_path, custom_dir.join("history.db").to_str().unwrap()); + assert_eq!(key_path, custom_dir.join("key").to_str().unwrap()); + assert_eq!( + record_store_path, + custom_dir.join("records.db").to_str().unwrap() + ); + assert_eq!(kv_db_path, custom_dir.join("kv.db").to_str().unwrap()); + assert_eq!( + scripts_db_path, + custom_dir.join("scripts.db").to_str().unwrap() + ); + assert_eq!(meta_db_path, custom_dir.join("meta.db").to_str().unwrap()); + assert_eq!( + daemon_socket_path, + atuin_common::utils::runtime_dir() + .join("atuin.sock") + .to_str() + .unwrap() + ); + assert_eq!( + daemon_pidfile_path, + custom_dir.join("atuin-daemon.pid").to_str().unwrap() + ); + assert!(!daemon_autostart); + + Ok(()) + } + + #[test] + fn effective_data_dir_returns_default_when_not_set() { + let effective = super::Settings::effective_data_dir(); + let default = atuin_common::utils::data_dir(); + + assert!(effective.to_str().is_some()); + assert!(effective.ends_with("atuin") || effective == default); + } + + #[test] + fn keymap_config_deserializes_simple_binding() { + let json = r#"{"emacs": {"ctrl-c": "exit"}}"#; + let config: super::KeymapConfig = serde_json::from_str(json).unwrap(); + assert_eq!(config.emacs.len(), 1); + match &config.emacs["ctrl-c"] { + super::KeyBindingConfig::Simple(s) => assert_eq!(s, "exit"), + _ => panic!("expected Simple variant"), + } + } + + #[test] + fn keymap_config_deserializes_conditional_binding() { + let json = r#"{ + "emacs": { + "left": [ + {"when": "cursor-at-start", "action": "exit"}, + {"action": "cursor-left"} + ] + } + }"#; + let config: super::KeymapConfig = serde_json::from_str(json).unwrap(); + match &config.emacs["left"] { + super::KeyBindingConfig::Rules(rules) => { + assert_eq!(rules.len(), 2); + assert_eq!(rules[0].when.as_deref(), Some("cursor-at-start")); + assert_eq!(rules[0].action, "exit"); + assert!(rules[1].when.is_none()); + assert_eq!(rules[1].action, "cursor-left"); + } + _ => panic!("expected Rules variant"), + } + } + + #[test] + fn keymap_config_deserializes_vim_normal() { + let json = r#"{"vim-normal": {"j": "select-next", "k": "select-previous"}}"#; + let config: super::KeymapConfig = serde_json::from_str(json).unwrap(); + assert_eq!(config.vim_normal.len(), 2); + assert!(config.emacs.is_empty()); + } + + #[test] + fn keymap_config_is_empty_when_default() { + let config = super::KeymapConfig::default(); + assert!(config.is_empty()); + } + + #[test] + fn keymap_config_mixed_modes() { + let json = r#"{ + "emacs": {"ctrl-c": "exit"}, + "vim-normal": {"q": "exit"}, + "inspector": {"d": "delete"} + }"#; + let config: super::KeymapConfig = serde_json::from_str(json).unwrap(); + assert!(!config.is_empty()); + assert_eq!(config.emacs.len(), 1); + assert_eq!(config.vim_normal.len(), 1); + assert_eq!(config.inspector.len(), 1); + assert!(config.vim_insert.is_empty()); + assert!(config.prefix.is_empty()); + } +} diff --git a/crates/atuin-client/src/settings/dotfiles.rs b/crates/atuin-client/src/settings/dotfiles.rs new file mode 100644 index 00000000..bbaf914f --- /dev/null +++ b/crates/atuin-client/src/settings/dotfiles.rs @@ -0,0 +1,7 @@ +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Serialize, Deserialize, Clone, Default)] +pub struct Settings { + #[serde(alias = "enable")] + pub enabled: bool, +} diff --git a/crates/atuin-client/src/settings/kv.rs b/crates/atuin-client/src/settings/kv.rs new file mode 100644 index 00000000..afc24a35 --- /dev/null +++ b/crates/atuin-client/src/settings/kv.rs @@ -0,0 +1,17 @@ +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Serialize, Deserialize, Clone)] +pub struct Settings { + pub db_path: String, +} + +impl Default for Settings { + fn default() -> Self { + let dir = atuin_common::utils::data_dir(); + let path = dir.join("kv.db"); + + Self { + db_path: path.to_string_lossy().to_string(), + } + } +} diff --git a/crates/atuin-client/src/settings/meta.rs b/crates/atuin-client/src/settings/meta.rs new file mode 100644 index 00000000..108d74ec --- /dev/null +++ b/crates/atuin-client/src/settings/meta.rs @@ -0,0 +1,17 @@ +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Serialize, Deserialize, Clone)] +pub struct Settings { + pub db_path: String, +} + +impl Default for Settings { + fn default() -> Self { + let dir = atuin_common::utils::data_dir(); + let path = dir.join("meta.db"); + + Self { + db_path: path.to_string_lossy().to_string(), + } + } +} diff --git a/crates/atuin-client/src/settings/scripts.rs b/crates/atuin-client/src/settings/scripts.rs new file mode 100644 index 00000000..a1a5ed6a --- /dev/null +++ b/crates/atuin-client/src/settings/scripts.rs @@ -0,0 +1,17 @@ +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Serialize, Deserialize, Clone)] +pub struct Settings { + pub db_path: String, +} + +impl Default for Settings { + fn default() -> Self { + let dir = atuin_common::utils::data_dir(); + let path = dir.join("scripts.db"); + + Self { + db_path: path.to_string_lossy().to_string(), + } + } +} diff --git a/crates/atuin-client/src/settings/watcher.rs b/crates/atuin-client/src/settings/watcher.rs new file mode 100644 index 00000000..740b8d12 --- /dev/null +++ b/crates/atuin-client/src/settings/watcher.rs @@ -0,0 +1,256 @@ +//! Config file watching for automatic settings reload. +//! +//! This module provides a `SettingsWatcher` that monitors the config file +//! for changes and broadcasts updated settings via a `tokio::sync::watch` channel. +//! +//! # Example +//! +//! ```no_run +//! use atuin_client::settings::watcher::global_settings_watcher; +//! +//! async fn example() -> eyre::Result<()> { +//! let watcher = global_settings_watcher()?; +//! let mut rx = watcher.subscribe(); +//! +//! // React to settings changes +//! while rx.changed().await.is_ok() { +//! let settings = rx.borrow(); +//! println!("Settings updated!"); +//! } +//! Ok(()) +//! } +//! ``` + +use std::{ + path::PathBuf, + sync::{Arc, OnceLock}, + time::Duration, +}; + +use eyre::{Result, WrapErr}; +use log::{debug, error, info, warn}; +use notify::{ + Config as NotifyConfig, RecommendedWatcher, RecursiveMode, Watcher, + event::{EventKind, ModifyKind}, +}; +use tokio::sync::watch; + +use super::Settings; + +/// Global singleton for the settings watcher. +static SETTINGS_WATCHER: OnceLock<Result<SettingsWatcher, String>> = OnceLock::new(); + +/// Get the global settings watcher singleton. +/// +/// Initializes the watcher on first call. Subsequent calls return the same instance. +/// The watcher monitors the config file for changes and broadcasts updates. +pub fn global_settings_watcher() -> Result<&'static SettingsWatcher> { + let result = SETTINGS_WATCHER.get_or_init(|| SettingsWatcher::new().map_err(|e| e.to_string())); + + match result { + Ok(watcher) => Ok(watcher), + Err(e) => Err(eyre::eyre!("{}", e)), + } +} + +/// Watches the config file for changes and broadcasts updated settings. +/// +/// Uses `notify` for cross-platform file watching and `tokio::sync::watch` +/// for efficient broadcast to multiple subscribers. +pub struct SettingsWatcher { + /// Receiver for settings updates. Clone this to subscribe. + rx: watch::Receiver<Arc<Settings>>, + /// Keeps the file watcher alive for the lifetime of this struct. + _watcher: RecommendedWatcher, +} + +impl SettingsWatcher { + /// Create a new settings watcher. + /// + /// Loads initial settings and starts watching the config file for changes. + /// Changes are debounced (500ms) to avoid multiple reloads during saves. + pub fn new() -> Result<Self> { + let initial_settings = Arc::new(Settings::new()?); + let (tx, rx) = watch::channel(initial_settings); + + let config_path = Self::config_path(); + info!("starting config file watcher: {:?}", config_path); + + let watcher = Self::create_watcher(tx, config_path)?; + + Ok(Self { + rx, + _watcher: watcher, + }) + } + + /// Subscribe to settings updates. + /// + /// Returns a receiver that will be notified when settings change. + /// Use `changed().await` to wait for the next update, then `borrow()` + /// to access the current settings. + pub fn subscribe(&self) -> watch::Receiver<Arc<Settings>> { + self.rx.clone() + } + + /// Get the current settings without subscribing to updates. + pub fn current(&self) -> Arc<Settings> { + self.rx.borrow().clone() + } + + /// Get the config file path. + fn config_path() -> PathBuf { + let config_dir = if let Ok(p) = std::env::var("ATUIN_CONFIG_DIR") { + PathBuf::from(p) + } else { + atuin_common::utils::config_dir() + }; + config_dir.join("config.toml") + } + + /// Create the file watcher with debouncing. + fn create_watcher( + tx: watch::Sender<Arc<Settings>>, + config_path: PathBuf, + ) -> Result<RecommendedWatcher> { + // Channel for debouncing file events + let (debounce_tx, debounce_rx) = std::sync::mpsc::channel::<()>(); + + // Spawn debounce thread + let config_path_clone = config_path.clone(); + std::thread::spawn(move || { + Self::debounce_loop(debounce_rx, tx, config_path_clone); + }); + + // Clone config_path for use in the watcher callback + let config_path_for_watcher = config_path.clone(); + + // Canonicalize config path for reliable comparison on macOS + // (handles symlinks like /var -> /private/var) + let canonical_config_path = config_path_for_watcher + .canonicalize() + .unwrap_or_else(|_| config_path_for_watcher.clone()); + + // Create file watcher + let mut watcher = RecommendedWatcher::new( + move |res: Result<notify::Event, notify::Error>| { + match res { + Ok(event) => { + // Defensive: if paths is empty, we can't filter, so assume + // it might be our config file and trigger a reload to be safe + if event.paths.is_empty() { + warn!( + "config watcher: event has no paths, triggering reload to be safe" + ); + let _ = debounce_tx.send(()); + return; + } + + // Only react to events for our specific config file + // (filter out editor temp files, backups, etc.) + let is_config_file = event.paths.iter().any(|path| { + // Canonicalize for reliable comparison (handles macOS symlinks) + let canonical_event_path = + path.canonicalize().unwrap_or_else(|_| path.clone()); + + // Check if this event is for our config file + // (either exact match or the file was renamed to our config) + canonical_event_path == canonical_config_path + || path.file_name() == config_path_for_watcher.file_name() + }); + + if !is_config_file { + return; + } + + // Only react to modify events (content changes) or creates + if matches!( + event.kind, + EventKind::Modify(ModifyKind::Data(_) | ModifyKind::Any) + | EventKind::Create(_) + ) { + debug!("config file event detected: {:?}", event); + // Send to debounce channel (ignore send errors - receiver might be gone) + let _ = debounce_tx.send(()); + } + } + Err(e) => { + error!("file watcher error: {}", e); + } + } + }, + NotifyConfig::default(), + ) + .wrap_err("failed to create file watcher")?; + + // Watch the config file's parent directory (some editors create new files) + let watch_path = config_path.parent().unwrap_or(&config_path); + + // Defensive: ensure watch path exists before trying to watch + if !watch_path.exists() { + warn!( + "config directory does not exist, creating it: {:?}", + watch_path + ); + std::fs::create_dir_all(watch_path) + .wrap_err_with(|| format!("failed to create config directory: {:?}", watch_path))?; + } + + watcher + .watch(watch_path, RecursiveMode::NonRecursive) + .wrap_err_with(|| format!("failed to watch config directory: {:?}", watch_path))?; + + info!("config file watcher initialized for: {:?}", watch_path); + Ok(watcher) + } + + /// Debounce loop that batches file events and reloads settings. + fn debounce_loop( + rx: std::sync::mpsc::Receiver<()>, + tx: watch::Sender<Arc<Settings>>, + config_path: PathBuf, + ) { + const DEBOUNCE_DURATION: Duration = Duration::from_millis(500); + + loop { + // Wait for first event + if rx.recv().is_err() { + // Channel closed, watcher was dropped + debug!("config watcher debounce loop exiting"); + return; + } + + // Drain any additional events within debounce window + while rx.recv_timeout(DEBOUNCE_DURATION).is_ok() { + // Keep draining + } + + // Defensive: check if config file exists before reloading + // (handles case where file was deleted - we'll get notified when it's recreated) + if !config_path.exists() { + debug!( + "config file does not exist, skipping reload: {:?}", + config_path + ); + continue; + } + + // Now reload settings + info!("config file changed, reloading settings: {:?}", config_path); + match Settings::new() { + Ok(settings) => { + if tx.send(Arc::new(settings)).is_err() { + // All receivers dropped + debug!("all settings subscribers dropped, exiting"); + return; + } + info!("settings reloaded successfully"); + } + Err(e) => { + warn!("failed to reload settings: {}", e); + // Keep the old settings, don't broadcast the error + } + } + } + } +} diff --git a/crates/atuin-client/src/sync.rs b/crates/atuin-client/src/sync.rs new file mode 100644 index 00000000..2c902794 --- /dev/null +++ b/crates/atuin-client/src/sync.rs @@ -0,0 +1,213 @@ +use std::collections::HashSet; +use std::iter::FromIterator; + +use eyre::Result; + +use atuin_common::api::AddHistoryRequest; +use crypto_secretbox::Key; +use time::OffsetDateTime; + +use crate::{ + api_client, + database::Database, + encryption::{decrypt, encrypt, load_key}, + settings::Settings, +}; + +pub fn hash_str(string: &str) -> String { + use sha2::{Digest, Sha256}; + let mut hasher = Sha256::new(); + hasher.update(string.as_bytes()); + hex::encode(hasher.finalize()) +} + +// Currently sync is kinda naive, and basically just pages backwards through +// history. This means newly added stuff shows up properly! We also just use +// the total count in each database to indicate whether a sync is needed. +// I think this could be massively improved! If we had a way of easily +// indicating count per time period (hour, day, week, year, etc) then we can +// easily pinpoint where we are missing data and what needs downloading. Start +// with year, then find the week, then the day, then the hour, then download it +// all! The current naive approach will do for now. + +// Check if remote has things we don't, and if so, download them. +// Returns (num downloaded, total local) +async fn sync_download( + key: &Key, + force: bool, + client: &api_client::Client<'_>, + db: &impl Database, +) -> Result<(i64, i64)> { + debug!("starting sync download"); + + let remote_status = client.status().await?; + let remote_count = remote_status.count; + + // useful to ensure we don't even save something that hasn't yet been synced + deleted + let remote_deleted = + HashSet::<&str>::from_iter(remote_status.deleted.iter().map(String::as_str)); + + let initial_local = db.history_count(true).await?; + let mut local_count = initial_local; + + let mut last_sync = if force { + OffsetDateTime::UNIX_EPOCH + } else { + Settings::last_sync().await? + }; + + let mut last_timestamp = OffsetDateTime::UNIX_EPOCH; + + let host = if force { Some(String::from("")) } else { None }; + + while remote_count > local_count { + let page = client + .get_history(last_sync, last_timestamp, host.clone()) + .await?; + + let history: Vec<_> = page + .history + .iter() + // TODO: handle deletion earlier in this chain + .map(|h| serde_json::from_str(h).expect("invalid base64")) + .map(|h| decrypt(h, key).expect("failed to decrypt history! check your key")) + .map(|mut h| { + if remote_deleted.contains(h.id.0.as_str()) { + h.deleted_at = Some(time::OffsetDateTime::now_utc()); + h.command = String::from(""); + } + + h + }) + .collect(); + + db.save_bulk(&history).await?; + + local_count = db.history_count(true).await?; + let remote_page_size = std::cmp::max(remote_status.page_size, 0) as usize; + + if history.len() < remote_page_size { + break; + } + + let page_last = history + .last() + .expect("could not get last element of page") + .timestamp; + + // in the case of a small sync frequency, it's possible for history to + // be "lost" between syncs. In this case we need to rewind the sync + // timestamps + if page_last == last_timestamp { + last_timestamp = OffsetDateTime::UNIX_EPOCH; + last_sync -= time::Duration::hours(1); + } else { + last_timestamp = page_last; + } + } + + for i in remote_status.deleted { + // we will update the stored history to have this data + // pretty much everything can be nullified + match db.load(i.as_str()).await? { + Some(h) => { + db.delete(h).await?; + } + _ => { + info!( + "could not delete history with id {}, not found locally", + i.as_str() + ); + } + } + } + + Ok((local_count - initial_local, local_count)) +} + +// Check if we have things remote doesn't, and if so, upload them +async fn sync_upload( + key: &Key, + _force: bool, + client: &api_client::Client<'_>, + db: &impl Database, +) -> Result<()> { + debug!("starting sync upload"); + + let remote_status = client.status().await?; + let remote_deleted: HashSet<String> = HashSet::from_iter(remote_status.deleted.clone()); + + let initial_remote_count = client.count().await?; + let mut remote_count = initial_remote_count; + + let local_count = db.history_count(true).await?; + + debug!("remote has {remote_count}, we have {local_count}"); + + // first just try the most recent set + let mut cursor = OffsetDateTime::now_utc(); + + while local_count > remote_count { + let last = db.before(cursor, remote_status.page_size).await?; + let mut buffer = Vec::new(); + + if last.is_empty() { + break; + } + + for i in last { + let data = encrypt(&i, key)?; + let data = serde_json::to_string(&data)?; + + let add_hist = AddHistoryRequest { + id: i.id.to_string(), + timestamp: i.timestamp, + data, + hostname: hash_str(&i.hostname), + }; + + buffer.push(add_hist); + } + + // anything left over outside of the 100 block size + client.post_history(&buffer).await?; + cursor = buffer.last().unwrap().timestamp; + remote_count = client.count().await?; + + debug!("upload cursor: {cursor:?}"); + } + + let deleted = db.deleted().await?; + + for i in deleted { + if remote_deleted.contains(&i.id.to_string()) { + continue; + } + + info!("deleting {} on remote", i.id); + client.delete_history(i).await?; + } + + Ok(()) +} + +pub async fn sync(settings: &Settings, force: bool, db: &impl Database) -> Result<()> { + let client = api_client::Client::new( + &settings.sync_address, + settings.sync_auth_token().await?, + settings.network_connect_timeout, + settings.network_timeout, + )?; + + Settings::save_sync_time().await?; + + let key = load_key(settings)?; // encryption key + + sync_upload(&key, force, &client, db).await?; + + let download = sync_download(&key, force, &client, db).await?; + + debug!("sync downloaded {}", download.0); + + Ok(()) +} diff --git a/crates/atuin-client/src/theme.rs b/crates/atuin-client/src/theme.rs new file mode 100644 index 00000000..a277ac13 --- /dev/null +++ b/crates/atuin-client/src/theme.rs @@ -0,0 +1,831 @@ +use config::{Config, File as ConfigFile, FileFormat}; +use log; +use palette::named; +use serde::{Deserialize, Serialize}; +use serde_json; +use std::collections::HashMap; +use std::error; +use std::io::{Error, ErrorKind}; +use std::path::PathBuf; +use std::sync::LazyLock; +use strum_macros; + +static DEFAULT_MAX_DEPTH: u8 = 10; + +// Collection of settable "meanings" that can have colors set. +// NOTE: You can add a new meaning here without breaking backwards compatibility but please: +// - update the atuin/docs repository, which has a list of available meanings +// - add a fallback in the MEANING_FALLBACKS below, so that themes which do not have it +// get a sensible fallback (see Title as an example) +#[derive( + Serialize, Deserialize, Copy, Clone, Hash, Debug, Eq, PartialEq, strum_macros::Display, +)] +#[strum(serialize_all = "camel_case")] +pub enum Meaning { + AlertInfo, + AlertWarn, + AlertError, + Annotation, + Base, + Guidance, + Important, + Title, + Muted, +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct ThemeConfig { + // Definition of the theme + pub theme: ThemeDefinitionConfigBlock, + + // Colors + pub colors: HashMap<Meaning, String>, +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct ThemeDefinitionConfigBlock { + /// Name of theme ("default" for base) + pub name: String, + + /// Whether any theme should be treated as a parent _if available_ + pub parent: Option<String>, +} + +use crossterm::style::{Attribute, Attributes, Color, ContentStyle}; + +// For now, a theme is loaded as a mapping of meanings to colors, but it may be desirable to +// expand that in the future to general styles, so we populate a Meaning->ContentStyle hashmap. +pub struct Theme { + pub name: String, + pub parent: Option<String>, + pub styles: HashMap<Meaning, ContentStyle>, +} + +// Themes have a number of convenience functions for the most commonly used meanings. +// The general purpose `as_style` routine gives back a style, but for ease-of-use and to keep +// theme-related boilerplate minimal, the convenience functions give a color. +impl Theme { + // This is the base "default" color, for general text + pub fn get_base(&self) -> ContentStyle { + self.styles[&Meaning::Base] + } + + pub fn get_info(&self) -> ContentStyle { + self.get_alert(log::Level::Info) + } + + pub fn get_warning(&self) -> ContentStyle { + self.get_alert(log::Level::Warn) + } + + pub fn get_error(&self) -> ContentStyle { + self.get_alert(log::Level::Error) + } + + // The alert meanings may be chosen by the Level enum, rather than the methods above + // or the full Meaning enum, to simplify programmatic selection of a log-level. + pub fn get_alert(&self, severity: log::Level) -> ContentStyle { + self.styles[ALERT_TYPES.get(&severity).unwrap()] + } + + pub fn new( + name: String, + parent: Option<String>, + styles: HashMap<Meaning, ContentStyle>, + ) -> Theme { + Theme { + name, + parent, + styles, + } + } + + pub fn closest_meaning<'a>(&self, meaning: &'a Meaning) -> &'a Meaning { + if self.styles.contains_key(meaning) { + meaning + } else if MEANING_FALLBACKS.contains_key(meaning) { + self.closest_meaning(&MEANING_FALLBACKS[meaning]) + } else { + &Meaning::Base + } + } + + // General access - if you have a meaning, this will give you a (crossterm) style + pub fn as_style(&self, meaning: Meaning) -> ContentStyle { + self.styles[self.closest_meaning(&meaning)] + } + + // Turns a map of meanings to colornames into a theme + // If theme-debug is on, then we will print any colornames that we cannot load, + // but we do not have this on in general, as it could print unfiltered text to the terminal + // from a theme TOML file. However, it will always return a theme, falling back to + // defaults on error, so that a TOML file does not break loading + pub fn from_foreground_colors( + name: String, + parent: Option<&Theme>, + foreground_colors: HashMap<Meaning, String>, + debug: bool, + ) -> Theme { + let styles: HashMap<Meaning, ContentStyle> = foreground_colors + .iter() + .map(|(name, color)| { + ( + *name, + StyleFactory::from_fg_string(color).unwrap_or_else(|err| { + if debug { + log::warn!("Tried to load string as a color unsuccessfully: ({name}={color}) {err}"); + } + ContentStyle::default() + }), + ) + }) + .collect(); + Theme::from_map(name, parent, &styles) + } + + // Boil down a meaning-color hashmap into a theme, by taking the defaults + // for any unknown colors + fn from_map( + name: String, + parent: Option<&Theme>, + overrides: &HashMap<Meaning, ContentStyle>, + ) -> Theme { + let styles = match parent { + Some(theme) => Box::new(theme.styles.clone()), + None => Box::new(DEFAULT_THEME.styles.clone()), + } + .iter() + .map(|(name, color)| match overrides.get(name) { + Some(value) => (*name, *value), + None => (*name, *color), + }) + .collect(); + Theme::new(name, parent.map(|p| p.name.clone()), styles) + } +} + +// Use palette to get a color from a string name, if possible +fn from_string(name: &str) -> Result<Color, String> { + if name.is_empty() { + return Err("Empty string".into()); + } + let first_char = name.chars().next().unwrap(); + match first_char { + '#' => { + let hexcode = &name[1..]; + let vec: Vec<u8> = hexcode + .chars() + .collect::<Vec<char>>() + .chunks(2) + .map(|pair| u8::from_str_radix(pair.iter().collect::<String>().as_str(), 16)) + .filter_map(|n| n.ok()) + .collect(); + if vec.len() != 3 { + return Err("Could not parse 3 hex values from string".into()); + } + Ok(Color::Rgb { + r: vec[0], + g: vec[1], + b: vec[2], + }) + } + '@' => { + // For full flexibility, we need to use serde_json, given + // crossterm's approach. + serde_json::from_str::<Color>(format!("\"{}\"", &name[1..]).as_str()) + .map_err(|_| format!("Could not convert color name {name} to Crossterm color")) + } + _ => { + let srgb = named::from_str(name).ok_or("No such color in palette")?; + Ok(Color::Rgb { + r: srgb.red, + g: srgb.green, + b: srgb.blue, + }) + } + } +} + +pub struct StyleFactory {} + +impl StyleFactory { + fn from_fg_string(name: &str) -> Result<ContentStyle, String> { + match from_string(name) { + Ok(color) => Ok(Self::from_fg_color(color)), + Err(err) => Err(err), + } + } + + // For succinctness, if we are confident that the name will be known, + // this routine is available to keep the code readable + fn known_fg_string(name: &str) -> ContentStyle { + Self::from_fg_string(name).unwrap() + } + + fn from_fg_color(color: Color) -> ContentStyle { + ContentStyle { + foreground_color: Some(color), + ..ContentStyle::default() + } + } + + fn from_fg_color_and_attributes(color: Color, attributes: Attributes) -> ContentStyle { + ContentStyle { + foreground_color: Some(color), + attributes, + ..ContentStyle::default() + } + } +} + +// Built-in themes. Rather than having extra files added before any theming +// is available, this gives a couple of basic options, demonstrating the use +// of themes: autumn and marine +static ALERT_TYPES: LazyLock<HashMap<log::Level, Meaning>> = LazyLock::new(|| { + HashMap::from([ + (log::Level::Info, Meaning::AlertInfo), + (log::Level::Warn, Meaning::AlertWarn), + (log::Level::Error, Meaning::AlertError), + ]) +}); + +static MEANING_FALLBACKS: LazyLock<HashMap<Meaning, Meaning>> = LazyLock::new(|| { + HashMap::from([ + (Meaning::Guidance, Meaning::AlertInfo), + (Meaning::Annotation, Meaning::AlertInfo), + (Meaning::Title, Meaning::Important), + ]) +}); + +static DEFAULT_THEME: LazyLock<Theme> = LazyLock::new(|| { + Theme::new( + "default".to_string(), + None, + HashMap::from([ + ( + Meaning::AlertError, + StyleFactory::from_fg_color(Color::DarkRed), + ), + ( + Meaning::AlertWarn, + StyleFactory::from_fg_color(Color::DarkYellow), + ), + ( + Meaning::AlertInfo, + StyleFactory::from_fg_color(Color::DarkGreen), + ), + ( + Meaning::Annotation, + StyleFactory::from_fg_color(Color::DarkGrey), + ), + ( + Meaning::Guidance, + StyleFactory::from_fg_color(Color::DarkBlue), + ), + ( + Meaning::Important, + StyleFactory::from_fg_color_and_attributes( + Color::White, + Attributes::from(Attribute::Bold), + ), + ), + (Meaning::Muted, StyleFactory::from_fg_color(Color::Grey)), + (Meaning::Base, ContentStyle::default()), + ]), + ) +}); + +static BUILTIN_THEMES: LazyLock<HashMap<&'static str, Theme>> = LazyLock::new(|| { + HashMap::from([ + ("default", HashMap::new()), + ( + "(none)", + HashMap::from([ + (Meaning::AlertError, ContentStyle::default()), + (Meaning::AlertWarn, ContentStyle::default()), + (Meaning::AlertInfo, ContentStyle::default()), + (Meaning::Annotation, ContentStyle::default()), + (Meaning::Guidance, ContentStyle::default()), + (Meaning::Important, ContentStyle::default()), + (Meaning::Muted, ContentStyle::default()), + (Meaning::Base, ContentStyle::default()), + ]), + ), + ( + "autumn", + HashMap::from([ + ( + Meaning::AlertError, + StyleFactory::known_fg_string("saddlebrown"), + ), + ( + Meaning::AlertWarn, + StyleFactory::known_fg_string("darkorange"), + ), + (Meaning::AlertInfo, StyleFactory::known_fg_string("gold")), + ( + Meaning::Annotation, + StyleFactory::from_fg_color(Color::DarkGrey), + ), + (Meaning::Guidance, StyleFactory::known_fg_string("brown")), + ]), + ), + ( + "marine", + HashMap::from([ + ( + Meaning::AlertError, + StyleFactory::known_fg_string("yellowgreen"), + ), + (Meaning::AlertWarn, StyleFactory::known_fg_string("cyan")), + ( + Meaning::AlertInfo, + StyleFactory::known_fg_string("turquoise"), + ), + ( + Meaning::Annotation, + StyleFactory::known_fg_string("steelblue"), + ), + ( + Meaning::Base, + StyleFactory::known_fg_string("lightsteelblue"), + ), + (Meaning::Guidance, StyleFactory::known_fg_string("teal")), + ]), + ), + ]) + .iter() + .map(|(name, theme)| (*name, Theme::from_map(name.to_string(), None, theme))) + .collect() +}); + +// To avoid themes being repeatedly loaded, we store them in a theme manager +pub struct ThemeManager { + loaded_themes: HashMap<String, Theme>, + debug: bool, + override_theme_dir: Option<String>, +} + +// Theme-loading logic +impl ThemeManager { + pub fn new(debug: Option<bool>, theme_dir: Option<String>) -> Self { + Self { + loaded_themes: HashMap::new(), + debug: debug.unwrap_or(false), + override_theme_dir: match theme_dir { + Some(theme_dir) => Some(theme_dir), + None => std::env::var("ATUIN_THEME_DIR").ok(), + }, + } + } + + // Try to load a theme from a `{name}.toml` file in the theme directory. If an override is set + // for the theme dir (via ATUIN_THEME_DIR env) we should load the theme from there + pub fn load_theme_from_file( + &mut self, + name: &str, + max_depth: u8, + ) -> Result<&Theme, Box<dyn error::Error>> { + let mut theme_file = if let Some(p) = &self.override_theme_dir { + if p.is_empty() { + return Err(Box::new(Error::new( + ErrorKind::NotFound, + "Empty theme directory override and could not find theme elsewhere", + ))); + } + PathBuf::from(p) + } else { + let config_dir = atuin_common::utils::config_dir(); + let mut theme_file = if let Ok(p) = std::env::var("ATUIN_CONFIG_DIR") { + PathBuf::from(p) + } else { + let mut theme_file = PathBuf::new(); + theme_file.push(config_dir); + theme_file + }; + theme_file.push("themes"); + theme_file + }; + + let theme_toml = format!["{name}.toml"]; + theme_file.push(theme_toml); + + let mut config_builder = Config::builder(); + + config_builder = config_builder.add_source(ConfigFile::new( + theme_file.to_str().unwrap(), + FileFormat::Toml, + )); + + let config = config_builder.build()?; + self.load_theme_from_config(name, config, max_depth) + } + + pub fn load_theme_from_config( + &mut self, + name: &str, + config: Config, + max_depth: u8, + ) -> Result<&Theme, Box<dyn error::Error>> { + let debug = self.debug; + let theme_config: ThemeConfig = match config.try_deserialize() { + Ok(tc) => tc, + Err(e) => { + return Err(Box::new(Error::new( + ErrorKind::InvalidInput, + format!( + "Failed to deserialize theme: {}", + if debug { + e.to_string() + } else { + "set theme debug on for more info".to_string() + } + ), + ))); + } + }; + let colors: HashMap<Meaning, String> = theme_config.colors; + let parent: Option<&Theme> = match theme_config.theme.parent { + Some(parent_name) => { + if max_depth == 0 { + return Err(Box::new(Error::new( + ErrorKind::InvalidInput, + "Parent requested but we hit the recursion limit", + ))); + } + Some(self.load_theme(parent_name.as_str(), Some(max_depth - 1))) + } + None => Some(self.load_theme("default", Some(max_depth - 1))), + }; + + if debug && name != theme_config.theme.name { + log::warn!( + "Your theme config name is not the name of your loaded theme {} != {}", + name, + theme_config.theme.name + ); + } + + let theme = Theme::from_foreground_colors(theme_config.theme.name, parent, colors, debug); + let name = name.to_string(); + self.loaded_themes.insert(name.clone(), theme); + let theme = self.loaded_themes.get(&name).unwrap(); + Ok(theme) + } + + // Check if the requested theme is loaded and, if not, then attempt to get it + // from the builtins or, if not there, from file + pub fn load_theme(&mut self, name: &str, max_depth: Option<u8>) -> &Theme { + if self.loaded_themes.contains_key(name) { + return self.loaded_themes.get(name).unwrap(); + } + let built_ins = &BUILTIN_THEMES; + match built_ins.get(name) { + Some(theme) => theme, + None => match self.load_theme_from_file(name, max_depth.unwrap_or(DEFAULT_MAX_DEPTH)) { + Ok(theme) => theme, + Err(err) => { + log::warn!("Could not load theme {name}: {err}"); + built_ins.get("(none)").unwrap() + } + }, + } + } +} + +#[cfg(test)] +mod theme_tests { + use super::*; + + #[test] + fn test_can_load_builtin_theme() { + let mut manager = ThemeManager::new(Some(false), Some("".to_string())); + let theme = manager.load_theme("autumn", None); + assert_eq!( + theme.as_style(Meaning::Guidance).foreground_color, + from_string("brown").ok() + ); + } + + #[test] + fn test_can_create_theme() { + let mut manager = ThemeManager::new(Some(false), Some("".to_string())); + let mytheme = Theme::new( + "mytheme".to_string(), + None, + HashMap::from([( + Meaning::AlertError, + StyleFactory::known_fg_string("yellowgreen"), + )]), + ); + manager.loaded_themes.insert("mytheme".to_string(), mytheme); + let theme = manager.load_theme("mytheme", None); + assert_eq!( + theme.as_style(Meaning::AlertError).foreground_color, + from_string("yellowgreen").ok() + ); + } + + #[test] + fn test_can_fallback_when_meaning_missing() { + let mut manager = ThemeManager::new(Some(false), Some("".to_string())); + + // We use title as an example of a meaning that is not defined + // even in the base theme. + assert!(!DEFAULT_THEME.styles.contains_key(&Meaning::Title)); + + let config = Config::builder() + .add_source(ConfigFile::from_str( + " + [theme] + name = \"title_theme\" + + [colors] + Guidance = \"white\" + AlertInfo = \"zomp\" + ", + FileFormat::Toml, + )) + .build() + .unwrap(); + let theme = manager + .load_theme_from_config("config_theme", config, 1) + .unwrap(); + + // Correctly picks overridden color. + assert_eq!( + theme.as_style(Meaning::Guidance).foreground_color, + from_string("white").ok() + ); + + // Does not fall back to any color. + assert_eq!(theme.as_style(Meaning::AlertInfo).foreground_color, None); + + // Even for the base. + assert_eq!(theme.as_style(Meaning::Base).foreground_color, None); + + // Falls back to red as meaning missing from theme, so picks base default. + assert_eq!( + theme.as_style(Meaning::AlertError).foreground_color, + Some(Color::DarkRed) + ); + + // Falls back to Important as Title not available. + assert_eq!( + theme.as_style(Meaning::Title).foreground_color, + theme.as_style(Meaning::Important).foreground_color, + ); + + let title_config = Config::builder() + .add_source(ConfigFile::from_str( + " + [theme] + name = \"title_theme\" + + [colors] + Title = \"white\" + AlertInfo = \"zomp\" + ", + FileFormat::Toml, + )) + .build() + .unwrap(); + let title_theme = manager + .load_theme_from_config("title_theme", title_config, 1) + .unwrap(); + + assert_eq!( + title_theme.as_style(Meaning::Title).foreground_color, + Some(Color::White) + ); + } + + #[test] + fn test_no_fallbacks_are_circular() { + let mytheme = Theme::new("mytheme".to_string(), None, HashMap::from([])); + MEANING_FALLBACKS + .iter() + .for_each(|pair| assert_eq!(mytheme.closest_meaning(pair.0), &Meaning::Base)) + } + + #[test] + fn test_can_get_colors_via_convenience_functions() { + let mut manager = ThemeManager::new(Some(true), Some("".to_string())); + let theme = manager.load_theme("default", None); + assert_eq!(theme.get_error().foreground_color.unwrap(), Color::DarkRed); + assert_eq!( + theme.get_warning().foreground_color.unwrap(), + Color::DarkYellow + ); + assert_eq!(theme.get_info().foreground_color.unwrap(), Color::DarkGreen); + assert_eq!(theme.get_base().foreground_color, None); + assert_eq!( + theme.get_alert(log::Level::Error).foreground_color.unwrap(), + Color::DarkRed + ) + } + + #[test] + fn test_can_use_parent_theme_for_fallbacks() { + testing_logger::setup(); + + let mut manager = ThemeManager::new(Some(false), Some("".to_string())); + + // First, we introduce a base theme + let solarized = Config::builder() + .add_source(ConfigFile::from_str( + " + [theme] + name = \"solarized\" + + [colors] + Guidance = \"white\" + AlertInfo = \"pink\" + ", + FileFormat::Toml, + )) + .build() + .unwrap(); + let solarized_theme = manager + .load_theme_from_config("solarized", solarized, 1) + .unwrap(); + + assert_eq!( + solarized_theme + .as_style(Meaning::AlertInfo) + .foreground_color, + from_string("pink").ok() + ); + + // Then we introduce a derived theme + let unsolarized = Config::builder() + .add_source(ConfigFile::from_str( + " + [theme] + name = \"unsolarized\" + parent = \"solarized\" + + [colors] + AlertInfo = \"red\" + ", + FileFormat::Toml, + )) + .build() + .unwrap(); + let unsolarized_theme = manager + .load_theme_from_config("unsolarized", unsolarized, 1) + .unwrap(); + + // It will take its own values + assert_eq!( + unsolarized_theme + .as_style(Meaning::AlertInfo) + .foreground_color, + from_string("red").ok() + ); + + // ...or fall back to the parent + assert_eq!( + unsolarized_theme + .as_style(Meaning::Guidance) + .foreground_color, + from_string("white").ok() + ); + + testing_logger::validate(|captured_logs| assert_eq!(captured_logs.len(), 0)); + + // If the parent is not found, we end up with the no theme colors or styling + // as this is considered a (soft) error state. + let nunsolarized = Config::builder() + .add_source(ConfigFile::from_str( + " + [theme] + name = \"nunsolarized\" + parent = \"nonsolarized\" + + [colors] + AlertInfo = \"red\" + ", + FileFormat::Toml, + )) + .build() + .unwrap(); + let nunsolarized_theme = manager + .load_theme_from_config("nunsolarized", nunsolarized, 1) + .unwrap(); + + assert_eq!( + nunsolarized_theme + .as_style(Meaning::Guidance) + .foreground_color, + None + ); + + testing_logger::validate(|captured_logs| { + assert_eq!(captured_logs.len(), 1); + assert_eq!( + captured_logs[0].body, + "Could not load theme nonsolarized: Empty theme directory override and could not find theme elsewhere" + ); + assert_eq!(captured_logs[0].level, log::Level::Warn) + }); + } + + #[test] + fn test_can_debug_theme() { + testing_logger::setup(); + [true, false].iter().for_each(|debug| { + let mut manager = ThemeManager::new(Some(*debug), Some("".to_string())); + let config = Config::builder() + .add_source(ConfigFile::from_str( + " + [theme] + name = \"mytheme\" + + [colors] + Guidance = \"white\" + AlertInfo = \"xinetic\" + ", + FileFormat::Toml, + )) + .build() + .unwrap(); + manager + .load_theme_from_config("config_theme", config, 1) + .unwrap(); + testing_logger::validate(|captured_logs| { + if *debug { + assert_eq!(captured_logs.len(), 2); + assert_eq!( + captured_logs[0].body, + "Your theme config name is not the name of your loaded theme config_theme != mytheme" + ); + assert_eq!(captured_logs[0].level, log::Level::Warn); + assert_eq!( + captured_logs[1].body, + "Tried to load string as a color unsuccessfully: (AlertInfo=xinetic) No such color in palette" + ); + assert_eq!(captured_logs[1].level, log::Level::Warn) + } else { + assert_eq!(captured_logs.len(), 0) + } + }) + }) + } + + #[test] + fn test_can_parse_color_strings_correctly() { + assert_eq!( + from_string("brown").unwrap(), + Color::Rgb { + r: 165, + g: 42, + b: 42 + } + ); + + assert_eq!(from_string(""), Err("Empty string".into())); + + ["manatee", "caput mortuum", "123456"] + .iter() + .for_each(|inp| { + assert_eq!(from_string(inp), Err("No such color in palette".into())); + }); + + assert_eq!( + from_string("#ff1122").unwrap(), + Color::Rgb { + r: 255, + g: 17, + b: 34 + } + ); + ["#1122", "#ffaa112", "#brown"].iter().for_each(|inp| { + assert_eq!( + from_string(inp), + Err("Could not parse 3 hex values from string".into()) + ); + }); + + assert_eq!(from_string("@dark_grey").unwrap(), Color::DarkGrey); + assert_eq!( + from_string("@rgb_(255,255,255)").unwrap(), + Color::Rgb { + r: 255, + g: 255, + b: 255 + } + ); + assert_eq!(from_string("@ansi_(255)").unwrap(), Color::AnsiValue(255)); + ["@", "@DarkGray", "@Dark 4ay", "@ansi(256)"] + .iter() + .for_each(|inp| { + assert_eq!( + from_string(inp), + Err(format!( + "Could not convert color name {inp} to Crossterm color" + )) + ); + }); + } +} diff --git a/crates/atuin-client/src/utils.rs b/crates/atuin-client/src/utils.rs new file mode 100644 index 00000000..35d7db26 --- /dev/null +++ b/crates/atuin-client/src/utils.rs @@ -0,0 +1,14 @@ +pub(crate) fn get_hostname() -> String { + std::env::var("ATUIN_HOST_NAME") + .unwrap_or_else(|_| whoami::hostname().unwrap_or_else(|_| "unknown-host".to_string())) +} + +pub(crate) fn get_username() -> String { + std::env::var("ATUIN_HOST_USER") + .unwrap_or_else(|_| whoami::username().unwrap_or_else(|_| "unknown-user".to_string())) +} + +/// Returns a pair of the hostname and username, separated by a colon. +pub(crate) fn get_host_user() -> String { + format!("{}:{}", get_hostname(), get_username()) +} diff --git a/crates/atuin-client/tests/data/xonsh-history.sqlite b/crates/atuin-client/tests/data/xonsh-history.sqlite Binary files differnew file mode 100644 index 00000000..744fcf86 --- /dev/null +++ b/crates/atuin-client/tests/data/xonsh-history.sqlite diff --git a/crates/atuin-client/tests/data/xonsh/xonsh-82eafbf5-9f43-489a-80d2-61c7dc6ef542.json b/crates/atuin-client/tests/data/xonsh/xonsh-82eafbf5-9f43-489a-80d2-61c7dc6ef542.json new file mode 100644 index 00000000..339a09f1 --- /dev/null +++ b/crates/atuin-client/tests/data/xonsh/xonsh-82eafbf5-9f43-489a-80d2-61c7dc6ef542.json @@ -0,0 +1,12 @@ +{"locs": [ 69, 3371, 3451, 3978], + "index": {"offsets":{"__total__":0,"cmds":[{"__total__":10,"cwd":18,"inp":78,"rtn":96,"ts":[106,125,105]},{"__total__":149,"cwd":157,"inp":217,"rtn":234,"ts":[244,263,243]},9],"env":{"ATUIN_SESSION":314,"BASH_COMPLETIONS":370,"COLORTERM":433,"DBUS_SESSION_BUS_ADDRESS":474,"DESKTOP_SESSION":529,"DISPLAY":550,"GDMSESSION":570,"GIO_LAUNCHED_DESKTOP_FILE":609,"GIO_LAUNCHED_DESKTOP_FILE_PID":704,"GJS_DEBUG_OUTPUT":734,"GJS_DEBUG_TOPICS":764,"GNOME_DESKTOP_SESSION_ID":811,"GNOME_SETUP_DISPLAY":856,"GNOME_SHELL_SESSION_MODE":890,"GTK_MODULES":915,"HOME":942,"IM_CONFIG_PHASE":976,"INVOCATION_ID":998,"JOURNAL_STREAM":1052,"LANG":1071,"LOGNAME":1097,"MANAGERPID":1118,"MOZ_ENABLE_WAYLAND":1148,"PATH":1161,"PWD":1736,"PYENV_DIR":1802,"PYENV_HOOK_PATH":1874,"PYENV_ROOT":2048,"PYENV_SHELL":2086,"PYENV_VERSION":2111,"QT_ACCESSIBILITY":2141,"QT_IM_MODULE":2162,"SESSION_MANAGER":2189,"SHELL":2279,"SHLVL":2303,"SSH_AGENT_LAUNCHER":2330,"SSH_AUTH_SOCK":2364,"SSL_CERT_DIR":2415,"SSL_CERT_FILE":2458,"SYSTEMD_EXEC_PID":2525,"TERM":2541,"TERM_PROGRAM":2575,"TERM_PROGRAM_VERSION":2610,"THREAD_SUBPROCS":2657,"USER":2670,"USERNAME":2689,"WAYLAND_DISPLAY":2715,"WEZTERM_CONFIG_DIR":2750,"WEZTERM_CONFIG_FILE":2806,"WEZTERM_EXECUTABLE":2874,"WEZTERM_EXECUTABLE_DIR":2927,"WEZTERM_PANE":2957,"WEZTERM_UNIX_SOCKET":2986,"XAUTHORITY":3047,"XDG_CONFIG_DIRS":3116,"XDG_CURRENT_DESKTOP":3176,"XDG_DATA_DIRS":3209,"XDG_MENU_PREFIX":3316,"XDG_RUNTIME_DIR":3345,"XDG_SESSION_CLASS":3387,"XDG_SESSION_DESKTOP":3418,"XDG_SESSION_TYPE":3448,"XMODIFIERS":3473,"XONSHRC":3496,"XONSHRC_DIR":3594,"XONSH_CAPTURE_ALWAYS":3674,"XONSH_CONFIG_DIR":3698,"XONSH_DATA_DIR":3747,"XONSH_INTERACTIVE":3805,"XONSH_LOGIN":3825,"XONSH_VERSION":3847,"__total__":296},"locked":3869,"sessionid":3889,"ts":[3936,3956,3935]},"sizes":{"__total__":3978,"cmds":[{"__total__":137,"cwd":51,"inp":9,"rtn":1,"ts":[17,18,40]},{"__total__":136,"cwd":51,"inp":8,"rtn":1,"ts":[17,18,40]},278],"env":{"ATUIN_SESSION":34,"BASH_COMPLETIONS":48,"COLORTERM":11,"DBUS_SESSION_BUS_ADDRESS":34,"DESKTOP_SESSION":8,"DISPLAY":4,"GDMSESSION":8,"GIO_LAUNCHED_DESKTOP_FILE":60,"GIO_LAUNCHED_DESKTOP_FILE_PID":8,"GJS_DEBUG_OUTPUT":8,"GJS_DEBUG_TOPICS":17,"GNOME_DESKTOP_SESSION_ID":20,"GNOME_SETUP_DISPLAY":4,"GNOME_SHELL_SESSION_MODE":8,"GTK_MODULES":17,"HOME":13,"IM_CONFIG_PHASE":3,"INVOCATION_ID":34,"JOURNAL_STREAM":9,"LANG":13,"LOGNAME":5,"MANAGERPID":6,"MOZ_ENABLE_WAYLAND":3,"PATH":566,"PWD":51,"PYENV_DIR":51,"PYENV_HOOK_PATH":158,"PYENV_ROOT":21,"PYENV_SHELL":6,"PYENV_VERSION":8,"QT_ACCESSIBILITY":3,"QT_IM_MODULE":6,"SESSION_MANAGER":79,"SHELL":13,"SHLVL":3,"SSH_AGENT_LAUNCHER":15,"SSH_AUTH_SOCK":33,"SSL_CERT_DIR":24,"SSL_CERT_FILE":45,"SYSTEMD_EXEC_PID":6,"TERM":16,"TERM_PROGRAM":9,"TERM_PROGRAM_VERSION":26,"THREAD_SUBPROCS":3,"USER":5,"USERNAME":5,"WAYLAND_DISPLAY":11,"WEZTERM_CONFIG_DIR":31,"WEZTERM_CONFIG_FILE":44,"WEZTERM_EXECUTABLE":25,"WEZTERM_EXECUTABLE_DIR":12,"WEZTERM_PANE":4,"WEZTERM_UNIX_SOCKET":45,"XAUTHORITY":48,"XDG_CONFIG_DIRS":35,"XDG_CURRENT_DESKTOP":14,"XDG_DATA_DIRS":86,"XDG_MENU_PREFIX":8,"XDG_RUNTIME_DIR":19,"XDG_SESSION_CLASS":6,"XDG_SESSION_DESKTOP":8,"XDG_SESSION_TYPE":9,"XMODIFIERS":10,"XONSHRC":81,"XONSHRC_DIR":54,"XONSH_CAPTURE_ALWAYS":2,"XONSH_CONFIG_DIR":29,"XONSH_DATA_DIR":35,"XONSH_INTERACTIVE":3,"XONSH_LOGIN":3,"XONSH_VERSION":8,"__total__":3561},"locked":5,"sessionid":38,"ts":[18,18,41]}}, + "data": {"cmds": [{"cwd": "\/home\/user\/Documents\/code\/atuin\/atuin-client", "inp": "false\n", "rtn": 1, "ts": [1707241291.142516, 1707241291.1527853] +} +, {"cwd": "\/home\/user\/Documents\/code\/atuin\/atuin-client", "inp": "exit\n", "rtn": 0, "ts": [1707241292.271584, 1707241292.2758434] +} +] +, "env": {"ATUIN_SESSION": "018d7f82ad167dc4888ca0bf294d2bfd", "BASH_COMPLETIONS": "\/usr\/share\/bash-completion\/bash_completion", "COLORTERM": "truecolor", "DBUS_SESSION_BUS_ADDRESS": "unix:path=\/run\/user\/1000\/bus", "DESKTOP_SESSION": "ubuntu", "DISPLAY": ":0", "GDMSESSION": "ubuntu", "GIO_LAUNCHED_DESKTOP_FILE": "\/usr\/share\/applications\/org.wezfurlong.wezterm.desktop", "GIO_LAUNCHED_DESKTOP_FILE_PID": "196859", "GJS_DEBUG_OUTPUT": "stderr", "GJS_DEBUG_TOPICS": "JS ERROR;JS LOG", "GNOME_DESKTOP_SESSION_ID": "this-is-deprecated", "GNOME_SETUP_DISPLAY": ":1", "GNOME_SHELL_SESSION_MODE": "ubuntu", "GTK_MODULES": "gail:atk-bridge", "HOME": "\/home\/user", "IM_CONFIG_PHASE": "1", "INVOCATION_ID": "4f121e7ad56c41a6b84aa3cbe1ad61fa", "JOURNAL_STREAM": "8:37187", "LANG": "en_US.UTF-8", "LOGNAME": "user", "MANAGERPID": "2118", "MOZ_ENABLE_WAYLAND": "1", "PATH": "\/home\/user\/.pyenv\/versions\/3.12.0\/bin:\/home\/user\/.pyenv\/libexec:\/home\/user\/.pyenv\/plugins\/python-build\/bin:\/home\/user\/.pyenv\/plugins\/pyenv-virtualenv\/bin:\/home\/user\/.pyenv\/plugins\/pyenv-update\/bin:\/home\/user\/.pyenv\/plugins\/pyenv-doctor\/bin:\/home\/user\/.cargo\/bin:\/home\/user\/.pyenv\/shims:\/home\/user\/.pyenv\/bin:\/home\/user\/bin:\/home\/user\/bin:\/usr\/local\/sbin:\/usr\/local\/bin:\/usr\/sbin:\/usr\/bin:\/sbin:\/bin:\/usr\/games:\/usr\/local\/games:\/snap\/bin:\/snap\/bin:\/home\/user\/.local\/share\/JetBrains\/Toolbox\/scripts", "PWD": "\/home\/user\/Documents\/code\/atuin\/atuin-client", "PYENV_DIR": "\/home\/user\/Documents\/code\/atuin\/atuin-client", "PYENV_HOOK_PATH": "\/home\/user\/.pyenv\/pyenv.d:\/usr\/local\/etc\/pyenv.d:\/etc\/pyenv.d:\/usr\/lib\/pyenv\/hooks:\/home\/user\/.pyenv\/plugins\/pyenv-virtualenv\/etc\/pyenv.d", "PYENV_ROOT": "\/home\/user\/.pyenv", "PYENV_SHELL": "bash", "PYENV_VERSION": "3.12.0", "QT_ACCESSIBILITY": "1", "QT_IM_MODULE": "ibus", "SESSION_MANAGER": "local\/box:@\/tmp\/.ICE-unix\/2452,unix\/box:\/tmp\/.ICE-unix\/2452", "SHELL": "\/bin\/bash", "SHLVL": "1", "SSH_AGENT_LAUNCHER": "gnome-keyring", "SSH_AUTH_SOCK": "\/run\/user\/1000\/keyring\/ssh", "SSL_CERT_DIR": "\/usr\/lib\/ssl\/certs", "SSL_CERT_FILE": "\/usr\/lib\/ssl\/certs\/ca-certificates.crt", "SYSTEMD_EXEC_PID": "2470", "TERM": "xterm-256color", "TERM_PROGRAM": "WezTerm", "TERM_PROGRAM_VERSION": "20240127-113634-bbcac864", "THREAD_SUBPROCS": "1", "USER": "user", "USERNAME": "user", "WAYLAND_DISPLAY": "wayland-0", "WEZTERM_CONFIG_DIR": "\/home\/user\/.config\/wezterm", "WEZTERM_CONFIG_FILE": "\/home\/user\/.config\/wezterm\/wezterm.lua", "WEZTERM_EXECUTABLE": "\/usr\/bin\/wezterm-gui", "WEZTERM_EXECUTABLE_DIR": "\/usr\/bin", "WEZTERM_PANE": "41", "WEZTERM_UNIX_SOCKET": "\/run\/user\/1000\/wezterm\/gui-sock-196859", "XAUTHORITY": "\/run\/user\/1000\/.mutter-Xwaylandauth.T986H2", "XDG_CONFIG_DIRS": "\/etc\/xdg\/xdg-ubuntu:\/etc\/xdg", "XDG_CURRENT_DESKTOP": "ubuntu:GNOME", "XDG_DATA_DIRS": "\/usr\/share\/ubuntu:\/usr\/local\/share\/:\/usr\/share\/:\/var\/lib\/snapd\/desktop", "XDG_MENU_PREFIX": "gnome-", "XDG_RUNTIME_DIR": "\/run\/user\/1000", "XDG_SESSION_CLASS": "user", "XDG_SESSION_DESKTOP": "ubuntu", "XDG_SESSION_TYPE": "wayland", "XMODIFIERS": "@im=ibus", "XONSHRC": "\/etc\/xonsh\/xonshrc:\/home\/user\/.config\/xonsh\/rc.xsh:\/home\/user\/.xonshrc", "XONSHRC_DIR": "\/etc\/xonsh\/rc.d:\/home\/user\/.config\/xonsh\/rc.d", "XONSH_CAPTURE_ALWAYS": "", "XONSH_CONFIG_DIR": "\/home\/user\/.config\/xonsh", "XONSH_DATA_DIR": "\/home\/user\/.local\/share\/xonsh", "XONSH_INTERACTIVE": "1", "XONSH_LOGIN": "1", "XONSH_VERSION": "0.14.2"} +, "locked": false, "sessionid": "82eafbf5-9f43-489a-80d2-61c7dc6ef542", "ts": [1707241286.9361255, 1707241292.3081477] +} + +} diff --git a/crates/atuin-client/tests/data/xonsh/xonsh-de16af90-9148-4461-8df3-5b5659c6420d.json b/crates/atuin-client/tests/data/xonsh/xonsh-de16af90-9148-4461-8df3-5b5659c6420d.json new file mode 100644 index 00000000..72694f04 --- /dev/null +++ b/crates/atuin-client/tests/data/xonsh/xonsh-de16af90-9148-4461-8df3-5b5659c6420d.json @@ -0,0 +1,12 @@ +{"locs": [ 69, 3372, 3452, 3936], + "index": {"offsets":{"__total__":0,"cmds":[{"__total__":10,"cwd":18,"inp":64,"rtn":94,"ts":[104,124,103]},{"__total__":148,"cwd":156,"inp":202,"rtn":220,"ts":[230,250,229]},9],"env":{"ATUIN_SESSION":300,"BASH_COMPLETIONS":356,"COLORTERM":419,"DBUS_SESSION_BUS_ADDRESS":460,"DESKTOP_SESSION":515,"DISPLAY":536,"GDMSESSION":556,"GIO_LAUNCHED_DESKTOP_FILE":595,"GIO_LAUNCHED_DESKTOP_FILE_PID":690,"GJS_DEBUG_OUTPUT":720,"GJS_DEBUG_TOPICS":750,"GNOME_DESKTOP_SESSION_ID":797,"GNOME_SETUP_DISPLAY":842,"GNOME_SHELL_SESSION_MODE":876,"GTK_MODULES":901,"HOME":928,"IM_CONFIG_PHASE":962,"INVOCATION_ID":984,"JOURNAL_STREAM":1038,"LANG":1057,"LOGNAME":1083,"MANAGERPID":1104,"MOZ_ENABLE_WAYLAND":1134,"PATH":1147,"PWD":1722,"PYENV_DIR":1774,"PYENV_HOOK_PATH":1832,"PYENV_ROOT":2006,"PYENV_SHELL":2044,"PYENV_VERSION":2069,"QT_ACCESSIBILITY":2099,"QT_IM_MODULE":2120,"SESSION_MANAGER":2147,"SHELL":2237,"SHLVL":2261,"SSH_AGENT_LAUNCHER":2288,"SSH_AUTH_SOCK":2322,"SSL_CERT_DIR":2373,"SSL_CERT_FILE":2416,"SYSTEMD_EXEC_PID":2483,"TERM":2499,"TERM_PROGRAM":2533,"TERM_PROGRAM_VERSION":2568,"THREAD_SUBPROCS":2615,"USER":2628,"USERNAME":2647,"WAYLAND_DISPLAY":2673,"WEZTERM_CONFIG_DIR":2708,"WEZTERM_CONFIG_FILE":2764,"WEZTERM_EXECUTABLE":2832,"WEZTERM_EXECUTABLE_DIR":2885,"WEZTERM_PANE":2915,"WEZTERM_UNIX_SOCKET":2944,"XAUTHORITY":3005,"XDG_CONFIG_DIRS":3074,"XDG_CURRENT_DESKTOP":3134,"XDG_DATA_DIRS":3167,"XDG_MENU_PREFIX":3274,"XDG_RUNTIME_DIR":3303,"XDG_SESSION_CLASS":3345,"XDG_SESSION_DESKTOP":3376,"XDG_SESSION_TYPE":3406,"XMODIFIERS":3431,"XONSHRC":3454,"XONSHRC_DIR":3552,"XONSH_CAPTURE_ALWAYS":3632,"XONSH_CONFIG_DIR":3656,"XONSH_DATA_DIR":3705,"XONSH_INTERACTIVE":3763,"XONSH_LOGIN":3783,"XONSH_VERSION":3805,"__total__":282},"locked":3827,"sessionid":3847,"ts":[3894,3914,3893]},"sizes":{"__total__":3936,"cmds":[{"__total__":136,"cwd":37,"inp":21,"rtn":1,"ts":[18,18,41]},{"__total__":123,"cwd":37,"inp":9,"rtn":1,"ts":[18,17,40]},264],"env":{"ATUIN_SESSION":34,"BASH_COMPLETIONS":48,"COLORTERM":11,"DBUS_SESSION_BUS_ADDRESS":34,"DESKTOP_SESSION":8,"DISPLAY":4,"GDMSESSION":8,"GIO_LAUNCHED_DESKTOP_FILE":60,"GIO_LAUNCHED_DESKTOP_FILE_PID":8,"GJS_DEBUG_OUTPUT":8,"GJS_DEBUG_TOPICS":17,"GNOME_DESKTOP_SESSION_ID":20,"GNOME_SETUP_DISPLAY":4,"GNOME_SHELL_SESSION_MODE":8,"GTK_MODULES":17,"HOME":13,"IM_CONFIG_PHASE":3,"INVOCATION_ID":34,"JOURNAL_STREAM":9,"LANG":13,"LOGNAME":5,"MANAGERPID":6,"MOZ_ENABLE_WAYLAND":3,"PATH":566,"PWD":37,"PYENV_DIR":37,"PYENV_HOOK_PATH":158,"PYENV_ROOT":21,"PYENV_SHELL":6,"PYENV_VERSION":8,"QT_ACCESSIBILITY":3,"QT_IM_MODULE":6,"SESSION_MANAGER":79,"SHELL":13,"SHLVL":3,"SSH_AGENT_LAUNCHER":15,"SSH_AUTH_SOCK":33,"SSL_CERT_DIR":24,"SSL_CERT_FILE":45,"SYSTEMD_EXEC_PID":6,"TERM":16,"TERM_PROGRAM":9,"TERM_PROGRAM_VERSION":26,"THREAD_SUBPROCS":3,"USER":5,"USERNAME":5,"WAYLAND_DISPLAY":11,"WEZTERM_CONFIG_DIR":31,"WEZTERM_CONFIG_FILE":44,"WEZTERM_EXECUTABLE":25,"WEZTERM_EXECUTABLE_DIR":12,"WEZTERM_PANE":4,"WEZTERM_UNIX_SOCKET":45,"XAUTHORITY":48,"XDG_CONFIG_DIRS":35,"XDG_CURRENT_DESKTOP":14,"XDG_DATA_DIRS":86,"XDG_MENU_PREFIX":8,"XDG_RUNTIME_DIR":19,"XDG_SESSION_CLASS":6,"XDG_SESSION_DESKTOP":8,"XDG_SESSION_TYPE":9,"XMODIFIERS":10,"XONSHRC":81,"XONSHRC_DIR":54,"XONSH_CAPTURE_ALWAYS":2,"XONSH_CONFIG_DIR":29,"XONSH_DATA_DIR":35,"XONSH_INTERACTIVE":3,"XONSH_LOGIN":3,"XONSH_VERSION":8,"__total__":3533},"locked":5,"sessionid":38,"ts":[18,18,41]}}, + "data": {"cmds": [{"cwd": "\/home\/user\/Documents\/code\/atuin", "inp": "echo hello world!\n", "rtn": 0, "ts": [1707193079.4782722, 1707193079.4829233] +} +, {"cwd": "\/home\/user\/Documents\/code\/atuin", "inp": "ls -l\n", "rtn": 0, "ts": [1707193081.7063284, 1707193081.727617] +} +] +, "env": {"ATUIN_SESSION": "018d7ca2e953742e9826012f30115040", "BASH_COMPLETIONS": "\/usr\/share\/bash-completion\/bash_completion", "COLORTERM": "truecolor", "DBUS_SESSION_BUS_ADDRESS": "unix:path=\/run\/user\/1000\/bus", "DESKTOP_SESSION": "ubuntu", "DISPLAY": ":0", "GDMSESSION": "ubuntu", "GIO_LAUNCHED_DESKTOP_FILE": "\/usr\/share\/applications\/org.wezfurlong.wezterm.desktop", "GIO_LAUNCHED_DESKTOP_FILE_PID": "196859", "GJS_DEBUG_OUTPUT": "stderr", "GJS_DEBUG_TOPICS": "JS ERROR;JS LOG", "GNOME_DESKTOP_SESSION_ID": "this-is-deprecated", "GNOME_SETUP_DISPLAY": ":1", "GNOME_SHELL_SESSION_MODE": "ubuntu", "GTK_MODULES": "gail:atk-bridge", "HOME": "\/home\/user", "IM_CONFIG_PHASE": "1", "INVOCATION_ID": "4f121e7ad56c41a6b84aa3cbe1ad61fa", "JOURNAL_STREAM": "8:37187", "LANG": "en_US.UTF-8", "LOGNAME": "user", "MANAGERPID": "2118", "MOZ_ENABLE_WAYLAND": "1", "PATH": "\/home\/user\/.pyenv\/versions\/3.12.0\/bin:\/home\/user\/.pyenv\/libexec:\/home\/user\/.pyenv\/plugins\/python-build\/bin:\/home\/user\/.pyenv\/plugins\/pyenv-virtualenv\/bin:\/home\/user\/.pyenv\/plugins\/pyenv-update\/bin:\/home\/user\/.pyenv\/plugins\/pyenv-doctor\/bin:\/home\/user\/.cargo\/bin:\/home\/user\/.pyenv\/shims:\/home\/user\/.pyenv\/bin:\/home\/user\/bin:\/home\/user\/bin:\/usr\/local\/sbin:\/usr\/local\/bin:\/usr\/sbin:\/usr\/bin:\/sbin:\/bin:\/usr\/games:\/usr\/local\/games:\/snap\/bin:\/snap\/bin:\/home\/user\/.local\/share\/JetBrains\/Toolbox\/scripts", "PWD": "\/home\/user\/Documents\/code\/atuin", "PYENV_DIR": "\/home\/user\/Documents\/code\/atuin", "PYENV_HOOK_PATH": "\/home\/user\/.pyenv\/pyenv.d:\/usr\/local\/etc\/pyenv.d:\/etc\/pyenv.d:\/usr\/lib\/pyenv\/hooks:\/home\/user\/.pyenv\/plugins\/pyenv-virtualenv\/etc\/pyenv.d", "PYENV_ROOT": "\/home\/user\/.pyenv", "PYENV_SHELL": "bash", "PYENV_VERSION": "3.12.0", "QT_ACCESSIBILITY": "1", "QT_IM_MODULE": "ibus", "SESSION_MANAGER": "local\/box:@\/tmp\/.ICE-unix\/2452,unix\/box:\/tmp\/.ICE-unix\/2452", "SHELL": "\/bin\/bash", "SHLVL": "1", "SSH_AGENT_LAUNCHER": "gnome-keyring", "SSH_AUTH_SOCK": "\/run\/user\/1000\/keyring\/ssh", "SSL_CERT_DIR": "\/usr\/lib\/ssl\/certs", "SSL_CERT_FILE": "\/usr\/lib\/ssl\/certs\/ca-certificates.crt", "SYSTEMD_EXEC_PID": "2470", "TERM": "xterm-256color", "TERM_PROGRAM": "WezTerm", "TERM_PROGRAM_VERSION": "20240127-113634-bbcac864", "THREAD_SUBPROCS": "1", "USER": "user", "USERNAME": "user", "WAYLAND_DISPLAY": "wayland-0", "WEZTERM_CONFIG_DIR": "\/home\/user\/.config\/wezterm", "WEZTERM_CONFIG_FILE": "\/home\/user\/.config\/wezterm\/wezterm.lua", "WEZTERM_EXECUTABLE": "\/usr\/bin\/wezterm-gui", "WEZTERM_EXECUTABLE_DIR": "\/usr\/bin", "WEZTERM_PANE": "38", "WEZTERM_UNIX_SOCKET": "\/run\/user\/1000\/wezterm\/gui-sock-196859", "XAUTHORITY": "\/run\/user\/1000\/.mutter-Xwaylandauth.T986H2", "XDG_CONFIG_DIRS": "\/etc\/xdg\/xdg-ubuntu:\/etc\/xdg", "XDG_CURRENT_DESKTOP": "ubuntu:GNOME", "XDG_DATA_DIRS": "\/usr\/share\/ubuntu:\/usr\/local\/share\/:\/usr\/share\/:\/var\/lib\/snapd\/desktop", "XDG_MENU_PREFIX": "gnome-", "XDG_RUNTIME_DIR": "\/run\/user\/1000", "XDG_SESSION_CLASS": "user", "XDG_SESSION_DESKTOP": "ubuntu", "XDG_SESSION_TYPE": "wayland", "XMODIFIERS": "@im=ibus", "XONSHRC": "\/etc\/xonsh\/xonshrc:\/home\/user\/.config\/xonsh\/rc.xsh:\/home\/user\/.xonshrc", "XONSHRC_DIR": "\/etc\/xonsh\/rc.d:\/home\/user\/.config\/xonsh\/rc.d", "XONSH_CAPTURE_ALWAYS": "", "XONSH_CONFIG_DIR": "\/home\/user\/.config\/xonsh", "XONSH_DATA_DIR": "\/home\/user\/.local\/share\/xonsh", "XONSH_INTERACTIVE": "1", "XONSH_LOGIN": "1", "XONSH_VERSION": "0.14.2"} +, "locked": false, "sessionid": "de16af90-9148-4461-8df3-5b5659c6420d", "ts": [1707193067.8615997, 1707193089.2513068] +} + +} diff --git a/crates/atuin-common/Cargo.toml b/crates/atuin-common/Cargo.toml new file mode 100644 index 00000000..811b0bdb --- /dev/null +++ b/crates/atuin-common/Cargo.toml @@ -0,0 +1,31 @@ +[package] +name = "atuin-common" +edition = "2024" +description = "common library for atuin" + +rust-version = { workspace = true } +version = { workspace = true } +authors = { workspace = true } +license = { workspace = true } +homepage = { workspace = true } +repository = { workspace = true } + +# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html + +[dependencies] +time = { workspace = true } +serde = { workspace = true } +uuid = { workspace = true } +typed-builder = { workspace = true } +eyre = { workspace = true } +sqlx = { workspace = true } +semver = { workspace = true } +thiserror = { workspace = true } +directories = { workspace = true } +sysinfo = "0.30.7" +base64 = { workspace = true } +getrandom = "0.2" +rustls = { workspace = true } + +[dev-dependencies] +pretty_assertions = { workspace = true } diff --git a/crates/atuin-common/src/api.rs b/crates/atuin-common/src/api.rs new file mode 100644 index 00000000..efc17163 --- /dev/null +++ b/crates/atuin-common/src/api.rs @@ -0,0 +1,136 @@ +use semver::Version; +use serde::{Deserialize, Serialize}; +use std::borrow::Cow; +use std::sync::LazyLock; +use time::OffsetDateTime; + +// the usage of X- has been deprecated for quite along time, it turns out +pub static ATUIN_HEADER_VERSION: &str = "Atuin-Version"; +pub static ATUIN_CARGO_VERSION: &str = env!("CARGO_PKG_VERSION"); + +pub static ATUIN_VERSION: LazyLock<Version> = + LazyLock::new(|| Version::parse(ATUIN_CARGO_VERSION).expect("failed to parse self semver")); + +#[derive(Debug, Serialize, Deserialize)] +pub struct UserResponse { + pub username: String, +} + +#[derive(Debug, Serialize, Deserialize)] +pub struct RegisterRequest { + pub email: String, + pub username: String, + pub password: String, +} + +#[derive(Debug, Serialize, Deserialize)] +pub struct RegisterResponse { + pub session: String, +} + +#[derive(Debug, Serialize, Deserialize)] +pub struct DeleteUserResponse {} + +#[derive(Debug, Serialize, Deserialize)] +pub struct ChangePasswordRequest { + pub current_password: String, + pub new_password: String, +} + +#[derive(Debug, Serialize, Deserialize)] +pub struct ChangePasswordResponse {} + +#[derive(Debug, Serialize, Deserialize)] +pub struct LoginRequest { + pub username: String, + pub password: String, +} + +#[derive(Debug, Serialize, Deserialize)] +pub struct LoginResponse { + pub session: String, +} + +#[derive(Debug, Serialize, Deserialize)] +pub struct AddHistoryRequest { + pub id: String, + #[serde(with = "time::serde::rfc3339")] + pub timestamp: OffsetDateTime, + pub data: String, + pub hostname: String, +} + +#[derive(Debug, Serialize, Deserialize)] +pub struct CountResponse { + pub count: i64, +} + +#[derive(Debug, Serialize, Deserialize)] +pub struct SyncHistoryRequest { + #[serde(with = "time::serde::rfc3339")] + pub sync_ts: OffsetDateTime, + #[serde(with = "time::serde::rfc3339")] + pub history_ts: OffsetDateTime, + pub host: String, +} + +#[derive(Debug, Serialize, Deserialize)] +pub struct SyncHistoryResponse { + pub history: Vec<String>, +} + +#[derive(Debug, Serialize, Deserialize)] +pub struct ErrorResponse<'a> { + pub reason: Cow<'a, str>, +} + +#[derive(Debug, Serialize, Deserialize)] +pub struct IndexResponse { + pub homage: String, + pub version: String, +} + +#[derive(Debug, Serialize, Deserialize)] +pub struct StatusResponse { + pub count: i64, + pub username: String, + pub deleted: Vec<String>, + + // These could/should also go on the index of the server + // However, we do not request the server index as a part of normal sync + // I'd rather slightly increase the size of this response, than add an extra HTTP request + pub page_size: i64, // max page size supported by the server + pub version: String, +} + +#[derive(Debug, Serialize, Deserialize)] +pub struct DeleteHistoryRequest { + pub client_id: String, +} + +#[derive(Debug, Serialize, Deserialize)] +pub struct MessageResponse { + pub message: String, +} + +#[derive(Debug, Serialize, Deserialize)] +pub struct MeResponse { + pub username: String, +} + +// Hub CLI authentication types + +/// Response from POST /auth/cli/code - generates a code for CLI auth +#[derive(Debug, Serialize, Deserialize)] +pub struct CliCodeResponse { + pub code: String, +} + +/// Response from GET /auth/cli/verify?code=<code> - polls for authorization +#[derive(Debug, Serialize, Deserialize)] +pub struct CliVerifyResponse { + /// Session token, present only when authorization is complete + pub token: Option<String>, + pub success: Option<bool>, + pub error: Option<String>, +} diff --git a/crates/atuin-common/src/calendar.rs b/crates/atuin-common/src/calendar.rs new file mode 100644 index 00000000..d3b1d921 --- /dev/null +++ b/crates/atuin-common/src/calendar.rs @@ -0,0 +1,16 @@ +// Calendar data +use serde::{Serialize, Deserialize}; + +pub enum TimePeriod { + YEAR, + MONTH, + DAY, +} + +#[derive(Debug, Serialize, Deserialize)] +pub struct TimePeriodInfo { + pub count: u64, + + // TODO: Use this for merkle tree magic + pub hash: String, +} diff --git a/crates/atuin-common/src/lib.rs b/crates/atuin-common/src/lib.rs new file mode 100644 index 00000000..91164a82 --- /dev/null +++ b/crates/atuin-common/src/lib.rs @@ -0,0 +1,60 @@ +#![deny(unsafe_code)] + +/// Defines a new UUID type wrapper +macro_rules! new_uuid { + ($name:ident) => { + #[derive( + Debug, + Copy, + Clone, + PartialEq, + Eq, + Hash, + PartialOrd, + Ord, + serde::Serialize, + serde::Deserialize, + )] + #[serde(transparent)] + pub struct $name(pub Uuid); + + impl<DB: sqlx::Database> sqlx::Type<DB> for $name + where + Uuid: sqlx::Type<DB>, + { + fn type_info() -> <DB as sqlx::Database>::TypeInfo { + Uuid::type_info() + } + } + + impl<'r, DB: sqlx::Database> sqlx::Decode<'r, DB> for $name + where + Uuid: sqlx::Decode<'r, DB>, + { + fn decode( + value: DB::ValueRef<'r>, + ) -> std::result::Result<Self, sqlx::error::BoxDynError> { + Uuid::decode(value).map(Self) + } + } + + impl<'q, DB: sqlx::Database> sqlx::Encode<'q, DB> for $name + where + Uuid: sqlx::Encode<'q, DB>, + { + fn encode_by_ref( + &self, + buf: &mut DB::ArgumentBuffer<'q>, + ) -> Result<sqlx::encode::IsNull, Box<dyn std::error::Error + Send + Sync + 'static>> + { + self.0.encode_by_ref(buf) + } + } + }; +} + +pub mod api; +pub mod record; +pub mod shell; +pub mod tls; +pub mod utils; diff --git a/crates/atuin-common/src/record.rs b/crates/atuin-common/src/record.rs new file mode 100644 index 00000000..e6ce2647 --- /dev/null +++ b/crates/atuin-common/src/record.rs @@ -0,0 +1,426 @@ +use std::collections::HashMap; + +use eyre::Result; +use serde::{Deserialize, Serialize}; +use typed_builder::TypedBuilder; +use uuid::Uuid; + +#[derive(Clone, Debug, PartialEq)] +pub struct DecryptedData(pub Vec<u8>); + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct EncryptedData { + pub data: String, + pub content_encryption_key: String, +} + +#[derive(Debug, PartialEq, PartialOrd, Ord, Eq)] +pub struct Diff { + pub host: HostId, + pub tag: String, + pub local: Option<RecordIdx>, + pub remote: Option<RecordIdx>, +} + +#[derive(Serialize, Deserialize, Clone, Debug, PartialEq)] +pub struct Host { + pub id: HostId, + pub name: String, +} + +impl Host { + pub fn new(id: HostId) -> Self { + Host { + id, + name: String::new(), + } + } +} + +new_uuid!(RecordId); +new_uuid!(HostId); + +pub type RecordIdx = u64; + +/// A single record stored inside of our local database +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, TypedBuilder)] +pub struct Record<Data> { + /// a unique ID + #[builder(default = RecordId(crate::utils::uuid_v7()))] + pub id: RecordId, + + /// The integer record ID. This is only unique per (host, tag). + pub idx: RecordIdx, + + /// The unique ID of the host. + // TODO(ellie): Optimize the storage here. We use a bunch of IDs, and currently store + // as strings. I would rather avoid normalization, so store as UUID binary instead of + // encoding to a string and wasting much more storage. + pub host: Host, + + /// The creation time in nanoseconds since unix epoch + #[builder(default = time::OffsetDateTime::now_utc().unix_timestamp_nanos() as u64)] + pub timestamp: u64, + + /// The version the data in the entry conforms to + // However we want to track versions for this tag, eg v2 + pub version: String, + + /// The type of data we are storing here. Eg, "history" + pub tag: String, + + /// Some data. This can be anything you wish to store. Use the tag field to know how to handle it. + pub data: Data, +} + +/// Extra data from the record that should be encoded in the data +#[derive(Debug, Copy, Clone)] +pub struct AdditionalData<'a> { + pub id: &'a RecordId, + pub idx: &'a u64, + pub version: &'a str, + pub tag: &'a str, + pub host: &'a HostId, +} + +impl<Data> Record<Data> { + pub fn append(&self, data: Vec<u8>) -> Record<DecryptedData> { + Record::builder() + .host(self.host.clone()) + .version(self.version.clone()) + .idx(self.idx + 1) + .tag(self.tag.clone()) + .data(DecryptedData(data)) + .build() + } +} + +/// An index representing the current state of the record stores +/// This can be both remote, or local, and compared in either direction +#[derive(Debug, Serialize, Deserialize)] +pub struct RecordStatus { + // A map of host -> tag -> max(idx) + pub hosts: HashMap<HostId, HashMap<String, RecordIdx>>, +} + +impl Default for RecordStatus { + fn default() -> Self { + Self::new() + } +} + +impl Extend<(HostId, String, RecordIdx)> for RecordStatus { + fn extend<T: IntoIterator<Item = (HostId, String, RecordIdx)>>(&mut self, iter: T) { + for (host, tag, tail_idx) in iter { + self.set_raw(host, tag, tail_idx); + } + } +} + +impl RecordStatus { + pub fn new() -> RecordStatus { + RecordStatus { + hosts: HashMap::new(), + } + } + + /// Insert a new tail record into the store + pub fn set(&mut self, tail: Record<DecryptedData>) { + self.set_raw(tail.host.id, tail.tag, tail.idx) + } + + pub fn set_raw(&mut self, host: HostId, tag: String, tail_id: RecordIdx) { + self.hosts.entry(host).or_default().insert(tag, tail_id); + } + + pub fn get(&self, host: HostId, tag: String) -> Option<RecordIdx> { + self.hosts.get(&host).and_then(|v| v.get(&tag)).cloned() + } + + /// Diff this index with another, likely remote index. + /// The two diffs can then be reconciled, and the optimal change set calculated + /// Returns a tuple, with (host, tag, Option(OTHER)) + /// OTHER is set to the value of the idx on the other machine. If it is greater than our index, + /// then we need to do some downloading. If it is smaller, then we need to do some uploading + /// Note that we cannot upload if we are not the owner of the record store - hosts can only + /// write to their own store. + pub fn diff(&self, other: &Self) -> Vec<Diff> { + let mut ret = Vec::new(); + + // First, we check if other has everything that self has + for (host, tag_map) in self.hosts.iter() { + for (tag, idx) in tag_map.iter() { + match other.get(*host, tag.clone()) { + // The other store is all up to date! No diff. + Some(t) if t.eq(idx) => continue, + + // The other store does exist, and it is either ahead or behind us. A diff regardless + Some(t) => ret.push(Diff { + host: *host, + tag: tag.clone(), + local: Some(*idx), + remote: Some(t), + }), + + // The other store does not exist :O + None => ret.push(Diff { + host: *host, + tag: tag.clone(), + local: Some(*idx), + remote: None, + }), + }; + } + } + + // At this point, there is a single case we have not yet considered. + // If the other store knows of a tag that we are not yet aware of, then the diff will be missed + + // account for that! + for (host, tag_map) in other.hosts.iter() { + for (tag, idx) in tag_map.iter() { + match self.get(*host, tag.clone()) { + // If we have this host/tag combo, the comparison and diff will have already happened above + Some(_) => continue, + + None => ret.push(Diff { + host: *host, + tag: tag.clone(), + remote: Some(*idx), + local: None, + }), + }; + } + } + + // Stability is a nice property to have + ret.sort(); + ret + } +} + +pub trait Encryption { + fn re_encrypt( + data: EncryptedData, + ad: AdditionalData, + old_key: &[u8; 32], + new_key: &[u8; 32], + ) -> Result<EncryptedData> { + let data = Self::decrypt(data, ad, old_key)?; + Ok(Self::encrypt(data, ad, new_key)) + } + fn encrypt(data: DecryptedData, ad: AdditionalData, key: &[u8; 32]) -> EncryptedData; + fn decrypt(data: EncryptedData, ad: AdditionalData, key: &[u8; 32]) -> Result<DecryptedData>; +} + +impl Record<DecryptedData> { + pub fn encrypt<E: Encryption>(self, key: &[u8; 32]) -> Record<EncryptedData> { + let ad = AdditionalData { + id: &self.id, + version: &self.version, + tag: &self.tag, + host: &self.host.id, + idx: &self.idx, + }; + Record { + data: E::encrypt(self.data, ad, key), + id: self.id, + host: self.host, + idx: self.idx, + timestamp: self.timestamp, + version: self.version, + tag: self.tag, + } + } +} + +impl Record<EncryptedData> { + pub fn decrypt<E: Encryption>(self, key: &[u8; 32]) -> Result<Record<DecryptedData>> { + let ad = AdditionalData { + id: &self.id, + version: &self.version, + tag: &self.tag, + host: &self.host.id, + idx: &self.idx, + }; + Ok(Record { + data: E::decrypt(self.data, ad, key)?, + id: self.id, + host: self.host, + idx: self.idx, + timestamp: self.timestamp, + version: self.version, + tag: self.tag, + }) + } + + pub fn re_encrypt<E: Encryption>( + self, + old_key: &[u8; 32], + new_key: &[u8; 32], + ) -> Result<Record<EncryptedData>> { + let ad = AdditionalData { + id: &self.id, + version: &self.version, + tag: &self.tag, + host: &self.host.id, + idx: &self.idx, + }; + Ok(Record { + data: E::re_encrypt(self.data, ad, old_key, new_key)?, + id: self.id, + host: self.host, + idx: self.idx, + timestamp: self.timestamp, + version: self.version, + tag: self.tag, + }) + } +} + +#[cfg(test)] +mod tests { + use crate::record::{Host, HostId}; + + use super::{DecryptedData, Diff, Record, RecordStatus}; + use pretty_assertions::assert_eq; + + fn test_record() -> Record<DecryptedData> { + Record::builder() + .host(Host::new(HostId(crate::utils::uuid_v7()))) + .version("v1".into()) + .tag(crate::utils::uuid_v7().simple().to_string()) + .data(DecryptedData(vec![0, 1, 2, 3])) + .idx(0) + .build() + } + + #[test] + fn record_index() { + let mut index = RecordStatus::new(); + let record = test_record(); + + index.set(record.clone()); + + let tail = index.get(record.host.id, record.tag); + + assert_eq!( + record.idx, + tail.expect("tail not in store"), + "tail in store did not match" + ); + } + + #[test] + fn record_index_overwrite() { + let mut index = RecordStatus::new(); + let record = test_record(); + let child = record.append(vec![1, 2, 3]); + + index.set(record.clone()); + index.set(child.clone()); + + let tail = index.get(record.host.id, record.tag); + + assert_eq!( + child.idx, + tail.expect("tail not in store"), + "tail in store did not match" + ); + } + + #[test] + fn record_index_no_diff() { + // Here, they both have the same version and should have no diff + + let mut index1 = RecordStatus::new(); + let mut index2 = RecordStatus::new(); + + let record1 = test_record(); + + index1.set(record1.clone()); + index2.set(record1); + + let diff = index1.diff(&index2); + + assert_eq!(0, diff.len(), "expected empty diff"); + } + + #[test] + fn record_index_single_diff() { + // Here, they both have the same stores, but one is ahead by a single record + + let mut index1 = RecordStatus::new(); + let mut index2 = RecordStatus::new(); + + let record1 = test_record(); + let record2 = record1.append(vec![1, 2, 3]); + + index1.set(record1); + index2.set(record2.clone()); + + let diff = index1.diff(&index2); + + assert_eq!(1, diff.len(), "expected single diff"); + assert_eq!( + diff[0], + Diff { + host: record2.host.id, + tag: record2.tag, + remote: Some(1), + local: Some(0) + } + ); + } + + #[test] + fn record_index_multi_diff() { + // A much more complex case, with a bunch more checks + let mut index1 = RecordStatus::new(); + let mut index2 = RecordStatus::new(); + + let store1record1 = test_record(); + let store1record2 = store1record1.append(vec![1, 2, 3]); + + let store2record1 = test_record(); + let store2record2 = store2record1.append(vec![1, 2, 3]); + + let store3record1 = test_record(); + + let store4record1 = test_record(); + + // index1 only knows about the first two entries of the first two stores + index1.set(store1record1); + index1.set(store2record1); + + // index2 is fully up to date with the first two stores, and knows of a third + index2.set(store1record2); + index2.set(store2record2); + index2.set(store3record1); + + // index1 knows of a 4th store + index1.set(store4record1); + + let diff1 = index1.diff(&index2); + let diff2 = index2.diff(&index1); + + // both diffs the same length + assert_eq!(4, diff1.len()); + assert_eq!(4, diff2.len()); + + dbg!(&diff1, &diff2); + + // both diffs should be ALMOST the same. They will agree on which hosts and tags + // require updating, but the "other" value will not be the same. + let smol_diff_1: Vec<(HostId, String)> = + diff1.iter().map(|v| (v.host, v.tag.clone())).collect(); + let smol_diff_2: Vec<(HostId, String)> = + diff1.iter().map(|v| (v.host, v.tag.clone())).collect(); + + assert_eq!(smol_diff_1, smol_diff_2); + + // diffing with yourself = no diff + assert_eq!(index1.diff(&index1).len(), 0); + assert_eq!(index2.diff(&index2).len(), 0); + } +} diff --git a/crates/atuin-common/src/shell.rs b/crates/atuin-common/src/shell.rs new file mode 100644 index 00000000..7f9a7b8f --- /dev/null +++ b/crates/atuin-common/src/shell.rs @@ -0,0 +1,183 @@ +use std::{ffi::OsStr, path::Path, process::Command}; + +use serde::Serialize; +use sysinfo::{Process, System, get_current_pid}; +use thiserror::Error; + +#[derive(PartialEq)] +pub enum Shell { + Sh, + Bash, + Fish, + Zsh, + Xonsh, + Nu, + Powershell, + + Unknown, +} + +impl std::fmt::Display for Shell { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let shell = match self { + Shell::Bash => "bash", + Shell::Fish => "fish", + Shell::Zsh => "zsh", + Shell::Nu => "nu", + Shell::Xonsh => "xonsh", + Shell::Sh => "sh", + Shell::Powershell => "powershell", + + Shell::Unknown => "unknown", + }; + + write!(f, "{shell}") + } +} + +#[derive(Debug, Error, Serialize)] +pub enum ShellError { + #[error("shell not supported")] + NotSupported, + + #[error("failed to execute shell command: {0}")] + ExecError(String), +} + +impl Shell { + pub fn current() -> Shell { + let sys = System::new_all(); + + let process = sys + .process(get_current_pid().expect("Failed to get current PID")) + .expect("Process with current pid does not exist"); + + let parent = sys + .process(process.parent().expect("Atuin running with no parent!")) + .expect("Process with parent pid does not exist"); + + let shell = parent.name().trim().to_lowercase(); + let shell = shell.strip_prefix('-').unwrap_or(&shell); + + Shell::from_string(shell.to_string()) + } + + pub fn from_env() -> Shell { + std::env::var("ATUIN_SHELL").map_or(Shell::Unknown, |shell| { + Shell::from_string(shell.trim().to_lowercase()) + }) + } + + pub fn config_file(&self) -> Option<std::path::PathBuf> { + let mut path = if let Some(base) = directories::BaseDirs::new() { + base.home_dir().to_owned() + } else { + return None; + }; + + // TODO: handle all shells + match self { + Shell::Bash => path.push(".bashrc"), + Shell::Zsh => path.push(".zshrc"), + Shell::Fish => path.push(".config/fish/config.fish"), + + _ => return None, + }; + + Some(path) + } + + /// Best-effort attempt to determine the default shell + /// This implementation will be different across different platforms + /// Caller should ensure to handle Shell::Unknown correctly + pub fn default_shell() -> Result<Shell, ShellError> { + let sys = System::name().unwrap_or("".to_string()).to_lowercase(); + + // TODO: Support Linux + // I'm pretty sure we can use /etc/passwd there, though there will probably be some issues + let path = if sys.contains("darwin") { + // This works in my testing so far + Shell::Sh.run_interactive([ + "dscl localhost -read \"/Local/Default/Users/$USER\" shell | awk '{print $2}'", + ])? + } else if cfg!(windows) { + return Ok(Shell::Powershell); + } else { + Shell::Sh.run_interactive(["getent passwd $LOGNAME | cut -d: -f7"])? + }; + + let path = Path::new(path.trim()); + let shell = path.file_name(); + + if shell.is_none() { + return Err(ShellError::NotSupported); + } + + Ok(Shell::from_string( + shell.unwrap().to_string_lossy().to_string(), + )) + } + + pub fn from_string(name: String) -> Shell { + match name.as_str() { + "bash" => Shell::Bash, + "fish" => Shell::Fish, + "zsh" => Shell::Zsh, + "xonsh" => Shell::Xonsh, + "nu" => Shell::Nu, + "sh" => Shell::Sh, + "powershell" => Shell::Powershell, + + _ => Shell::Unknown, + } + } + + /// Returns true if the shell is posix-like + /// Note that while fish is not posix compliant, it behaves well enough for our current + /// featureset that this does not matter. + pub fn is_posixish(&self) -> bool { + matches!(self, Shell::Bash | Shell::Fish | Shell::Zsh) + } + + pub fn run_interactive<I, S>(&self, args: I) -> Result<String, ShellError> + where + I: IntoIterator<Item = S>, + S: AsRef<OsStr>, + { + let shell = self.to_string(); + let output = if self == &Self::Powershell { + Command::new(shell) + .args(args) + .output() + .map_err(|e| ShellError::ExecError(e.to_string()))? + } else { + Command::new(shell) + .arg("-ic") + .args(args) + .output() + .map_err(|e| ShellError::ExecError(e.to_string()))? + }; + + Ok(String::from_utf8(output.stdout).unwrap()) + } +} + +pub fn shell_name(parent: Option<&Process>) -> String { + let sys = System::new_all(); + + let parent = if let Some(parent) = parent { + parent + } else { + let process = sys + .process(get_current_pid().expect("Failed to get current PID")) + .expect("Process with current pid does not exist"); + + sys.process(process.parent().expect("Atuin running with no parent!")) + .expect("Process with parent pid does not exist") + }; + + let shell = parent.name().trim().to_lowercase(); + let shell = shell.strip_prefix('-').unwrap_or(&shell); + + shell.to_string() +} diff --git a/crates/atuin-common/src/tls.rs b/crates/atuin-common/src/tls.rs new file mode 100644 index 00000000..e8c840e0 --- /dev/null +++ b/crates/atuin-common/src/tls.rs @@ -0,0 +1,15 @@ +use std::sync::Once; + +static INIT: Once = Once::new(); + +/// Ensure the rustls crypto provider (ring) is installed. +/// +/// Must be called before creating any reqwest clients. Safe to call +/// multiple times — only the first call installs the provider. +pub fn ensure_crypto_provider() { + INIT.call_once(|| { + rustls::crypto::ring::default_provider() + .install_default() + .expect("Failed to install rustls crypto provider"); + }); +} diff --git a/crates/atuin-common/src/utils.rs b/crates/atuin-common/src/utils.rs new file mode 100644 index 00000000..b885423e --- /dev/null +++ b/crates/atuin-common/src/utils.rs @@ -0,0 +1,301 @@ +use std::borrow::Cow; +use std::env; +use std::path::PathBuf; + +use eyre::{Result, eyre}; + +use base64::prelude::{BASE64_URL_SAFE_NO_PAD, Engine}; +use getrandom::getrandom; +use uuid::Uuid; + +/// Generate N random bytes, using a cryptographically secure source +pub fn crypto_random_bytes<const N: usize>() -> [u8; N] { + // rand say they are in principle safe for crypto purposes, but that it is perhaps a better + // idea to use getrandom for things such as passwords. + let mut ret = [0u8; N]; + + getrandom(&mut ret).expect("Failed to generate random bytes!"); + + ret +} + +/// Generate N random bytes using a cryptographically secure source, return encoded as a string +pub fn crypto_random_string<const N: usize>() -> String { + let bytes = crypto_random_bytes::<N>(); + + // We only use this to create a random string, and won't be reversing it to find the original + // data - no padding is OK there. It may be in URLs. + BASE64_URL_SAFE_NO_PAD.encode(bytes) +} + +pub fn uuid_v7() -> Uuid { + Uuid::now_v7() +} + +pub fn uuid_v4() -> String { + Uuid::new_v4().as_simple().to_string() +} + +pub fn has_git_dir(path: &str) -> bool { + let mut gitdir = PathBuf::from(path); + gitdir.push(".git"); + + gitdir.exists() +} + +// detect if any parent dir has a git repo in it +// I really don't want to bring in libgit for something simple like this +// If we start to do anything more advanced, then perhaps +pub fn in_git_repo(path: &str) -> Option<PathBuf> { + let mut gitdir = PathBuf::from(path); + + while gitdir.parent().is_some() && !has_git_dir(gitdir.to_str().unwrap()) { + gitdir.pop(); + } + + // No parent? then we hit root, finding no git + if gitdir.parent().is_some() { + return Some(gitdir); + } + + None +} + +// TODO: more reliable, more tested +// I don't want to use ProjectDirs, it puts config in awkward places on +// mac. Data too. Seems to be more intended for GUI apps. + +pub fn home_dir() -> PathBuf { + directories::BaseDirs::new() + .map(|d| d.home_dir().to_path_buf()) + .expect("could not determine home directory") +} + +pub fn config_dir() -> PathBuf { + let config_dir = + std::env::var("XDG_CONFIG_HOME").map_or_else(|_| home_dir().join(".config"), PathBuf::from); + config_dir.join("atuin") +} + +pub fn data_dir() -> PathBuf { + let data_dir = std::env::var("XDG_DATA_HOME") + .map_or_else(|_| home_dir().join(".local").join("share"), PathBuf::from); + + data_dir.join("atuin") +} + +pub fn runtime_dir() -> PathBuf { + std::env::var("XDG_RUNTIME_DIR").map_or_else(|_| data_dir(), PathBuf::from) +} + +pub fn logs_dir() -> PathBuf { + home_dir().join(".atuin").join("logs") +} + +pub fn dotfiles_cache_dir() -> PathBuf { + // In most cases, this will be ~/.local/share/atuin/dotfiles/cache + let data_dir = std::env::var("XDG_DATA_HOME") + .map_or_else(|_| home_dir().join(".local").join("share"), PathBuf::from); + + data_dir.join("atuin").join("dotfiles").join("cache") +} + +pub fn get_current_dir() -> String { + // Prefer PWD environment variable over cwd if available to better support symbolic links + match env::var("PWD") { + Ok(v) => v, + Err(_) => match env::current_dir() { + Ok(dir) => dir.display().to_string(), + Err(_) => String::from(""), + }, + } +} + +pub fn broken_symlink<P: Into<PathBuf>>(path: P) -> bool { + let path = path.into(); + path.is_symlink() && !path.exists() +} + +/// Extension trait for anything that can behave like a string to make it easy to escape control +/// characters. +/// +/// Intended to help prevent control characters being printed and interpreted by the terminal when +/// printing history as well as to ensure the commands that appear in the interactive search +/// reflect the actual command run rather than just the printable characters. +pub trait Escapable: AsRef<str> { + fn escape_control(&self) -> Cow<'_, str> { + if !self.as_ref().contains(|c: char| c.is_ascii_control()) { + self.as_ref().into() + } else { + let mut remaining = self.as_ref(); + // Not a perfect way to reserve space but should reduce the allocations + let mut buf = String::with_capacity(remaining.len()); + while let Some(i) = remaining.find(|c: char| c.is_ascii_control()) { + // safe to index with `..i`, `i` and `i+1..` as part[i] is a single byte ascii char + buf.push_str(&remaining[..i]); + buf.push('^'); + buf.push(match remaining.as_bytes()[i] { + 0x7F => '?', + code => char::from_u32(u32::from(code) + 64).unwrap(), + }); + remaining = &remaining[i + 1..]; + } + buf.push_str(remaining); + buf.into() + } + } +} + +pub fn unquote(s: &str) -> Result<String> { + if s.chars().count() < 2 { + return Err(eyre!("not enough chars")); + } + + let quote = s.chars().next().unwrap(); + + // not quoted, do nothing + if quote != '"' && quote != '\'' && quote != '`' { + return Ok(s.to_string()); + } + + if s.chars().last().unwrap() != quote { + return Err(eyre!("unexpected eof, quotes do not match")); + } + + // removes quote characters + // the sanity checks performed above ensure that the quotes will be ASCII and this will not + // panic + let s = &s[1..s.len() - 1]; + + Ok(s.to_string()) +} + +impl<T: AsRef<str>> Escapable for T {} + +#[allow(unsafe_code)] +#[cfg(test)] +mod tests { + use pretty_assertions::assert_ne; + + use super::*; + + use std::collections::HashSet; + + #[cfg(not(windows))] + #[test] + fn test_dirs() { + // these tests need to be run sequentially to prevent race condition + test_config_dir_xdg(); + test_config_dir(); + test_data_dir_xdg(); + test_data_dir(); + } + + #[cfg(not(windows))] + fn test_config_dir_xdg() { + // TODO: Audit that the environment access only happens in single-threaded code. + unsafe { env::remove_var("HOME") }; + // TODO: Audit that the environment access only happens in single-threaded code. + unsafe { env::set_var("XDG_CONFIG_HOME", "/home/user/custom_config") }; + assert_eq!( + config_dir(), + PathBuf::from("/home/user/custom_config/atuin") + ); + // TODO: Audit that the environment access only happens in single-threaded code. + unsafe { env::remove_var("XDG_CONFIG_HOME") }; + } + + #[cfg(not(windows))] + fn test_config_dir() { + // TODO: Audit that the environment access only happens in single-threaded code. + unsafe { env::set_var("HOME", "/home/user") }; + // TODO: Audit that the environment access only happens in single-threaded code. + unsafe { env::remove_var("XDG_CONFIG_HOME") }; + + assert_eq!(config_dir(), PathBuf::from("/home/user/.config/atuin")); + + // TODO: Audit that the environment access only happens in single-threaded code. + unsafe { env::remove_var("HOME") }; + } + + #[cfg(not(windows))] + fn test_data_dir_xdg() { + // TODO: Audit that the environment access only happens in single-threaded code. + unsafe { env::remove_var("HOME") }; + // TODO: Audit that the environment access only happens in single-threaded code. + unsafe { env::set_var("XDG_DATA_HOME", "/home/user/custom_data") }; + assert_eq!(data_dir(), PathBuf::from("/home/user/custom_data/atuin")); + // TODO: Audit that the environment access only happens in single-threaded code. + unsafe { env::remove_var("XDG_DATA_HOME") }; + } + + #[cfg(not(windows))] + fn test_data_dir() { + // TODO: Audit that the environment access only happens in single-threaded code. + unsafe { env::set_var("HOME", "/home/user") }; + // TODO: Audit that the environment access only happens in single-threaded code. + unsafe { env::remove_var("XDG_DATA_HOME") }; + assert_eq!(data_dir(), PathBuf::from("/home/user/.local/share/atuin")); + // TODO: Audit that the environment access only happens in single-threaded code. + unsafe { env::remove_var("HOME") }; + } + + #[test] + fn uuid_is_unique() { + let how_many: usize = 1000000; + + // for peace of mind + let mut uuids: HashSet<Uuid> = HashSet::with_capacity(how_many); + + // there will be many in the same millisecond + for _ in 0..how_many { + let uuid = uuid_v7(); + uuids.insert(uuid); + } + + assert_eq!(uuids.len(), how_many); + } + + #[test] + fn escape_control_characters() { + use super::Escapable; + // CSI colour sequence + assert_eq!("\x1b[31mfoo".escape_control(), "^[[31mfoo"); + + // Tabs count as control chars + assert_eq!("foo\tbar".escape_control(), "foo^Ibar"); + + // space is in control char range but should be excluded + assert_eq!("two words".escape_control(), "two words"); + + // unicode multi-byte characters + let s = "🐢\x1b[32m🦀"; + assert_eq!(s.escape_control(), s.replace("\x1b", "^[")); + } + + #[test] + fn escape_no_control_characters() { + use super::Escapable as _; + assert!(matches!( + "no control characters".escape_control(), + Cow::Borrowed(_) + )); + assert!(matches!( + "with \x1b[31mcontrol\x1b[0m characters".escape_control(), + Cow::Owned(_) + )); + } + + #[test] + fn dumb_random_test() { + // Obviously not a test of randomness, but make sure we haven't made some + // catastrophic error + + assert_ne!(crypto_random_string::<1>(), crypto_random_string::<1>()); + assert_ne!(crypto_random_string::<2>(), crypto_random_string::<2>()); + assert_ne!(crypto_random_string::<4>(), crypto_random_string::<4>()); + assert_ne!(crypto_random_string::<8>(), crypto_random_string::<8>()); + assert_ne!(crypto_random_string::<16>(), crypto_random_string::<16>()); + assert_ne!(crypto_random_string::<32>(), crypto_random_string::<32>()); + } +} diff --git a/crates/atuin-daemon/Cargo.toml b/crates/atuin-daemon/Cargo.toml new file mode 100644 index 00000000..b9d1d8fd --- /dev/null +++ b/crates/atuin-daemon/Cargo.toml @@ -0,0 +1,53 @@ +[package] +name = "atuin-daemon" +edition = "2024" +version = { workspace = true } +description = "The daemon crate for Atuin" + +authors.workspace = true +rust-version.workspace = true +license.workspace = true +homepage.workspace = true +repository.workspace = true +readme.workspace = true + +# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html + +[dependencies] +atuin-client = { path = "../atuin-client", version = "18.13.2" } +atuin-common = { path = "../atuin-common", version = "18.13.2" } +atuin-dotfiles = { path = "../atuin-dotfiles", version = "18.13.2" } +atuin-history = { path = "../atuin-history", version = "18.13.2" } + +time = { workspace = true } +uuid = { workspace = true } +tokio = { workspace = true } +tower = { workspace = true } +eyre = { workspace = true } +tracing = { workspace = true } +tracing-subscriber = { workspace = true } + +dashmap = "6.1.0" +lasso = { version = "0.7", features = ["multi-threaded"] } +tonic-types = "0.14" +tonic = "0.14" +tonic-prost = "0.14" +prost = "0.14" +prost-types = "0.14" +tokio-stream = { version = "0.1.14", features = ["net"] } +hyper-util = "0.1" + +rand.workspace = true +nucleo = { git = "https://github.com/atuinsh/nucleo-ext.git", rev="74bd786" } + + +[target.'cfg(target_os = "linux")'.dependencies] +listenfd = "1.0.1" + +[dev-dependencies] +tempfile = { workspace = true } + +[build-dependencies] +protox = "0.9" +tonic-build = "0.14" +tonic-prost-build = "0.14" diff --git a/crates/atuin-daemon/build.rs b/crates/atuin-daemon/build.rs new file mode 100644 index 00000000..7034aa04 --- /dev/null +++ b/crates/atuin-daemon/build.rs @@ -0,0 +1,24 @@ +use std::{env, fs, path::PathBuf}; + +use protox::prost::Message; + +fn main() -> std::io::Result<()> { + let proto_paths = [ + "proto/history.proto", + "proto/search.proto", + "proto/control.proto", + ]; + let proto_include_dirs = ["proto"]; + + let file_descriptors = protox::compile(proto_paths, proto_include_dirs).unwrap(); + + let file_descriptor_path = PathBuf::from(env::var_os("OUT_DIR").expect("OUT_DIR not set")) + .join("file_descriptor_set.bin"); + fs::write(&file_descriptor_path, file_descriptors.encode_to_vec()).unwrap(); + + tonic_prost_build::configure() + .build_server(true) + .file_descriptor_set_path(&file_descriptor_path) + .skip_protoc_run() + .compile_protos(&proto_paths, &proto_include_dirs) +} diff --git a/crates/atuin-daemon/proto/control.proto b/crates/atuin-daemon/proto/control.proto new file mode 100644 index 00000000..06347902 --- /dev/null +++ b/crates/atuin-daemon/proto/control.proto @@ -0,0 +1,62 @@ +syntax = "proto3"; +package control; + +// The Control service allows external processes (CLI commands, etc.) +// to inject events into the running daemon. +service Control { + // Send an event to the daemon's event bus + rpc SendEvent(SendEventRequest) returns (SendEventResponse); +} + +message SendEventRequest { + oneof event { + // History was pruned - search index needs full rebuild + HistoryPrunedEvent history_pruned = 1; + + // Specific history items were deleted + HistoryDeletedEvent history_deleted = 2; + + // Request immediate sync + ForceSyncEvent force_sync = 3; + + // Settings have changed, reload if needed + SettingsReloadedEvent settings_reloaded = 4; + + // Request graceful shutdown + ShutdownEvent shutdown = 5; + + // History was rebuilt - search index needs full rebuild + HistoryRebuiltEvent history_rebuilt = 6; + } +} + +message SendEventResponse { + // Empty on success; errors come through gRPC status +} + +// Individual event message types + +message HistoryPrunedEvent { + // No fields needed - just signals that pruning happened +} + +message HistoryRebuiltEvent { + // No fields needed - just signals that rebuilding happened +} + +message HistoryDeletedEvent { + // IDs of deleted history items (UUIDs as strings) + repeated string ids = 1; +} + +message ForceSyncEvent { + // No fields needed - just triggers sync +} + +message SettingsReloadedEvent { + // No fields needed - components should re-read settings +} + +message ShutdownEvent { + // No fields needed - triggers graceful shutdown +} diff --git a/crates/atuin-daemon/proto/history.proto b/crates/atuin-daemon/proto/history.proto new file mode 100644 index 00000000..2a45b7cf --- /dev/null +++ b/crates/atuin-daemon/proto/history.proto @@ -0,0 +1,54 @@ +syntax = "proto3"; +package history; + +message StartHistoryRequest { + // If people are still using my software in ~530 years, they can figure out a u128 migration + uint64 timestamp = 1; // nanosecond unix epoch + string command = 2; + string cwd = 3; + string session = 4; + string hostname = 5; + string author = 6; + string intent = 7; +} + +message EndHistoryRequest { + string id = 1; + int64 exit = 2; + uint64 duration = 3; +} + +message StartHistoryReply { + string id = 1; + string version = 2; + uint32 protocol = 3; +} + +message EndHistoryReply { + string id = 1; + uint64 idx = 2; + string version = 3; + uint32 protocol = 4; +} + +message StatusRequest {} + +message StatusReply { + bool healthy = 1; + string version = 2; + uint32 pid = 3; + uint32 protocol = 4; +} + +message ShutdownRequest {} + +message ShutdownReply { + bool accepted = 1; +} + +service History { + rpc StartHistory(StartHistoryRequest) returns (StartHistoryReply); + rpc EndHistory(EndHistoryRequest) returns (EndHistoryReply); + rpc Status(StatusRequest) returns (StatusReply); + rpc Shutdown(ShutdownRequest) returns (ShutdownReply); +} diff --git a/crates/atuin-daemon/proto/search.proto b/crates/atuin-daemon/proto/search.proto new file mode 100644 index 00000000..6b84acbd --- /dev/null +++ b/crates/atuin-daemon/proto/search.proto @@ -0,0 +1,35 @@ +syntax = "proto3"; +package search; + +enum FilterMode { + GLOBAL = 0; + HOST = 1; + SESSION = 2; + DIRECTORY = 3; + WORKSPACE = 4; + SESSION_PRELOAD = 5; +} + +message SearchContext { + string session_id = 1; + string cwd = 2; + string hostname = 3; + string host_id = 4; + optional string git_root = 5; +} + +message SearchRequest { + string query = 1; + uint64 query_id = 2; // Incrementing ID to match responses to queries + FilterMode filter_mode = 3; + SearchContext context = 4; +} + +message SearchResponse { + uint64 query_id = 1; // Echo back the query ID + repeated bytes ids = 2; +} + +service Search { + rpc Search(stream SearchRequest) returns (stream SearchResponse); +} diff --git a/crates/atuin-daemon/src/client.rs b/crates/atuin-daemon/src/client.rs new file mode 100644 index 00000000..2f492f6b --- /dev/null +++ b/crates/atuin-daemon/src/client.rs @@ -0,0 +1,419 @@ +use atuin_client::database::Context; +use atuin_client::settings::{FilterMode, Settings}; +use eyre::{Context as EyreContext, Result}; +#[cfg(windows)] +use tokio::net::TcpStream; +use tonic::Code; +use tonic::transport::{Channel, Endpoint, Uri}; +use tower::service_fn; + +use hyper_util::rt::TokioIo; + +#[cfg(unix)] +use tokio::net::UnixStream; + +use atuin_client::history::History; +use tracing::{Level, instrument, span}; + +use crate::control::HistoryRebuiltEvent; +use crate::control::{ + ForceSyncEvent, HistoryDeletedEvent, HistoryPrunedEvent, SendEventRequest, + SettingsReloadedEvent, ShutdownEvent, control_client::ControlClient as ControlServiceClient, +}; +use crate::events::DaemonEvent; +use crate::history::{ + EndHistoryReply, EndHistoryRequest, ShutdownRequest, StartHistoryReply, StartHistoryRequest, + StatusReply, StatusRequest, history_client::HistoryClient as HistoryServiceClient, +}; +use crate::search::{ + FilterMode as RpcFilterMode, SearchContext as RpcSearchContext, SearchRequest, SearchResponse, + search_client::SearchClient as SearchServiceClient, +}; + +pub struct HistoryClient { + client: HistoryServiceClient<Channel>, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum DaemonClientErrorKind { + Connect, + Unavailable, + Unimplemented, + Other, +} + +#[must_use] +pub fn classify_error(error: &eyre::Report) -> DaemonClientErrorKind { + for cause in error.chain() { + if cause.downcast_ref::<tonic::transport::Error>().is_some() { + return DaemonClientErrorKind::Connect; + } + + if let Some(status) = cause.downcast_ref::<tonic::Status>() { + return match status.code() { + Code::Unavailable => DaemonClientErrorKind::Unavailable, + Code::Unimplemented => DaemonClientErrorKind::Unimplemented, + _ => DaemonClientErrorKind::Other, + }; + } + } + + DaemonClientErrorKind::Other +} + +// Wrap the grpc client +impl HistoryClient { + #[cfg(unix)] + pub async fn new(path: String) -> Result<Self> { + use eyre::Context; + + let log_path = path.clone(); + let channel = Endpoint::try_from("http://atuin_local_daemon:0")? + .connect_with_connector(service_fn(move |_: Uri| { + let path = path.clone(); + + async move { + Ok::<_, std::io::Error>(TokioIo::new(UnixStream::connect(path.clone()).await?)) + } + })) + .await + .wrap_err_with(|| { + format!( + "failed to connect to local atuin daemon at {}. Is it running?", + &log_path + ) + })?; + + let client = HistoryServiceClient::new(channel); + + Ok(HistoryClient { client }) + } + + #[cfg(not(unix))] + pub async fn new(port: u64) -> Result<Self> { + let channel = Endpoint::try_from("http://atuin_local_daemon:0")? + .connect_with_connector(service_fn(move |_: Uri| { + let url = format!("127.0.0.1:{port}"); + + async move { + Ok::<_, std::io::Error>(TokioIo::new(TcpStream::connect(url.clone()).await?)) + } + })) + .await + .wrap_err_with(|| { + format!( + "failed to connect to local atuin daemon at 127.0.0.1:{port}. Is it running?" + ) + })?; + + let client = HistoryServiceClient::new(channel); + + Ok(HistoryClient { client }) + } + + pub async fn start_history(&mut self, h: History) -> Result<StartHistoryReply> { + let req = StartHistoryRequest { + command: h.command, + cwd: h.cwd, + hostname: h.hostname, + session: h.session, + timestamp: h.timestamp.unix_timestamp_nanos() as u64, + author: h.author, + intent: h.intent.unwrap_or_default(), + }; + + Ok(self.client.start_history(req).await?.into_inner()) + } + + pub async fn end_history( + &mut self, + id: String, + duration: u64, + exit: i64, + ) -> Result<EndHistoryReply> { + let req = EndHistoryRequest { id, duration, exit }; + + Ok(self.client.end_history(req).await?.into_inner()) + } + + pub async fn status(&mut self) -> Result<StatusReply> { + Ok(self.client.status(StatusRequest {}).await?.into_inner()) + } + + pub async fn shutdown(&mut self) -> Result<bool> { + let resp = self.client.shutdown(ShutdownRequest {}).await?.into_inner(); + Ok(resp.accepted) + } +} + +pub struct SearchClient { + client: SearchServiceClient<Channel>, +} + +impl SearchClient { + #[cfg(unix)] + pub async fn new(path: String) -> Result<Self> { + let log_path = path.clone(); + let channel = Endpoint::try_from("http://atuin_local_daemon:0")? + .connect_with_connector(service_fn(move |_: Uri| { + let path = path.clone(); + + async move { + Ok::<_, std::io::Error>(TokioIo::new(UnixStream::connect(path.clone()).await?)) + } + })) + .await + .wrap_err_with(|| { + format!( + "failed to connect to local atuin daemon at {}. Is it running?", + &log_path + ) + })?; + + let client = SearchServiceClient::new(channel); + + Ok(SearchClient { client }) + } + + #[cfg(not(unix))] + pub async fn new(port: u64) -> Result<Self> { + let channel = Endpoint::try_from("http://atuin_local_daemon:0")? + .connect_with_connector(service_fn(move |_: Uri| { + let url = format!("127.0.0.1:{port}"); + + async move { + Ok::<_, std::io::Error>(TokioIo::new(TcpStream::connect(url.clone()).await?)) + } + })) + .await + .wrap_err_with(|| { + format!( + "failed to connect to local atuin daemon at 127.0.0.1:{port}. Is it running?" + ) + })?; + + let client = SearchServiceClient::new(channel); + + Ok(SearchClient { client }) + } + + #[instrument(skip_all, level = Level::TRACE, name = "daemon_client_search", fields(query = %query, query_id = query_id))] + pub async fn search( + &mut self, + query: String, + query_id: u64, + filter_mode: FilterMode, + context: Option<Context>, + ) -> Result<tonic::Streaming<SearchResponse>> { + let request = SearchRequest { + query, + query_id, + filter_mode: RpcFilterMode::from(filter_mode).into(), + context: context.map(RpcSearchContext::from), + }; + let request_stream = tokio_stream::once(request); + let response = span!(Level::TRACE, "daemon_client_search.request") + .in_scope(async || self.client.search(request_stream).await) + .await?; + + Ok(response.into_inner()) + } +} + +impl From<FilterMode> for RpcFilterMode { + fn from(filter_mode: FilterMode) -> Self { + match filter_mode { + FilterMode::Global => RpcFilterMode::Global, + FilterMode::Host => RpcFilterMode::Host, + FilterMode::Session => RpcFilterMode::Session, + FilterMode::Directory => RpcFilterMode::Directory, + FilterMode::Workspace => RpcFilterMode::Workspace, + FilterMode::SessionPreload => RpcFilterMode::SessionPreload, + } + } +} + +impl From<Context> for RpcSearchContext { + fn from(context: Context) -> Self { + RpcSearchContext { + session_id: context.session, + cwd: context.cwd, + hostname: context.hostname, + host_id: context.host_id, + git_root: context + .git_root + .map(|path| path.to_string_lossy().to_string()), + } + } +} + +// ============================================================================ +// Control Client +// ============================================================================ + +/// Client for the Control gRPC service. +/// +/// Used to inject events into a running daemon from external processes. +pub struct ControlClient { + client: ControlServiceClient<Channel>, +} + +impl ControlClient { + /// Connect to the daemon's control service. + #[cfg(unix)] + pub async fn new(path: String) -> Result<Self> { + let log_path = path.clone(); + let channel = Endpoint::try_from("http://atuin_local_daemon:0")? + .connect_with_connector(service_fn(move |_: Uri| { + let path = path.clone(); + + async move { + Ok::<_, std::io::Error>(TokioIo::new(UnixStream::connect(path.clone()).await?)) + } + })) + .await + .wrap_err_with(|| { + format!( + "failed to connect to local atuin daemon at {}. Is it running?", + &log_path + ) + })?; + + let client = ControlServiceClient::new(channel); + + Ok(ControlClient { client }) + } + + /// Connect to the daemon's control service. + #[cfg(not(unix))] + pub async fn new(port: u64) -> Result<Self> { + let channel = Endpoint::try_from("http://atuin_local_daemon:0")? + .connect_with_connector(service_fn(move |_: Uri| { + let url = format!("127.0.0.1:{port}"); + + async move { + Ok::<_, std::io::Error>(TokioIo::new(TcpStream::connect(url.clone()).await?)) + } + })) + .await + .wrap_err_with(|| { + format!( + "failed to connect to local atuin daemon at 127.0.0.1:{port}. Is it running?" + ) + })?; + + let client = ControlServiceClient::new(channel); + + Ok(ControlClient { client }) + } + + /// Connect using settings. + #[cfg(unix)] + pub async fn from_settings(settings: &Settings) -> Result<Self> { + Self::new(settings.daemon.socket_path.clone()).await + } + + /// Connect using settings. + #[cfg(not(unix))] + pub async fn from_settings(settings: &Settings) -> Result<Self> { + Self::new(settings.daemon.tcp_port).await + } + + /// Send an event to the daemon. + pub async fn send_event(&mut self, event: DaemonEvent) -> Result<()> { + let proto_event = daemon_event_to_proto(event); + let request = SendEventRequest { + event: Some(proto_event), + }; + self.client.send_event(request).await?; + Ok(()) + } +} + +/// Convert a daemon event to its proto representation. +fn daemon_event_to_proto(event: DaemonEvent) -> crate::control::send_event_request::Event { + use crate::control::send_event_request::Event; + + match event { + DaemonEvent::HistoryPruned => Event::HistoryPruned(HistoryPrunedEvent {}), + DaemonEvent::HistoryRebuilt => Event::HistoryRebuilt(HistoryRebuiltEvent {}), + DaemonEvent::HistoryDeleted { ids } => Event::HistoryDeleted(HistoryDeletedEvent { + ids: ids.into_iter().map(|id| id.0).collect(), + }), + DaemonEvent::ForceSync => Event::ForceSync(ForceSyncEvent {}), + DaemonEvent::SettingsReloaded => Event::SettingsReloaded(SettingsReloadedEvent {}), + DaemonEvent::ShutdownRequested => Event::Shutdown(ShutdownEvent {}), + // These events are internal and not sent via the control service + DaemonEvent::HistoryStarted(_) + | DaemonEvent::HistoryEnded(_) + | DaemonEvent::RecordsAdded(_) + | DaemonEvent::SyncCompleted { .. } + | DaemonEvent::SyncFailed { .. } => { + // Use shutdown as a fallback, though this shouldn't happen + tracing::warn!("attempted to send internal event via control service"); + Event::Shutdown(ShutdownEvent {}) + } + } +} + +// ============================================================================ +// Convenience Functions +// ============================================================================ + +/// Emit an event to the daemon. +/// +/// This is a fire-and-forget helper for sending events to the daemon from +/// external processes like CLI commands. If the daemon isn't running, this +/// will silently succeed (returns Ok). +/// +/// # Example +/// +/// ```ignore +/// // After pruning history +/// emit_event(DaemonEvent::HistoryPruned).await?; +/// +/// // After deleting specific history items +/// emit_event(DaemonEvent::HistoryDeleted { ids: vec![...] }).await?; +/// +/// // Request immediate sync +/// emit_event(DaemonEvent::ForceSync).await?; +/// ``` +pub async fn emit_event(event: DaemonEvent) -> Result<()> { + emit_event_with_settings(event, None).await +} + +/// Emit an event to the daemon with explicit settings. +/// +/// If settings are not provided, they will be loaded from the default location. +/// If the daemon isn't running, this will silently succeed. +pub async fn emit_event_with_settings( + event: DaemonEvent, + settings: Option<&Settings>, +) -> Result<()> { + // Load settings if not provided + let owned_settings; + let settings = match settings { + Some(s) => s, + None => { + owned_settings = Settings::new()?; + &owned_settings + } + }; + + // Try to connect - if daemon isn't running, that's fine + let mut client = match ControlClient::from_settings(settings).await { + Ok(c) => c, + Err(e) => { + tracing::debug!(?e, "daemon not running, skipping event emission"); + return Ok(()); + } + }; + + // Send the event + if let Err(e) = client.send_event(event).await { + tracing::debug!(?e, "failed to send event to daemon"); + // Don't fail - this is fire-and-forget + } + + Ok(()) +} diff --git a/crates/atuin-daemon/src/components/history.rs b/crates/atuin-daemon/src/components/history.rs new file mode 100644 index 00000000..23d48c5e --- /dev/null +++ b/crates/atuin-daemon/src/components/history.rs @@ -0,0 +1,252 @@ +//! History component. +//! +//! Handles command history lifecycle (start/end) and provides the History gRPC service. + +use std::sync::Arc; + +use atuin_client::{ + database::Database, + history::{History, HistoryId, store::HistoryStore}, + settings::Settings, +}; +use dashmap::DashMap; +use eyre::Result; +use time::OffsetDateTime; +use tonic::{Request, Response, Status}; +use tracing::{Level, instrument}; + +use crate::{ + daemon::{Component, DaemonHandle}, + events::DaemonEvent, + history::{ + EndHistoryReply, EndHistoryRequest, ShutdownReply, ShutdownRequest, StartHistoryReply, + StartHistoryRequest, StatusReply, StatusRequest, + history_server::{History as HistorySvc, HistoryServer}, + }, +}; + +const DAEMON_PROTOCOL_VERSION: u32 = 1; + +/// History component - manages command history lifecycle. +/// +/// This component: +/// - Tracks currently running commands (stored in memory) +/// - Saves completed commands to the database and record store +/// - Emits history events for other components (e.g., search indexing) +/// - Provides the History gRPC service +pub struct HistoryComponent { + inner: Arc<HistoryComponentInner>, +} + +struct HistoryComponentInner { + /// Commands currently running (not yet completed). + running: DashMap<HistoryId, History>, + + /// Handle to the daemon (set during start). + handle: tokio::sync::RwLock<Option<DaemonHandle>>, + + /// History store for pushing records (set during start). + history_store: tokio::sync::RwLock<Option<HistoryStore>>, +} + +impl HistoryComponent { + /// Create a new history component. + pub fn new() -> Self { + Self { + inner: Arc::new(HistoryComponentInner { + running: DashMap::new(), + handle: tokio::sync::RwLock::new(None), + history_store: tokio::sync::RwLock::new(None), + }), + } + } + + /// Get the gRPC service for this component. + /// + /// This returns a tonic service that can be added to a gRPC server. + pub fn grpc_service(&self) -> HistoryServer<HistoryGrpcService> { + HistoryServer::new(HistoryGrpcService { + inner: self.inner.clone(), + }) + } +} + +impl Default for HistoryComponent { + fn default() -> Self { + Self::new() + } +} + +#[tonic::async_trait] +impl Component for HistoryComponent { + fn name(&self) -> &'static str { + "history" + } + + async fn start(&mut self, handle: DaemonHandle) -> Result<()> { + // Create the history store + let host_id = Settings::host_id().await?; + let history_store = + HistoryStore::new(handle.store().clone(), host_id, *handle.encryption_key()); + + *self.inner.history_store.write().await = Some(history_store); + *self.inner.handle.write().await = Some(handle); + + tracing::info!("history component started"); + Ok(()) + } + + async fn handle_event(&mut self, _event: &DaemonEvent) -> Result<()> { + // History component produces events but doesn't need to react to them + Ok(()) + } + + async fn stop(&mut self) -> Result<()> { + tracing::info!("history component stopped"); + Ok(()) + } +} + +/// The gRPC service implementation. +/// +/// This is a thin wrapper that delegates to the component's shared state. +pub struct HistoryGrpcService { + inner: Arc<HistoryComponentInner>, +} + +#[tonic::async_trait] +impl HistorySvc for HistoryGrpcService { + #[instrument(skip_all, level = Level::INFO)] + async fn start_history( + &self, + request: Request<StartHistoryRequest>, + ) -> Result<Response<StartHistoryReply>, Status> { + let req = request.into_inner(); + + let timestamp = + OffsetDateTime::from_unix_timestamp_nanos(req.timestamp as i128).map_err(|_| { + Status::invalid_argument( + "failed to parse timestamp as unix time (expected nanos since epoch)", + ) + })?; + + let h: History = History::daemon() + .timestamp(timestamp) + .command(req.command) + .cwd(req.cwd) + .session(req.session) + .hostname(req.hostname) + .build() + .into(); + + // Emit the event + if let Some(handle) = self.inner.handle.read().await.as_ref() { + handle.emit(DaemonEvent::HistoryStarted(h.clone())); + } + + let id = h.id.clone(); + tracing::info!(id = id.to_string(), "start history"); + self.inner.running.insert(id.clone(), h); + + let reply = StartHistoryReply { + id: id.to_string(), + version: env!("CARGO_PKG_VERSION").to_string(), + protocol: DAEMON_PROTOCOL_VERSION, + }; + + Ok(Response::new(reply)) + } + + #[instrument(skip_all, level = Level::INFO)] + async fn end_history( + &self, + request: Request<EndHistoryRequest>, + ) -> Result<Response<EndHistoryReply>, Status> { + let req = request.into_inner(); + let id = HistoryId(req.id); + + if let Some((_, mut history)) = self.inner.running.remove(&id) { + history.exit = req.exit; + history.duration = match req.duration { + 0 => i64::try_from( + (OffsetDateTime::now_utc() - history.timestamp).whole_nanoseconds(), + ) + .expect("failed to convert calculated duration to i64"), + value => i64::try_from(value).expect("failed to get i64 duration"), + }; + + // Get the handle and store to save the history + let handle_guard = self.inner.handle.read().await; + let handle = handle_guard + .as_ref() + .ok_or_else(|| Status::internal("component not initialized"))?; + + let store_guard = self.inner.history_store.read().await; + let history_store = store_guard + .as_ref() + .ok_or_else(|| Status::internal("component not initialized"))?; + + // Save to database + handle + .history_db() + .save(&history) + .await + .map_err(|e| Status::internal(format!("failed to write to db: {e:?}")))?; + + tracing::info!( + id = id.0.to_string(), + duration = history.duration, + "end history" + ); + + // Push to record store + let (record_id, idx) = history_store + .push(history.clone()) + .await + .map_err(|e| Status::internal(format!("failed to push record to store: {e:?}")))?; + + // Emit the event + handle.emit(DaemonEvent::HistoryEnded(history)); + + let reply = EndHistoryReply { + id: record_id.0.to_string(), + idx, + version: env!("CARGO_PKG_VERSION").to_string(), + protocol: DAEMON_PROTOCOL_VERSION, + }; + + return Ok(Response::new(reply)); + } + + Err(Status::not_found(format!( + "could not find history with id: {id}" + ))) + } + + #[instrument(skip_all, level = Level::INFO)] + async fn status( + &self, + _request: Request<StatusRequest>, + ) -> Result<Response<StatusReply>, Status> { + let reply = StatusReply { + healthy: true, + version: env!("CARGO_PKG_VERSION").to_string(), + pid: std::process::id(), + protocol: DAEMON_PROTOCOL_VERSION, + }; + + Ok(Response::new(reply)) + } + + #[instrument(skip_all, level = Level::INFO)] + async fn shutdown( + &self, + _request: Request<ShutdownRequest>, + ) -> Result<Response<ShutdownReply>, Status> { + // Use the daemon handle to request shutdown + if let Some(handle) = self.inner.handle.read().await.as_ref() { + handle.shutdown(); + } + Ok(Response::new(ShutdownReply { accepted: true })) + } +} diff --git a/crates/atuin-daemon/src/components/mod.rs b/crates/atuin-daemon/src/components/mod.rs new file mode 100644 index 00000000..5950d5d5 --- /dev/null +++ b/crates/atuin-daemon/src/components/mod.rs @@ -0,0 +1,22 @@ +//! Daemon components. +//! +//! Components are the building blocks of the daemon. Each component handles +//! a specific domain and can: +//! +//! - Expose gRPC services +//! - React to events +//! - Spawn background tasks +//! +//! Available components: +//! +//! - [`history::HistoryComponent`]: Command history lifecycle management +//! - [`search::SearchComponent`]: Fuzzy search over history +//! - [`sync::SyncComponent`]: Cloud sync + +pub mod history; +pub mod search; +pub mod sync; + +pub use history::HistoryComponent; +pub use search::SearchComponent; +pub use sync::SyncComponent; diff --git a/crates/atuin-daemon/src/components/search.rs b/crates/atuin-daemon/src/components/search.rs new file mode 100644 index 00000000..9fc87fae --- /dev/null +++ b/crates/atuin-daemon/src/components/search.rs @@ -0,0 +1,413 @@ +//! Search component. +//! +//! Provides fuzzy search over command history using the Nucleo search library +//! with frecency-based ranking and dynamic filtering. + +use std::{pin::Pin, sync::Arc}; + +use atuin_client::database::Database; +use eyre::Result; +use tokio::sync::RwLock; +use tokio_stream::Stream; +use tonic::{Request, Response, Status, Streaming}; +use tracing::{Level, debug, info, instrument, span, trace}; +use uuid::Uuid; + +use crate::{ + daemon::{Component, DaemonHandle}, + events::DaemonEvent, + search::{ + FilterMode, IndexFilterMode, QueryContext, SearchIndex, SearchRequest, SearchResponse, + search_server::{Search as SearchSvc, SearchServer}, + }, +}; + +const PAGE_SIZE: usize = 5000; +const RESULTS_LIMIT: u32 = 200; +/// How often to rebuild the frecency map (in seconds). +const FRECENCY_REFRESH_INTERVAL_SECS: u64 = 60; + +/// Search component - provides fuzzy search over command history. +/// +/// This component: +/// - Maintains a deduplicated search index with frecency ranking +/// - Loads history from the database on startup +/// - Updates the index when history events occur +/// - Provides the Search gRPC service +pub struct SearchComponent { + index: Arc<RwLock<SearchIndex>>, + handle: tokio::sync::RwLock<Option<DaemonHandle>>, + loader_handle: Option<tokio::task::JoinHandle<()>>, + frecency_handle: Option<tokio::task::JoinHandle<()>>, +} + +impl SearchComponent { + /// Create a new search component. + pub fn new() -> Self { + Self { + index: Arc::new(RwLock::new(SearchIndex::new())), + handle: tokio::sync::RwLock::new(None), + loader_handle: None, + frecency_handle: None, + } + } + + /// Get the gRPC service for this component. + pub fn grpc_service(&self) -> SearchServer<SearchGrpcService> { + SearchServer::new(SearchGrpcService { + index: self.index.clone(), + }) + } + + /// Rebuild the entire search index from the database. + async fn rebuild_index(&self) -> Result<()> { + let handle_guard = self.handle.read().await; + let handle = handle_guard + .as_ref() + .ok_or_else(|| eyre::eyre!("component not initialized"))?; + + info!("Rebuilding search index from database"); + + // Create a new index + let new_index = SearchIndex::new(); + + // Load all history into the new index + let db = handle.history_db().clone(); + let mut pager = db.all_paged(PAGE_SIZE, false, true); + loop { + match pager.next().await { + Ok(Some(histories)) => { + info!( + "Loading {} history entries into search index", + histories.len() + ); + new_index.add_histories(&histories); + } + Ok(None) => break, + Err(e) => { + tracing::error!("Failed to load history during rebuild: {}", e); + break; + } + } + } + + info!( + "Search index rebuild complete; {} unique commands", + new_index.command_count() + ); + + // Replace the old index with the new one + *self.index.write().await = new_index; + Ok(()) + } +} + +impl Default for SearchComponent { + fn default() -> Self { + Self::new() + } +} + +#[tonic::async_trait] +impl Component for SearchComponent { + fn name(&self) -> &'static str { + "search" + } + + async fn start(&mut self, handle: DaemonHandle) -> Result<()> { + *self.handle.write().await = Some(handle.clone()); + + // Spawn background task to load history into index + let index = self.index.clone(); + let db = handle.history_db().clone(); + let handle_for_loader = handle.clone(); + + self.loader_handle = Some(tokio::spawn(async move { + info!( + "Loading history into search index; page size = {}", + PAGE_SIZE + ); + let mut pager = db.all_paged(PAGE_SIZE, false, true); + loop { + match pager.next().await { + Ok(Some(histories)) => { + info!( + "Loading {} history entries into search index", + histories.len() + ); + index.read().await.add_histories(&histories); + } + Ok(None) => { + info!( + "Initial history load complete; {} unique commands indexed", + index.read().await.command_count() + ); + // Build initial frecency map with current settings + let settings = handle_for_loader.settings().await; + index.read().await.rebuild_frecency(&settings.search).await; + info!("Initial frecency map built"); + break; + } + Err(e) => { + tracing::error!("Failed to load history: {}", e); + break; + } + } + } + })); + + // Spawn background task to periodically refresh frecency + let index_for_frecency = self.index.clone(); + let handle_for_frecency = handle.clone(); + self.frecency_handle = Some(tokio::spawn(async move { + let mut interval = tokio::time::interval(std::time::Duration::from_secs( + FRECENCY_REFRESH_INTERVAL_SECS, + )); + loop { + interval.tick().await; + trace!("Refreshing frecency map"); + let settings = handle_for_frecency.settings().await; + index_for_frecency + .read() + .await + .rebuild_frecency(&settings.search) + .await; + } + })); + + tracing::info!("search component started"); + Ok(()) + } + + async fn handle_event(&mut self, event: &DaemonEvent) -> Result<()> { + match event { + DaemonEvent::RecordsAdded(records) => { + debug!( + count = records.len(), + "Processing added records for search index" + ); + + let handle_guard = self.handle.read().await; + if let Some(handle) = handle_guard.as_ref() { + let histories: Vec<_> = handle + .history_db() + .query_history( + format!( + "select * from history where id in ({})", + records + .iter() + .map(|record| record.0.to_string()) + .collect::<Vec<_>>() + .join(",") + ) + .as_str(), + ) + .await + .unwrap_or_default(); + + span!(Level::TRACE, "inject_records", count = histories.len()) + .in_scope(async || { + self.index.read().await.add_histories(&histories); + }) + .await; + } + } + DaemonEvent::HistoryStarted(history) => { + debug!(id = %history.id, command = %history.command, "History started (no index action)"); + } + DaemonEvent::HistoryEnded(history) => { + span!(Level::TRACE, "inject_history_ended") + .in_scope(async || { + self.index.read().await.add_history(history); + }) + .await; + } + DaemonEvent::HistoryPruned | DaemonEvent::HistoryRebuilt => { + info!("History store pruned or rebuilt, rebuilding search index"); + if let Err(e) = self.rebuild_index().await { + tracing::error!("Failed to rebuild search index: {}", e); + } + } + DaemonEvent::HistoryDeleted { ids } => { + info!( + count = ids.len(), + "History deleted, rebuilding search index" + ); + // For now, just rebuild the entire index. A more efficient implementation + // would remove specific items from the index. + if let Err(e) = self.rebuild_index().await { + tracing::error!("Failed to rebuild search index: {}", e); + } + } + DaemonEvent::SettingsReloaded => { + info!("Settings reloaded, rebuilding frecency map with new multipliers"); + let handle_guard = self.handle.read().await; + if let Some(handle) = handle_guard.as_ref() { + let settings = handle.settings().await; + self.index + .read() + .await + .rebuild_frecency(&settings.search) + .await; + } + } + // Events we don't care about + DaemonEvent::SyncCompleted { .. } + | DaemonEvent::SyncFailed { .. } + | DaemonEvent::ForceSync + | DaemonEvent::ShutdownRequested => {} + } + Ok(()) + } + + async fn stop(&mut self) -> Result<()> { + if let Some(handle) = self.loader_handle.take() { + handle.abort(); + } + if let Some(handle) = self.frecency_handle.take() { + handle.abort(); + } + tracing::info!("search component stopped"); + Ok(()) + } +} + +/// The gRPC service implementation. +pub struct SearchGrpcService { + index: Arc<RwLock<SearchIndex>>, +} + +#[tonic::async_trait] +impl SearchSvc for SearchGrpcService { + type SearchStream = Pin<Box<dyn Stream<Item = Result<SearchResponse, Status>> + Send>>; + + #[instrument(skip_all, level = Level::TRACE, name = "search_rpc")] + async fn search( + &self, + request: Request<Streaming<SearchRequest>>, + ) -> Result<Response<Self::SearchStream>, Status> { + let mut in_stream = request.into_inner(); + let index = self.index.clone(); + + // Create output channel + let (tx, rx) = tokio::sync::mpsc::channel::<Result<SearchResponse, Status>>(128); + + // Spawn task to handle incoming requests and send responses + tokio::spawn(async move { + while let Some(req) = in_stream.message().await.transpose() { + match req { + Ok(search_req) => { + let query = search_req.query; + let query_id = search_req.query_id; + let filter_mode: FilterMode = search_req + .filter_mode + .try_into() + .unwrap_or(FilterMode::Global); + let proto_context = search_req.context; + + debug!( + "search request: query = {}, query_id = {}, filter_mode = {}, context = {:?}", + query, + query_id, + filter_mode.as_str_name(), + proto_context + ); + + // Convert proto FilterMode + context to IndexFilterMode + let index_filter = convert_filter_mode(filter_mode, &proto_context); + + // Build QueryContext from proto context + let query_context = proto_context + .map(|ctx| QueryContext { + cwd: Some(with_trailing_slash(&ctx.cwd)), + git_root: ctx.git_root.map(|s| with_trailing_slash(&s)), + hostname: Some(ctx.hostname), + session_id: Some(ctx.session_id), + }) + .unwrap_or_default(); + + // Perform the search + let history_ids = + span!(Level::TRACE, "daemon_search_query", %query, query_id) + .in_scope(|| async { + let index = index.read().await; + index + .search(&query, index_filter, &query_context, RESULTS_LIMIT) + .await + }) + .await; + + // Convert history IDs to bytes + let ids: Vec<Vec<u8>> = history_ids + .iter() + .filter_map(|id| { + Uuid::parse_str(id) + .ok() + .map(|uuid| uuid.as_bytes().to_vec()) + }) + .collect(); + + if tx.send(Ok(SearchResponse { query_id, ids })).await.is_err() { + break; // Client disconnected + } + } + Err(e) => { + let _ = tx.send(Err(e)).await; + break; + } + } + } + }); + + // Convert receiver to stream + let out_stream = tokio_stream::wrappers::ReceiverStream::new(rx); + Ok(Response::new(Box::pin(out_stream))) + } +} + +/// Convert proto FilterMode and context to IndexFilterMode. +fn convert_filter_mode( + mode: FilterMode, + context: &Option<crate::search::SearchContext>, +) -> IndexFilterMode { + match (mode, context) { + (FilterMode::Global, _) => IndexFilterMode::Global, + (FilterMode::Directory, Some(ctx)) => { + IndexFilterMode::Directory(with_trailing_slash(&ctx.cwd)) + } + (FilterMode::Workspace, Some(ctx)) => { + if let Some(ref git_root) = ctx.git_root { + IndexFilterMode::Workspace(with_trailing_slash(git_root)) + } else { + // Fall back to directory if no git root + IndexFilterMode::Directory(with_trailing_slash(&ctx.cwd)) + } + } + (FilterMode::Host, Some(ctx)) => IndexFilterMode::Host(ctx.hostname.clone()), + (FilterMode::Session, Some(ctx)) => IndexFilterMode::Session(ctx.session_id.clone()), + (FilterMode::SessionPreload, Some(ctx)) => { + // SessionPreload is similar to Session - filter by session + IndexFilterMode::Session(ctx.session_id.clone()) + } + // If no context provided, fall back to global + _ => IndexFilterMode::Global, + } +} + +#[cfg(windows)] +pub fn with_trailing_slash(s: &str) -> String { + if s.ends_with('\\') { + s.to_string() + } else { + format!("{}\\", s) + } +} + +#[cfg(not(windows))] +pub fn with_trailing_slash(s: &str) -> String { + if s.ends_with('/') { + s.to_string() + } else { + format!("{}/", s) + } +} diff --git a/crates/atuin-daemon/src/components/sync.rs b/crates/atuin-daemon/src/components/sync.rs new file mode 100644 index 00000000..314b375e --- /dev/null +++ b/crates/atuin-daemon/src/components/sync.rs @@ -0,0 +1,296 @@ +//! Sync component. +//! +//! Handles periodic synchronization with the Atuin cloud server. + +use std::time::Duration; + +use eyre::Result; +use rand::Rng; +use tokio::sync::mpsc; +use tokio::time::{self, MissedTickBehavior}; + +use atuin_client::{history::store::HistoryStore, record::sync, settings::Settings}; +use atuin_dotfiles::store::{AliasStore, var::VarStore}; + +use crate::{ + daemon::{Component, DaemonHandle}, + events::DaemonEvent, +}; + +/// Commands that can be sent to the sync task. +enum SyncCommand { + /// Trigger an immediate sync. + ForceSync, + /// Stop the sync loop. + Stop, +} + +/// Sync state - tracks whether we're in normal operation or retrying after failure. +#[derive(Clone, Copy, PartialEq, Eq)] +enum SyncState { + /// Normal operation. Periodic syncs only run if auto_sync is enabled. + Idle, + /// Retrying after a sync failure. Retries continue regardless of auto_sync + /// until the sync succeeds. + Retrying, +} + +/// Sync component - handles periodic cloud synchronization. +/// +/// This component: +/// - Runs a background sync loop on a configurable interval +/// - Implements exponential backoff on sync failures +/// - Responds to ForceSync events for immediate sync +/// - Emits SyncCompleted/SyncFailed events +pub struct SyncComponent { + task_handle: Option<tokio::task::JoinHandle<()>>, + command_tx: Option<mpsc::Sender<SyncCommand>>, +} + +impl SyncComponent { + /// Create a new sync component. + pub fn new() -> Self { + Self { + task_handle: None, + command_tx: None, + } + } +} + +impl Default for SyncComponent { + fn default() -> Self { + Self::new() + } +} + +#[tonic::async_trait] +impl Component for SyncComponent { + fn name(&self) -> &'static str { + "sync" + } + + async fn start(&mut self, handle: DaemonHandle) -> Result<()> { + let (cmd_tx, cmd_rx) = mpsc::channel(16); + self.command_tx = Some(cmd_tx); + + // Spawn the sync loop with its own copy of the handle + self.task_handle = Some(tokio::spawn(sync_loop(handle, cmd_rx))); + + tracing::info!("sync component started"); + Ok(()) + } + + async fn handle_event(&mut self, event: &DaemonEvent) -> Result<()> { + if let DaemonEvent::ForceSync = event { + tracing::info!("force sync requested"); + if let Some(tx) = &self.command_tx { + let _ = tx.send(SyncCommand::ForceSync).await; + } + } + Ok(()) + } + + async fn stop(&mut self) -> Result<()> { + if let Some(tx) = &self.command_tx { + let _ = tx.send(SyncCommand::Stop).await; + } + if let Some(handle) = self.task_handle.take() { + // Give the task a moment to shut down gracefully + let _ = tokio::time::timeout(std::time::Duration::from_secs(5), handle).await; + } + tracing::info!("sync component stopped"); + Ok(()) + } +} + +/// The main sync loop. +/// +/// This runs in a spawned task and handles periodic sync as well as +/// force sync requests. +async fn sync_loop(handle: DaemonHandle, mut cmd_rx: mpsc::Receiver<SyncCommand>) { + tracing::info!("sync loop starting"); + + // Clone settings since we need them across await points + let settings = handle.settings().await.clone(); + let host_id = match Settings::host_id().await { + Ok(id) => id, + Err(e) => { + tracing::error!("failed to get host id, sync disabled: {e}"); + return; + } + }; + + // Create the stores we need + let encryption_key = *handle.encryption_key(); + let history_store = HistoryStore::new(handle.store().clone(), host_id, encryption_key); + let alias_store = AliasStore::new(handle.store().clone(), host_id, encryption_key); + let var_store = VarStore::new(handle.store().clone(), host_id, encryption_key); + + // Don't backoff by more than 30 mins (with a random jitter of up to 1 min) + let max_interval: f64 = 60.0 * 30.0 + rand::thread_rng().gen_range(0.0..60.0); + + let mut ticker = time::interval(time::Duration::from_secs(settings.daemon.sync_frequency)); + + // IMPORTANT: without this, if we miss ticks because a sync takes ages or is otherwise delayed, + // we may end up running a lot of syncs in a hot loop. + ticker.set_missed_tick_behavior(MissedTickBehavior::Skip); + + let mut sync_state = SyncState::Idle; + + loop { + tokio::select! { + _ = ticker.tick() => { + let settings = handle.settings().await; + + // Skip periodic ticks if auto_sync is disabled AND we're not retrying + // a previous failure. Retries must continue regardless of auto_sync. + if !settings.auto_sync && sync_state == SyncState::Idle { + tracing::debug!("auto_sync disabled, skipping periodic sync tick"); + continue; + } + + sync_state = do_sync_tick( + &handle, + &history_store, + &alias_store, + &var_store, + &mut ticker, + max_interval, + &settings, + ).await; + } + cmd = cmd_rx.recv() => { + match cmd { + Some(SyncCommand::ForceSync) => { + tracing::info!("executing force sync"); + let settings = handle.settings().await; + sync_state = do_sync_tick( + &handle, + &history_store, + &alias_store, + &var_store, + &mut ticker, + max_interval, + &settings, + ).await; + } + Some(SyncCommand::Stop) | None => { + tracing::info!("sync loop stopping"); + break; + } + } + } + } + } +} + +/// Execute a single sync tick. +/// +/// Returns the new sync state: `Idle` on success, `Retrying` on failure. +async fn do_sync_tick( + handle: &DaemonHandle, + history_store: &HistoryStore, + alias_store: &AliasStore, + var_store: &VarStore, + ticker: &mut time::Interval, + max_interval: f64, + settings: &Settings, +) -> SyncState { + tracing::info!("sync tick"); + + // Check if logged in + let logged_in = match settings.logged_in().await { + Ok(v) => v, + Err(e) => { + tracing::warn!("failed to check login status, skipping sync tick: {e}"); + return SyncState::Idle; + } + }; + + if !logged_in { + tracing::debug!("not logged in, skipping sync tick"); + return SyncState::Idle; + } + + // Perform the sync + let res = sync::sync(settings, handle.store()).await; + + match res { + Err(e) => { + tracing::error!("sync tick failed with {e}"); + + // Emit failure event + handle.emit(DaemonEvent::SyncFailed { + error: e.to_string(), + }); + + // Exponential backoff + let mut rng = rand::thread_rng(); + let mut new_interval = ticker.period().as_secs_f64() * rng.gen_range(2.0..2.2); + + if new_interval > max_interval { + new_interval = max_interval; + } + + *ticker = time::interval_at( + tokio::time::Instant::now() + Duration::from_secs(new_interval as u64), + time::Duration::from_secs(new_interval as u64), + ); + ticker.reset_after(time::Duration::from_secs(new_interval as u64)); + ticker.set_missed_tick_behavior(MissedTickBehavior::Skip); + + tracing::error!("backing off, next sync tick in {new_interval}"); + + SyncState::Retrying + } + Ok((uploaded_count, downloaded_records)) => { + tracing::info!( + uploaded = uploaded_count, + downloaded = downloaded_records.len(), + "sync complete" + ); + + // Build history from downloaded records + if let Err(e) = history_store + .incremental_build(handle.history_db(), &downloaded_records) + .await + { + tracing::error!("failed to build history from downloaded records: {e}"); + } + + // Emit the records added event (for search indexing) + handle.emit(DaemonEvent::RecordsAdded(downloaded_records.clone())); + + // Emit sync completed event + handle.emit(DaemonEvent::SyncCompleted { + uploaded: uploaded_count as usize, + downloaded: downloaded_records.len(), + }); + + // Rebuild alias and var stores + if let Err(e) = alias_store.build().await { + tracing::error!("failed to rebuild alias store: {e}"); + } + if let Err(e) = var_store.build().await { + tracing::error!("failed to rebuild var store: {e}"); + } + + // Reset backoff on success + if ticker.period().as_secs() != settings.daemon.sync_frequency { + *ticker = time::interval_at( + tokio::time::Instant::now() + + Duration::from_secs(settings.daemon.sync_frequency), + time::Duration::from_secs(settings.daemon.sync_frequency), + ); + ticker.set_missed_tick_behavior(MissedTickBehavior::Skip); + } + + // Store sync time + if let Err(e) = Settings::save_sync_time().await { + tracing::error!("failed to save sync time: {e}"); + } + + SyncState::Idle + } + } +} diff --git a/crates/atuin-daemon/src/control/mod.rs b/crates/atuin-daemon/src/control/mod.rs new file mode 100644 index 00000000..afb29c57 --- /dev/null +++ b/crates/atuin-daemon/src/control/mod.rs @@ -0,0 +1,12 @@ +//! Control module for external event injection. +//! +//! This module provides the gRPC service that allows external processes +//! (like CLI commands) to inject events into the daemon's event bus. + +mod service; + +// Include the generated proto code +tonic::include_proto!("control"); + +// Re-export the service +pub use service::ControlService; diff --git a/crates/atuin-daemon/src/control/service.rs b/crates/atuin-daemon/src/control/service.rs new file mode 100644 index 00000000..2e7403ce --- /dev/null +++ b/crates/atuin-daemon/src/control/service.rs @@ -0,0 +1,71 @@ +//! Control service implementation. +//! +//! This gRPC service allows external processes (like CLI commands) to inject +//! events into the daemon's event bus. + +use atuin_client::history::HistoryId; +use tonic::{Request, Response, Status}; +use tracing::{Level, info, instrument}; + +use super::{ + SendEventRequest, SendEventResponse, + control_server::{Control, ControlServer}, + send_event_request::Event, +}; +use crate::{daemon::DaemonHandle, events::DaemonEvent}; + +/// The Control gRPC service. +/// +/// This service is used by external processes to inject events into the daemon. +/// It's not a component - it's part of the daemon's core infrastructure. +pub struct ControlService { + handle: DaemonHandle, +} + +impl ControlService { + /// Create a new control service with the given daemon handle. + pub fn new(handle: DaemonHandle) -> Self { + Self { handle } + } + + /// Get a tonic server for this service. + pub fn into_server(self) -> ControlServer<Self> { + ControlServer::new(self) + } +} + +#[tonic::async_trait] +impl Control for ControlService { + #[instrument(skip_all, level = Level::INFO, name = "control_send_event")] + async fn send_event( + &self, + request: Request<SendEventRequest>, + ) -> Result<Response<SendEventResponse>, Status> { + let req = request.into_inner(); + + let event = req + .event + .ok_or_else(|| Status::invalid_argument("event is required"))?; + + let daemon_event = proto_event_to_daemon_event(event)?; + + info!(?daemon_event, "received control event"); + self.handle.emit(daemon_event); + + Ok(Response::new(SendEventResponse {})) + } +} + +/// Convert a proto event to a daemon event. +fn proto_event_to_daemon_event(event: Event) -> Result<DaemonEvent, Status> { + match event { + Event::HistoryPruned(_) => Ok(DaemonEvent::HistoryPruned), + Event::HistoryRebuilt(_) => Ok(DaemonEvent::HistoryRebuilt), + Event::HistoryDeleted(e) => Ok(DaemonEvent::HistoryDeleted { + ids: e.ids.into_iter().map(HistoryId).collect(), + }), + Event::ForceSync(_) => Ok(DaemonEvent::ForceSync), + Event::SettingsReloaded(_) => Ok(DaemonEvent::SettingsReloaded), + Event::Shutdown(_) => Ok(DaemonEvent::ShutdownRequested), + } +} diff --git a/crates/atuin-daemon/src/daemon.rs b/crates/atuin-daemon/src/daemon.rs new file mode 100644 index 00000000..625ca205 --- /dev/null +++ b/crates/atuin-daemon/src/daemon.rs @@ -0,0 +1,458 @@ +//! Core daemon infrastructure. +//! +//! This module provides the foundational types for building the atuin daemon: +//! +//! - [`DaemonState`]: Shared state owned by the daemon +//! - [`DaemonHandle`]: A lightweight, cloneable handle for accessing daemon state +//! - [`Component`]: A trait for implementing daemon components +//! - [`Daemon`]: The main daemon orchestrator +//! - [`DaemonBuilder`]: Builder for constructing and configuring the daemon + +use std::sync::Arc; + +use atuin_client::{ + database::Sqlite as HistoryDatabase, encryption, record::sqlite_store::SqliteStore, + settings::Settings, +}; +use eyre::{Context, Result}; +use tokio::sync::{RwLock, broadcast}; + +use crate::events::DaemonEvent; + +// ============================================================================ +// DaemonState +// ============================================================================ + +/// Shared state owned by the daemon. +/// +/// This contains all the resources that components and services need access to. +/// The state is wrapped in an `Arc` and accessed via [`DaemonHandle`]. +pub struct DaemonState { + // Event bus + event_tx: broadcast::Sender<DaemonEvent>, + + // Configuration (mutable - can be reloaded) + settings: RwLock<Settings>, + + // Encryption key (immutable - derived at startup) + encryption_key: [u8; 32], + + // Database handles + history_db: HistoryDatabase, + store: SqliteStore, +} + +// ============================================================================ +// DaemonHandle +// ============================================================================ + +/// A lightweight handle to the daemon's shared state. +/// +/// This is the primary way for components, gRPC services, and spawned tasks to +/// interact with the daemon. It provides access to: +/// +/// - Event emission and subscription +/// - Configuration (settings, encryption key) +/// - Database handles +/// +/// The handle is cheaply cloneable (wraps an `Arc`) and can be freely passed +/// around to any code that needs daemon access. +/// +/// # Example +/// +/// ```ignore +/// // Emit an event +/// handle.emit(DaemonEvent::HistoryPruned); +/// +/// // Access settings +/// let settings = handle.settings().await; +/// let sync_freq = settings.daemon.sync_frequency; +/// +/// // Access database +/// let history = handle.history_db().load(id).await?; +/// ``` +#[derive(Clone)] +pub struct DaemonHandle { + state: Arc<DaemonState>, +} + +impl DaemonHandle { + // ---- Events ---- + + /// Emit an event to the daemon's event bus. + /// + /// This is fire-and-forget - if no receivers are listening (which shouldn't + /// happen in normal operation), the event is dropped silently. + pub fn emit(&self, event: DaemonEvent) { + if let Err(e) = self.state.event_tx.send(event) { + tracing::warn!("failed to emit event (no receivers?): {e}"); + } + } + + /// Subscribe to the event bus. + /// + /// Returns a receiver that will receive all events emitted after this call. + /// Useful for components that need to listen for events outside of the + /// normal `handle_event` callback flow. + pub fn subscribe(&self) -> broadcast::Receiver<DaemonEvent> { + self.state.event_tx.subscribe() + } + + /// Request graceful shutdown of the daemon. + pub fn shutdown(&self) { + self.emit(DaemonEvent::ShutdownRequested); + } + + // ---- Configuration ---- + + /// Get the current settings. + /// + /// This acquires a read lock on the settings. For most use cases, clone + /// the settings if you need to hold onto them. + pub async fn settings(&self) -> tokio::sync::RwLockReadGuard<'_, Settings> { + self.state.settings.read().await + } + + /// Reload settings from disk and emit a SettingsReloaded event. + /// + /// Components listening for `SettingsReloaded` can then re-read settings + /// via `handle.settings()` to pick up the changes. + pub async fn reload_settings(&self) -> Result<()> { + let new_settings = Settings::new()?; + self.apply_settings(new_settings).await; + Ok(()) + } + + /// Apply already-loaded settings and emit a SettingsReloaded event. + /// + /// Use this when settings have already been loaded (e.g., from a file watcher) + /// to avoid parsing the config file twice. + pub async fn apply_settings(&self, settings: Settings) { + *self.state.settings.write().await = settings; + self.emit(DaemonEvent::SettingsReloaded); + tracing::info!("settings applied"); + } + + /// Get the encryption key. + pub fn encryption_key(&self) -> &[u8; 32] { + &self.state.encryption_key + } + + // ---- Database ---- + + /// Get a reference to the history database. + pub fn history_db(&self) -> &HistoryDatabase { + &self.state.history_db + } + + /// Get a reference to the record store. + pub fn store(&self) -> &SqliteStore { + &self.state.store + } +} + +impl std::fmt::Debug for DaemonHandle { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("DaemonHandle").finish_non_exhaustive() + } +} + +// ============================================================================ +// Component Trait +// ============================================================================ + +/// A daemon component that handles a specific domain. +/// +/// Components are the building blocks of the daemon. Each component: +/// +/// - Has a unique name for logging and debugging +/// - Can optionally expose gRPC services +/// - Receives a [`DaemonHandle`] on startup for accessing daemon resources +/// - Handles events from the event bus +/// - Performs cleanup on shutdown +/// +/// # Lifecycle +/// +/// 1. **Construction**: Component is created (usually via `new()`) +/// 2. **Start**: `start()` is called with a [`DaemonHandle`] +/// 3. **Running**: `handle_event()` is called for each event on the bus +/// 4. **Shutdown**: `stop()` is called for cleanup +/// +/// # Example +/// +/// ```ignore +/// pub struct MyComponent { +/// handle: Option<DaemonHandle>, +/// } +/// +/// #[async_trait] +/// impl Component for MyComponent { +/// fn name(&self) -> &'static str { "my-component" } +/// +/// async fn start(&mut self, handle: DaemonHandle) -> Result<()> { +/// self.handle = Some(handle); +/// Ok(()) +/// } +/// +/// async fn handle_event(&mut self, event: &DaemonEvent) -> Result<()> { +/// match event { +/// DaemonEvent::SomeEvent => { +/// // Handle the event +/// if let Some(handle) = &self.handle { +/// handle.emit(DaemonEvent::ResponseEvent); +/// } +/// } +/// _ => {} +/// } +/// Ok(()) +/// } +/// +/// async fn stop(&mut self) -> Result<()> { +/// Ok(()) +/// } +/// } +/// ``` +#[tonic::async_trait] +pub trait Component: Send + Sync { + /// Human-readable name for logging and debugging. + fn name(&self) -> &'static str; + + /// Called once at startup. + /// + /// Store the handle if you need to emit events or access daemon resources + /// later. The handle is cheaply cloneable, so feel free to clone it for + /// spawned tasks. + async fn start(&mut self, handle: DaemonHandle) -> Result<()>; + + /// Handle an incoming event. + /// + /// Called for every event on the bus. To emit new events in response, + /// use the handle stored during `start()`. Events emitted here will be + /// processed in subsequent event loop iterations. + async fn handle_event(&mut self, event: &DaemonEvent) -> Result<()>; + + /// Called on graceful shutdown. + /// + /// Use this to clean up resources, abort spawned tasks, etc. + async fn stop(&mut self) -> Result<()>; +} + +// ============================================================================ +// Daemon +// ============================================================================ + +/// The main daemon orchestrator. +/// +/// The daemon manages components, runs the event loop, and coordinates startup +/// and shutdown. It is constructed via [`DaemonBuilder`]. +/// +/// # Event Loop +/// +/// The daemon runs a simple event loop: +/// +/// 1. Wait for an event on the bus +/// 2. Dispatch the event to all components (in registration order) +/// 3. Components may emit new events in response +/// 4. Repeat until `ShutdownRequested` is received +/// +/// Events emitted during handling are queued and processed in subsequent +/// iterations, ensuring the loop eventually drains. +pub struct Daemon { + components: Vec<Box<dyn Component>>, + handle: DaemonHandle, +} + +impl Daemon { + /// Create a new daemon builder. + pub fn builder(settings: Settings) -> DaemonBuilder { + DaemonBuilder::new(settings) + } + + /// Get a clone of the daemon handle. + /// + /// The handle can be used to emit events, access settings, etc. + pub fn handle(&self) -> DaemonHandle { + self.handle.clone() + } + + /// Start all components. + /// + /// This must be called before `run_event_loop()`. It initializes all + /// registered components with the daemon handle. + pub async fn start_components(&mut self) -> Result<()> { + for component in &mut self.components { + tracing::info!(component = component.name(), "starting component"); + component + .start(self.handle.clone()) + .await + .with_context(|| format!("failed to start component: {}", component.name()))?; + } + Ok(()) + } + + /// Run the daemon event loop. + /// + /// This processes events until a ShutdownRequested event is received. + /// Components must be started first via `start_components()`. + pub async fn run_event_loop(&mut self) -> Result<()> { + let mut event_rx = self.handle.subscribe(); + loop { + match event_rx.recv().await { + Ok(DaemonEvent::ShutdownRequested) => { + tracing::info!("shutdown requested, stopping daemon"); + break; + } + Ok(event) => { + tracing::debug!(?event, "processing event"); + self.dispatch_event(&event).await; + } + Err(broadcast::error::RecvError::Lagged(n)) => { + tracing::warn!( + skipped = n, + "event receiver lagged, some events were dropped" + ); + } + Err(broadcast::error::RecvError::Closed) => { + tracing::info!("event bus closed, stopping daemon"); + break; + } + } + } + Ok(()) + } + + /// Stop all components. + /// + /// This performs graceful shutdown of all components. + pub async fn stop_components(&mut self) { + for component in &mut self.components { + tracing::info!(component = component.name(), "stopping component"); + if let Err(e) = component.stop().await { + tracing::error!( + component = component.name(), + error = ?e, + "error stopping component" + ); + } + } + tracing::info!("all components stopped"); + } + + /// Run the daemon. + /// + /// This is a convenience method that starts components, runs the event loop, + /// and handles shutdown. It does not return until the daemon is shut down. + pub async fn run(mut self) -> Result<()> { + self.start_components().await?; + self.run_event_loop().await?; + self.stop_components().await; + tracing::info!("daemon stopped"); + Ok(()) + } + + async fn dispatch_event(&mut self, event: &DaemonEvent) { + for component in &mut self.components { + if let Err(e) = component.handle_event(event).await { + tracing::error!( + component = component.name(), + error = ?e, + "error handling event" + ); + } + } + } +} + +// ============================================================================ +// DaemonBuilder +// ============================================================================ + +/// Builder for constructing a [`Daemon`]. +/// +/// # Example +/// +/// ```ignore +/// let daemon = Daemon::builder(settings) +/// .store(store) +/// .history_db(history_db) +/// .component(HistoryComponent::new()) +/// .component(SearchComponent::new()) +/// .component(SyncComponent::new()) +/// .build() +/// .await?; +/// +/// daemon.run().await?; +/// ``` +pub struct DaemonBuilder { + settings: Settings, + store: Option<SqliteStore>, + history_db: Option<HistoryDatabase>, + components: Vec<Box<dyn Component>>, +} + +impl DaemonBuilder { + /// Create a new daemon builder with the given settings. + pub fn new(settings: Settings) -> Self { + Self { + settings, + store: None, + history_db: None, + components: Vec::new(), + } + } + + /// Set the record store. + pub fn store(mut self, store: SqliteStore) -> Self { + self.store = Some(store); + self + } + + /// Set the history database. + pub fn history_db(mut self, db: HistoryDatabase) -> Self { + self.history_db = Some(db); + self + } + + /// Register a component. + /// + /// Components are started in registration order and stopped in reverse order. + pub fn component(mut self, component: impl Component + 'static) -> Self { + self.components.push(Box::new(component)); + self + } + + /// Build the daemon. + /// + /// This loads the encryption key and creates the daemon state. + pub async fn build(self) -> Result<Daemon> { + let store = self.store.ok_or_else(|| eyre::eyre!("store is required"))?; + let history_db = self + .history_db + .ok_or_else(|| eyre::eyre!("history_db is required"))?; + + // Load encryption key + let encryption_key: [u8; 32] = encryption::load_key(&self.settings) + .context("could not load encryption key")? + .into(); + + // Create the event bus + let (event_tx, _) = broadcast::channel(64); + + // Create the shared state + let state = Arc::new(DaemonState { + event_tx, + settings: RwLock::new(self.settings), + encryption_key, + history_db, + store, + }); + + // Create the handle (just a reference to the state) + let handle = DaemonHandle { state }; + + Ok(Daemon { + components: self.components, + handle, + }) + } +} diff --git a/crates/atuin-daemon/src/events.rs b/crates/atuin-daemon/src/events.rs new file mode 100644 index 00000000..4e6c6ff3 --- /dev/null +++ b/crates/atuin-daemon/src/events.rs @@ -0,0 +1,74 @@ +//! Daemon events. +//! +//! Events are the primary communication mechanism within the daemon. +//! Components emit events to notify others of state changes, and handle +//! events to react to changes elsewhere in the system. +//! +//! External processes (like CLI commands) can also inject events via the +//! Control gRPC service. + +use atuin_client::history::{History, HistoryId}; +use atuin_common::record::RecordId; + +/// Events that flow through the daemon's event bus. +/// +/// Events are broadcast to all components. Each component decides which +/// events it cares about in its `handle_event` implementation. +#[derive(Debug, Clone)] +pub enum DaemonEvent { + // ---- History lifecycle ---- + /// A command has started running. + HistoryStarted(History), + + /// A command has finished running. + HistoryEnded(History), + + // ---- Sync ---- + /// Records were synced from the server. + /// + /// The search component uses this to update its index with new history. + RecordsAdded(Vec<RecordId>), + + /// Sync completed successfully. + SyncCompleted { + /// Number of records uploaded. + uploaded: usize, + /// Number of records downloaded. + downloaded: usize, + }, + + /// Sync failed. + SyncFailed { + /// Error message describing what went wrong. + error: String, + }, + + /// Request an immediate sync (external trigger). + ForceSync, + + // ---- External commands ---- + /// History was pruned - search index needs a full rebuild. + /// + /// Emitted when the user runs `atuin history prune` or similar. + HistoryPruned, + + /// History was rebuilt - search index needs a full rebuild. + /// + /// Emitted when the user runs `atuin store rebuild history` or similar. + HistoryRebuilt, + + /// Specific history items were deleted. + /// + /// The search component should remove these from its index. + HistoryDeleted { + /// IDs of the deleted history entries. + ids: Vec<HistoryId>, + }, + + /// Settings have changed, components should reload if needed. + SettingsReloaded, + + // ---- Lifecycle ---- + /// Request graceful shutdown of the daemon. + ShutdownRequested, +} diff --git a/crates/atuin-daemon/src/history/mod.rs b/crates/atuin-daemon/src/history/mod.rs new file mode 100644 index 00000000..b71853df --- /dev/null +++ b/crates/atuin-daemon/src/history/mod.rs @@ -0,0 +1,6 @@ +//! History module for the daemon gRPC history service. +//! +//! This module contains the proto-generated types for the history gRPC service. + +// Include the generated proto code +tonic::include_proto!("history"); diff --git a/crates/atuin-daemon/src/lib.rs b/crates/atuin-daemon/src/lib.rs new file mode 100644 index 00000000..84f808e4 --- /dev/null +++ b/crates/atuin-daemon/src/lib.rs @@ -0,0 +1,131 @@ +use atuin_client::database::Sqlite as HistoryDatabase; +use atuin_client::record::sqlite_store::SqliteStore; +use atuin_client::settings::{Settings, watcher::global_settings_watcher}; +use eyre::Result; + +pub mod client; +pub mod components; +pub mod control; +pub mod daemon; +pub mod events; +pub mod history; +pub mod search; +pub mod server; + +// Re-export core daemon types for convenience +pub use daemon::{Component, Daemon, DaemonBuilder, DaemonHandle}; +pub use events::DaemonEvent; + +// Re-export components +pub use components::{HistoryComponent, SearchComponent, SyncComponent}; + +// Re-export client helpers +pub use client::{ControlClient, emit_event, emit_event_with_settings}; + +/// Boot the daemon using the new component-based architecture. +/// +/// This creates a daemon with the standard components (history, search, sync), +/// starts the gRPC server with their services, and runs the event loop. +pub async fn boot( + settings: Settings, + store: SqliteStore, + history_db: HistoryDatabase, +) -> Result<()> { + // Create the components + let history_component = HistoryComponent::new(); + let search_component = SearchComponent::new(); + let sync_component = SyncComponent::new(); + + // Get the gRPC services before moving components into the daemon + // (The services share state with the components via Arc) + let history_service = history_component.grpc_service(); + let search_service = search_component.grpc_service(); + + // Build the daemon + let mut daemon = Daemon::builder(settings.clone()) + .store(store) + .history_db(history_db) + .component(history_component) + .component(search_component) + .component(sync_component) + .build() + .await?; + + // Get a handle for the control service and gRPC server shutdown + let handle = daemon.handle(); + + // Create the control service + let control_service = control::ControlService::new(handle.clone()); + + // Start all components first (so gRPC services can work) + daemon.start_components().await?; + + // Spawn config file watcher to reload settings on changes + if let Ok(watcher) = global_settings_watcher() { + let mut settings_rx = watcher.subscribe(); + let watcher_handle = handle.clone(); + tokio::spawn(async move { + tracing::info!("config file watcher started"); + while settings_rx.changed().await.is_ok() { + // Use the already-loaded settings from the watcher + // (avoids parsing the config file twice) + let new_settings = (*settings_rx.borrow()).clone(); + watcher_handle.apply_settings((*new_settings).clone()).await; + } + tracing::debug!("config file watcher stopped"); + }); + } else { + tracing::warn!( + "failed to start config file watcher; settings changes will require daemon restart" + ); + } + + // Spawn signal handler to emit ShutdownRequested on Ctrl+C/SIGTERM + let signal_handle = handle.clone(); + tokio::spawn(async move { + shutdown_signal().await; + tracing::info!("received shutdown signal"); + signal_handle.shutdown(); + }); + + // Start the gRPC server in the background + server::run_grpc_server( + settings, + history_service, + search_service, + control_service.into_server(), + handle, + ) + .await?; + + // Run the daemon event loop + daemon.run_event_loop().await?; + + // Stop all components on shutdown + daemon.stop_components().await; + + tracing::info!("daemon shut down complete"); + Ok(()) +} + +/// Wait for a shutdown signal (Ctrl+C or SIGTERM). +#[cfg(unix)] +async fn shutdown_signal() { + let mut term = tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()) + .expect("failed to register sigterm handler"); + let mut int = tokio::signal::unix::signal(tokio::signal::unix::SignalKind::interrupt()) + .expect("failed to register sigint handler"); + + tokio::select! { + _ = term.recv() => {}, + _ = int.recv() => {}, + } +} + +/// Wait for a shutdown signal (Ctrl+C). +#[cfg(not(unix))] +async fn shutdown_signal() { + tokio::signal::ctrl_c() + .await + .expect("failed to listen for ctrl+c"); +} diff --git a/crates/atuin-daemon/src/search/index.rs b/crates/atuin-daemon/src/search/index.rs new file mode 100644 index 00000000..1445871e --- /dev/null +++ b/crates/atuin-daemon/src/search/index.rs @@ -0,0 +1,679 @@ +//! Search index with frecency-based ranking. +//! +//! This module provides a deduplicated search index where each unique command +//! is stored once, with metadata about all its invocations. This enables: +//! +//! - Efficient fuzzy matching (fewer items to match) +//! - Frecency-based ranking (frequency + recency) +//! - Dynamic filtering by directory, host, session, etc. + +use std::{ + collections::{HashMap, HashSet}, + sync::Arc, +}; + +use atuin_client::history::History; +use atuin_client::settings::Search; +use dashmap::DashMap; +use lasso::{Spur, ThreadedRodeo}; +use nucleo::{Injector, Nucleo, pattern}; +use time::OffsetDateTime; +use tokio::sync::RwLock; +use tracing::{Level, instrument}; +use uuid::Uuid; + +use crate::components::search::with_trailing_slash; + +/// Parse a UUID string into a 16-byte array. +/// Returns None if the string is not a valid UUID. +fn parse_uuid_bytes(s: &str) -> Option<[u8; 16]> { + Uuid::parse_str(s).ok().map(|u| *u.as_bytes()) +} + +/// Format a 16-byte array as a UUID string. +fn format_uuid_bytes(bytes: &[u8; 16]) -> String { + Uuid::from_bytes(*bytes).to_string() +} + +/// Pre-computed frecency data for O(1) lookup. +#[derive(Debug, Clone, Default)] +pub struct FrecencyData { + /// Total number of times this command was used. + pub count: u32, + /// Most recent usage timestamp (unix seconds). + pub last_used: i64, +} + +impl FrecencyData { + /// Record a new usage of this command. + pub fn record_use(&mut self, timestamp: i64) { + self.count += 1; + if timestamp > self.last_used { + self.last_used = timestamp; + } + } + + /// Compute frecency score based on count and recency. + /// + /// Uses a decay function where more recent commands score higher. + /// The formula balances frequency (how often) with recency (how recent). + /// + /// Multipliers allow tuning the relative weights: + /// - `recency_mul`: Multiplier for recency score (default: 1.0) + /// - `frequency_mul`: Multiplier for frequency score (default: 1.0) + /// + /// A multiplier of 0.0 disables that component, 1.0 is unchanged, 2.0 doubles weight. + /// Values like 0.5 reduce weight by half, 1.5 increases by 50%, etc. + #[instrument(level = tracing::Level::TRACE, name = "index_frecency_compute")] + pub fn compute(&self, now: i64, recency_mul: f64, frequency_mul: f64) -> u32 { + if self.count == 0 { + return 0; + } + + // Time-based decay: score decreases as time passes + let age_seconds = (now - self.last_used).max(0) as u64; + let age_hours = age_seconds / 3600; + + // Decay factor: recent commands get higher scores + // - Last hour: multiplier ~1.0 + // - Last day: multiplier ~0.5 + // - Last week: multiplier ~0.1 + // - Older: multiplier approaches 0 + let recency_score: f64 = match age_hours { + 0 => 100.0, + 1..=6 => 90.0, + 7..=24 => 70.0, + 25..=72 => 50.0, + 73..=168 => 30.0, + 169..=720 => 15.0, + _ => 5.0, + }; + + // Frequency boost: more uses = higher score (with diminishing returns) + let frequency_score = ((self.count as f64).ln() * 20.0).min(100.0); + + // Apply multipliers and combine scores, then round to u32 + ((recency_score * recency_mul) + (frequency_score * frequency_mul)).round() as u32 + } +} + +/// Data for a unique command. +pub struct CommandData { + /// History ID of the most recent invocation (16-byte UUID). + most_recent_id: [u8; 16], + /// Timestamp of the most recent invocation. + most_recent_timestamp: i64, + /// Pre-computed global frecency. + pub global_frecency: FrecencyData, + + // Pre-computed indexes for O(1) filter lookups + // Using HashSet instead of DashSet since CommandData lives inside DashMap (already synchronized) + /// All directories where this command has been run (interned keys). + directories: HashSet<Spur>, + /// All hostnames where this command has been run (interned keys). + hosts: HashSet<Spur>, + /// All sessions where this command has been run (as 16-byte UUIDs). + sessions: HashSet<[u8; 16]>, +} + +impl CommandData { + /// Create a new CommandData from a history entry. + /// Returns None if the history entry has invalid UUIDs. + pub fn new(history: &History, interner: &ThreadedRodeo) -> Option<Self> { + let history_id = parse_uuid_bytes(&history.id.0)?; + let session = parse_uuid_bytes(&history.session)?; + let timestamp = history.timestamp.unix_timestamp(); + + let dir_key = interner.get_or_intern(with_trailing_slash(&history.cwd)); + let host_key = interner.get_or_intern(&history.hostname); + + let mut directories = HashSet::new(); + directories.insert(dir_key); + + let mut hosts = HashSet::new(); + hosts.insert(host_key); + + let mut sessions = HashSet::new(); + sessions.insert(session); + + let mut global_frecency = FrecencyData::default(); + global_frecency.record_use(timestamp); + + Some(Self { + most_recent_id: history_id, + most_recent_timestamp: timestamp, + global_frecency, + directories, + hosts, + sessions, + }) + } + + /// Add an invocation from a history entry. + /// Returns false if the history entry has invalid UUIDs. + pub fn add_invocation(&mut self, history: &History, interner: &ThreadedRodeo) -> bool { + let Some(history_id) = parse_uuid_bytes(&history.id.0) else { + return false; + }; + let Some(session) = parse_uuid_bytes(&history.session) else { + return false; + }; + + let timestamp = history.timestamp.unix_timestamp(); + + // Update global frecency + self.global_frecency.record_use(timestamp); + + // Update pre-computed indexes for O(1) filter lookups + let dir_key = interner.get_or_intern(with_trailing_slash(&history.cwd)); + self.directories.insert(dir_key); + self.hosts.insert(interner.get_or_intern(&history.hostname)); + self.sessions.insert(session); + + // Update most recent if this invocation is newer + if timestamp > self.most_recent_timestamp { + self.most_recent_id = history_id; + self.most_recent_timestamp = timestamp; + } + + true + } + + /// Get the most recent history ID for this command. + pub fn most_recent_id(&self) -> String { + format_uuid_bytes(&self.most_recent_id) + } + + /// Check if any invocation matches a directory filter (exact match). + /// O(1) lookup using pre-computed index. + pub fn has_invocation_in_dir(&self, dir: &str, interner: &ThreadedRodeo) -> bool { + interner + .get(dir) + .is_some_and(|spur| self.directories.contains(&spur)) + } + + /// Check if any invocation matches a directory prefix (workspace/git root). + /// O(n) where n = number of unique directories for this command. + pub fn has_invocation_in_workspace(&self, prefix: &str, interner: &ThreadedRodeo) -> bool { + self.directories + .iter() + .any(|&spur| interner.resolve(&spur).starts_with(prefix)) + } + + /// Check if any invocation matches a hostname. + /// O(1) lookup using pre-computed index. + pub fn has_invocation_on_host(&self, hostname: &str, interner: &ThreadedRodeo) -> bool { + interner + .get(hostname) + .is_some_and(|spur| self.hosts.contains(&spur)) + } + + /// Check if any invocation matches a session. + /// O(1) lookup using pre-computed index. + pub fn has_invocation_in_session(&self, session: &str) -> bool { + parse_uuid_bytes(session).is_some_and(|bytes| self.sessions.contains(&bytes)) + } +} + +/// Filter mode for search queries. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum IndexFilterMode { + /// No filtering - search all commands. + Global, + /// Filter to commands run in a specific directory. + Directory(String), + /// Filter to commands run in a workspace (directory prefix). + Workspace(String), + /// Filter to commands run on a specific host. + Host(String), + /// Filter to commands run in a specific session. + Session(String), +} + +/// Context for search queries. +#[derive(Debug, Clone, Default)] +pub struct QueryContext { + pub cwd: Option<String>, + pub git_root: Option<String>, + pub hostname: Option<String>, + pub session_id: Option<String>, +} + +/// Shareable frecency map: command -> frecency score. +/// Wrapped in Arc for zero-copy sharing with scorer callbacks. +type FrecencyMap = Arc<HashMap<Arc<str>, u32>>; + +/// A deduplicated search index with frecency-based ranking. +/// +/// Commands are stored by their text, with metadata about all invocations. +/// Nucleo handles fuzzy matching, while frecency is computed via scorer callback. +/// +/// Global frecency is precomputed by a background task and used for scoring. +/// If frecency data is not available, search still works but without frecency ranking; +/// although this should never happen due to precomputing the frecency map. +pub struct SearchIndex { + /// Map from command text to command data. + /// Using DashMap for concurrent read/write access, wrapped in Arc for sharing with scorer. + /// Keys are Arc<str> to enable zero-copy sharing with frecency_map. + commands: Arc<DashMap<Arc<str>, CommandData>>, + /// Nucleo fuzzy matcher - items are command strings. + nucleo: RwLock<Nucleo<String>>, + /// Injector for adding new commands to Nucleo. + injector: Injector<String>, + /// Precomputed global frecency map. Updated by background task. + frecency_map: RwLock<Option<FrecencyMap>>, + /// String interner for deduplicating cwd, hostname, and directory paths. + interner: Arc<ThreadedRodeo>, +} + +impl SearchIndex { + /// Create a new empty search index. + pub fn new() -> Self { + let nucleo_config = nucleo::Config::DEFAULT; + // Single column for command text + let nucleo = Nucleo::<String>::new(nucleo_config, Arc::new(|| {}), None, 1); + let injector = nucleo.injector(); + + Self { + commands: Arc::new(DashMap::new()), + nucleo: RwLock::new(nucleo), + injector, + frecency_map: RwLock::new(None), + interner: Arc::new(ThreadedRodeo::new()), + } + } + + /// Add a history entry to the index. + /// + /// If the command already exists, updates its invocation data. + /// If it's a new command, adds it to both the map and Nucleo. + pub fn add_history(&self, history: &History) { + let command = history.command.as_str(); + + // DashMap with Arc<str> keys can be looked up with &str via Borrow trait + if let Some(mut entry) = self.commands.get_mut(command) { + // Existing command - just update invocations + entry.add_invocation(history, &self.interner); + } else { + // New command - create Arc<str> once and share it + let Some(data) = CommandData::new(history, &self.interner) else { + return; // Invalid UUIDs, skip this entry + }; + let command_arc: Arc<str> = command.into(); + self.commands.insert(Arc::clone(&command_arc), data); + // Nucleo still needs String (unavoidable copy for fuzzy matching) + self.injector.push(command_arc.to_string(), |cmd, cols| { + cols[0] = cmd.clone().into(); + }); + } + // Note: frecency_map is rebuilt by background task, not invalidated here + } + + /// Add multiple history entries to the index. + pub fn add_histories(&self, histories: &[History]) { + for history in histories { + self.add_history(history); + } + } + + /// Get the number of unique commands in the index. + pub fn command_count(&self) -> usize { + self.commands.len() + } + + /// Get the number of items in Nucleo (should match command_count). + pub async fn nucleo_item_count(&self) -> u32 { + self.nucleo.read().await.snapshot().item_count() + } + + /// Search for commands matching a query. + /// + /// Returns a list of history IDs (most recent invocation per command). + /// Uses precomputed global frecency for scoring if available. + #[instrument(skip_all, level = tracing::Level::TRACE, name = "index_search", fields(query = %query))] + pub async fn search( + &self, + query: &str, + filter_mode: IndexFilterMode, + _context: &QueryContext, + limit: u32, + ) -> Vec<String> { + let mut nucleo = self.nucleo.write().await; + + // Get precomputed frecency map (may be None if not yet computed) + let frecency_map = self.frecency_map.read().await.clone(); + + // Build filter based on mode + let filter = self.build_filter(&filter_mode); + nucleo.set_filter(filter); + + // Build scorer from precomputed frecency (or None if not available) + let scorer = Self::build_scorer(frecency_map); + nucleo.set_scorer(scorer); + + // Update pattern + nucleo.pattern.reparse( + 0, + query, + pattern::CaseMatching::Smart, + pattern::Normalization::Smart, + false, + ); + + tracing::span!(Level::TRACE, "index_search_tick").in_scope(|| { + // Tick until complete + while nucleo.tick(10).running {} + }); + + // Collect results + let snapshot = nucleo.snapshot(); + let matched_count = snapshot.matched_item_count().min(limit); + + tracing::span!(Level::TRACE, "index_search_results").in_scope(|| { + snapshot + .matched_items(..matched_count) + .filter_map(|item| { + let cmd = item.data; + // DashMap<Arc<str>, _>::get accepts &str via Borrow trait + self.commands + .get(cmd.as_str()) + .map(|data| data.most_recent_id()) + }) + .collect() + }) + } + + /// Rebuild the global frecency map. + /// + /// This should be called by a background task periodically. + /// The map is used for scoring search results. + /// + /// Uses multipliers from search settings: + /// - `recency_score_multiplier`: Weight for recency component + /// - `frequency_score_multiplier`: Weight for frequency component + /// - `frecency_score_multiplier`: Overall multiplier for final score + #[instrument(skip_all, level = tracing::Level::DEBUG, name = "rebuild_frecency")] + pub async fn rebuild_frecency(&self, search_settings: &Search) { + let now = OffsetDateTime::now_utc().unix_timestamp(); + let mut frecency_map: HashMap<Arc<str>, u32> = HashMap::new(); + + // Clamp multipliers to non-negative values to prevent broken frecency ranking + // (negative values would produce unexpected results when cast to u32) + let recency_mul = search_settings.recency_score_multiplier.max(0.0); + let frequency_mul = search_settings.frequency_score_multiplier.max(0.0); + let frecency_mul = search_settings.frecency_score_multiplier.max(0.0); + + for entry in self.commands.iter() { + let frecency = entry + .global_frecency + .compute(now, recency_mul, frequency_mul); + // Apply overall frecency multiplier and round to u32 + let frecency = (frecency as f64 * frecency_mul).round() as u32; + // Arc::clone is cheap - just increments reference count + frecency_map.insert(Arc::clone(entry.key()), frecency); + } + + *self.frecency_map.write().await = Some(Arc::new(frecency_map)); + } + + /// Build filter predicate for the given mode. + fn build_filter(&self, mode: &IndexFilterMode) -> Option<nucleo::Filter<String>> { + // For Global mode, no filter needed + if matches!(mode, IndexFilterMode::Global) { + return None; + } + + // Pre-compute which commands pass the filter + // Use HashSet<String> for the short-lived filter (simpler than Arc lookup) + let passing_commands: Arc<HashSet<String>> = { + let mut set = HashSet::new(); + for entry in self.commands.iter() { + let passes = match mode { + IndexFilterMode::Global => unreachable!(), + IndexFilterMode::Directory(dir) => { + entry.has_invocation_in_dir(dir, &self.interner) + } + IndexFilterMode::Workspace(prefix) => { + entry.has_invocation_in_workspace(prefix, &self.interner) + } + IndexFilterMode::Host(hostname) => { + entry.has_invocation_on_host(hostname, &self.interner) + } + IndexFilterMode::Session(session) => entry.has_invocation_in_session(session), + }; + if passes { + // Convert Arc<str> to String for filter lookup + set.insert(entry.key().to_string()); + } + } + Arc::new(set) + }; + + Some(Arc::new(move |cmd: &String| passing_commands.contains(cmd))) + } + + /// Build scorer from precomputed frecency map. + /// + /// Returns None if frecency map is not available (search still works, just without frecency ranking). + fn build_scorer(frecency_map: Option<FrecencyMap>) -> Option<nucleo::Scorer<String>> { + let map = frecency_map?; + Some(Arc::new(move |cmd: &String, fuzzy_score: u32| { + // HashMap<Arc<str>, _>::get accepts &str via Borrow trait + let frecency = map.get(cmd.as_str()).copied().unwrap_or(0); + fuzzy_score + frecency + })) + } +} + +impl Default for SearchIndex { + fn default() -> Self { + Self::new() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use time::macros::datetime; + + fn make_history(command: &str, cwd: &str, timestamp: OffsetDateTime) -> History { + History::import() + .timestamp(timestamp) + .command(command) + .cwd(cwd) + .build() + .into() + } + + #[test] + fn frecency_data_compute() { + let now = 1000000i64; + + // Recent command (with default multipliers of 1.0) + let recent = FrecencyData { + count: 5, + last_used: now - 60, // 1 minute ago + }; + assert!(recent.compute(now, 1.0, 1.0) > 100); // High score + + // Old command + let old = FrecencyData { + count: 5, + last_used: now - 86400 * 30, // 30 days ago + }; + assert!(old.compute(now, 1.0, 1.0) < recent.compute(now, 1.0, 1.0)); + + // Frequently used old command + let frequent_old = FrecencyData { + count: 100, + last_used: now - 86400 * 7, // 1 week ago + }; + // Should still have decent score due to frequency + assert!(frequent_old.compute(now, 1.0, 1.0) > 50); + } + + #[test] + fn frecency_data_compute_with_multipliers() { + let now = 1000000i64; + + let data = FrecencyData { + count: 5, + last_used: now - 60, // 1 minute ago (recency_score = 100) + }; + + // Default multipliers (1.0, 1.0) + let default_score = data.compute(now, 1.0, 1.0); + + // Double recency weight + let double_recency = data.compute(now, 2.0, 1.0); + assert!(double_recency > default_score); + + // Double frequency weight + let double_frequency = data.compute(now, 1.0, 2.0); + assert!(double_frequency > default_score); + + // Zero out recency (only frequency counts) + let no_recency = data.compute(now, 0.0, 1.0); + assert!(no_recency < default_score); + + // Zero out frequency (only recency counts) + let no_frequency = data.compute(now, 1.0, 0.0); + assert!(no_frequency < default_score); + + // Zero both (should be zero) + let no_score = data.compute(now, 0.0, 0.0); + assert_eq!(no_score, 0); + + // Fractional multipliers + let half_recency = data.compute(now, 0.5, 1.0); + assert!(half_recency < default_score); + assert!(half_recency > no_recency); + + // 1.5x multiplier + let boost_recency = data.compute(now, 1.5, 1.0); + assert!(boost_recency > default_score); + assert!(boost_recency < double_recency); + } + + #[test] + fn command_data_add_invocation() { + let interner = ThreadedRodeo::new(); + + let (dir1, dir2) = if cfg!(windows) { + ("C:\\Users\\User\\project", "C:\\Users\\User\\other") + } else { + ("/home/user/project", "/home/user/other") + }; + + let history1 = make_history("git status", dir1, datetime!(2024-01-01 10:00 UTC)); + let history2 = make_history("git status", dir2, datetime!(2024-01-01 12:00 UTC)); + + let mut data = CommandData::new(&history1, &interner).unwrap(); + assert_eq!(data.global_frecency.count, 1); + let id1 = data.most_recent_id(); + + data.add_invocation(&history2, &interner); + assert_eq!(data.global_frecency.count, 2); + + // Most recent ID should update to history2 (newer timestamp) + let id2 = data.most_recent_id(); + assert_ne!(id1, id2); + } + + #[test] + fn command_data_filters() { + let interner = ThreadedRodeo::new(); + + let (dir1, dir2) = if cfg!(windows) { + ("C:\\Users\\User\\project", "C:\\Users\\User\\other") + } else { + ("/home/user/project", "/home/user/other") + }; + + let h1 = make_history("git status", dir1, datetime!(2024-01-01 10:00 UTC)); + let h2 = make_history("git status", dir2, datetime!(2024-01-01 12:00 UTC)); + + let mut data = CommandData::new(&h1, &interner).unwrap(); + data.add_invocation(&h2, &interner); + + let (check1, check2, check3) = if cfg!(windows) { + ( + with_trailing_slash("C:\\Users\\User\\project"), + with_trailing_slash("C:\\Users\\User\\other"), + with_trailing_slash("C:\\Users\\User\\missing"), + ) + } else { + ( + with_trailing_slash("/home/user/project"), + with_trailing_slash("/home/user/other"), + with_trailing_slash("/home/user/missing"), + ) + }; + + assert!(data.has_invocation_in_dir(&check1, &interner)); + assert!(data.has_invocation_in_dir(&check2, &interner)); + assert!(!data.has_invocation_in_dir(&check3, &interner)); + + let (check1, check2, check3) = if cfg!(windows) { + ( + with_trailing_slash("C:\\Users\\User"), + with_trailing_slash("C:\\Users"), + with_trailing_slash("C:\\Users\\User\\var"), + ) + } else { + ( + with_trailing_slash("/home/user"), + with_trailing_slash("/home"), + with_trailing_slash("/var"), + ) + }; + + assert!(data.has_invocation_in_workspace(&check1, &interner)); + assert!(data.has_invocation_in_workspace(&check2, &interner)); + assert!(!data.has_invocation_in_workspace(&check3, &interner)); + } + + #[tokio::test] + async fn search_index_add_and_search() { + let index = SearchIndex::new(); + + let h1 = make_history( + "git status", + "/home/user/project", + datetime!(2024-01-01 10:00 UTC), + ); + let h2 = make_history( + "git commit -m 'test'", + "/home/user/project", + datetime!(2024-01-01 10:05 UTC), + ); + let h3 = make_history( + "ls -la", + "/home/user/other", + datetime!(2024-01-01 10:10 UTC), + ); + + index.add_history(&h1); + index.add_history(&h2); + index.add_history(&h3); + + assert_eq!(index.command_count(), 3); + + // Search for "git" - should match 2 commands + let results = index + .search("git", IndexFilterMode::Global, &QueryContext::default(), 10) + .await; + assert_eq!(results.len(), 2); + + // Search with directory filter + let results = index + .search( + "", + IndexFilterMode::Directory(with_trailing_slash("/home/user/project")), + &QueryContext::default(), + 10, + ) + .await; + assert_eq!(results.len(), 2); // git status and git commit + } +} diff --git a/crates/atuin-daemon/src/search/mod.rs b/crates/atuin-daemon/src/search/mod.rs new file mode 100644 index 00000000..4d261956 --- /dev/null +++ b/crates/atuin-daemon/src/search/mod.rs @@ -0,0 +1,11 @@ +//! Search module for the daemon gRPC search service. +//! +//! This module provides fuzzy search over command history using Nucleo. + +mod index; + +// Include the generated proto code +tonic::include_proto!("search"); + +// Re-export the service and index +pub use index::{IndexFilterMode, QueryContext, SearchIndex}; diff --git a/crates/atuin-daemon/src/server.rs b/crates/atuin-daemon/src/server.rs new file mode 100644 index 00000000..a11de612 --- /dev/null +++ b/crates/atuin-daemon/src/server.rs @@ -0,0 +1,164 @@ +use eyre::Result; + +use crate::components::history::HistoryGrpcService; +use crate::components::search::SearchGrpcService; +use crate::control::{ControlService, control_server::ControlServer}; +use crate::daemon::DaemonHandle; +use crate::history::history_server::HistoryServer; +use crate::search::search_server::SearchServer; + +use atuin_client::settings::Settings; + +/// Run the gRPC server with the given services. +/// +/// This starts the gRPC server in the background and returns immediately. +/// The server will shut down when a ShutdownRequested event is received. +#[cfg(unix)] +pub async fn run_grpc_server( + settings: Settings, + history_service: HistoryServer<HistoryGrpcService>, + search_service: SearchServer<SearchGrpcService>, + control_service: ControlServer<ControlService>, + handle: DaemonHandle, +) -> Result<()> { + use tokio::net::UnixListener; + use tokio_stream::wrappers::UnixListenerStream; + + let socket_path = settings.daemon.socket_path.clone(); + + let (uds, cleanup) = if cfg!(target_os = "linux") && settings.daemon.systemd_socket { + #[cfg(target_os = "linux")] + { + use eyre::{OptionExt, WrapErr}; + use std::os::unix::net::SocketAddr; + use std::path::PathBuf; + tracing::info!("getting systemd socket"); + let listener = listenfd::ListenFd::from_env() + .take_unix_listener(0)? + .ok_or_eyre("missing systemd socket")?; + listener.set_nonblocking(true)?; + let actual_path: Result<PathBuf, eyre::Report> = listener + .local_addr() + .context("getting systemd socket's path") + .and_then(|addr: SocketAddr| { + addr.as_pathname() + .ok_or_eyre("systemd socket missing path") + .map(|path: &std::path::Path| path.to_owned()) + }); + match actual_path { + Ok(actual_path) => { + tracing::info!("listening on systemd socket: {actual_path:?}"); + if actual_path != std::path::Path::new(&socket_path) { + tracing::warn!( + "systemd socket is not at configured client path: {socket_path:?}" + ); + } + } + Err(err) => { + tracing::warn!( + "could not detect systemd socket path, ensure that it's at the configured path: {socket_path:?}, error: {err:?}" + ); + } + } + (UnixListener::from_std(listener)?, false) + } + #[cfg(not(target_os = "linux"))] + unreachable!() + } else { + tracing::info!("listening on unix socket {socket_path:?}"); + (UnixListener::bind(socket_path.clone())?, true) + }; + + let uds_stream = UnixListenerStream::new(uds); + + // Create shutdown signal from daemon handle + let shutdown_signal = async move { + let mut rx = handle.subscribe(); + loop { + use crate::DaemonEvent; + + match rx.recv().await { + Ok(DaemonEvent::ShutdownRequested) => break, + Ok(_) => continue, + Err(_) => break, // Channel closed + } + } + if cleanup { + eprintln!("Removing socket..."); + if let Err(e) = std::fs::remove_file(&socket_path) + && e.kind() != std::io::ErrorKind::NotFound + { + eprintln!("failed to remove socket: {e}"); + } + } + eprintln!("Shutting down gRPC server..."); + }; + + // Spawn the server in the background + tokio::spawn(async move { + use tonic::transport::Server; + + if let Err(e) = Server::builder() + .add_service(history_service) + .add_service(search_service) + .add_service(control_service) + .serve_with_incoming_shutdown(uds_stream, shutdown_signal) + .await + { + tracing::error!("gRPC server error: {e}"); + } + }); + + Ok(()) +} + +/// Run the gRPC server with the given services (Windows/TCP version). +#[cfg(not(unix))] +pub async fn run_grpc_server( + settings: Settings, + history_service: HistoryServer<HistoryGrpcService>, + search_service: SearchServer<SearchGrpcService>, + control_service: ControlServer<ControlService>, + handle: DaemonHandle, +) -> Result<()> { + use tokio::net::TcpListener; + use tokio_stream::wrappers::TcpListenerStream; + use tonic::transport::Server; + + let port = settings.daemon.tcp_port; + let url = format!("127.0.0.1:{port}"); + let tcp = TcpListener::bind(&url).await?; + let tcp_stream = TcpListenerStream::new(tcp); + + tracing::info!("listening on tcp port {:?}", port); + + // Create shutdown signal from daemon handle + let shutdown_signal = async move { + use crate::DaemonEvent; + + let mut rx = handle.subscribe(); + loop { + match rx.recv().await { + Ok(DaemonEvent::ShutdownRequested) => break, + Ok(_) => continue, + Err(_) => break, // Channel closed + } + } + eprintln!("Shutting down gRPC server..."); + }; + + // Spawn the server in the background + tokio::spawn(async move { + if let Err(e) = Server::builder() + .add_service(history_service) + .add_service(search_service) + .add_service(control_service) + .serve_with_incoming_shutdown(tcp_stream, shutdown_signal) + .await + { + tracing::error!("gRPC server error: {e}"); + } + }); + + Ok(()) +} diff --git a/crates/atuin-daemon/tests/lifecycle.rs b/crates/atuin-daemon/tests/lifecycle.rs new file mode 100644 index 00000000..3b6952de --- /dev/null +++ b/crates/atuin-daemon/tests/lifecycle.rs @@ -0,0 +1,172 @@ +//! Integration tests for the daemon server lifecycle. +//! +//! Each test spins up a real gRPC server on a temporary unix socket, +//! connects a client, and exercises the daemon RPCs. + +#[cfg(unix)] +mod unix { + use std::time::Duration; + + use atuin_client::database::Sqlite; + use atuin_client::record::sqlite_store::SqliteStore; + use atuin_client::settings::{Settings, init_meta_config_for_testing}; + use atuin_daemon::client::HistoryClient; + use atuin_daemon::components::HistoryComponent; + use atuin_daemon::{Daemon, DaemonHandle}; + use tempfile::TempDir; + use tokio::net::UnixListener; + use tokio_stream::wrappers::UnixListenerStream; + use tonic::transport::Server; + + /// Spins up a daemon server on a temp socket and returns a connected client, + /// the daemon handle (for shutdown), and the temp dir (must be held to keep paths alive). + async fn start_test_daemon() -> (HistoryClient, DaemonHandle, TempDir) { + let tmp = tempfile::tempdir().unwrap(); + + let db_path = tmp.path().join("history.db"); + let record_path = tmp.path().join("records.db"); + let key_path = tmp.path().join("key"); + let socket_path = tmp.path().join("test.sock"); + let meta_path = tmp.path().join("meta.db"); + + // Initialize the meta store config for testing (required for Settings::host_id()) + init_meta_config_for_testing(meta_path.to_str().unwrap(), 5.0); + + // Build settings with test paths + let settings: Settings = Settings::builder() + .expect("could not build settings builder") + .set_override("db_path", db_path.to_str().unwrap()) + .expect("failed to set db_path") + .set_override("record_store_path", record_path.to_str().unwrap()) + .expect("failed to set record_store_path") + .set_override("key_path", key_path.to_str().unwrap()) + .expect("failed to set key_path") + .set_override("daemon.socket_path", socket_path.to_str().unwrap()) + .expect("failed to set socket_path") + .set_override("meta.db_path", meta_path.to_str().unwrap()) + .expect("failed to set meta.db_path") + .build() + .expect("could not build settings") + .try_deserialize() + .expect("could not deserialize settings"); + + // Create databases + let history_db = Sqlite::new(&db_path, 5.0).await.unwrap(); + let store = SqliteStore::new(&record_path, 5.0).await.unwrap(); + + // Create the history component and get its gRPC service + let history_component = HistoryComponent::new(); + let history_service = history_component.grpc_service(); + + // Build and start the daemon + let mut daemon = Daemon::builder(settings) + .store(store) + .history_db(history_db) + .component(history_component) + .build() + .await + .unwrap(); + + let handle = daemon.handle(); + + // Start components (this initializes the history component with the handle) + daemon.start_components().await.unwrap(); + + // Start the gRPC server + let uds = UnixListener::bind(&socket_path).unwrap(); + let stream = UnixListenerStream::new(uds); + + let server_handle = handle.clone(); + tokio::spawn(async move { + let mut rx = server_handle.subscribe(); + Server::builder() + .add_service(history_service) + .serve_with_incoming_shutdown(stream, async move { + loop { + match rx.recv().await { + Ok(atuin_daemon::DaemonEvent::ShutdownRequested) => break, + Ok(_) => continue, + Err(_) => break, + } + } + }) + .await + .unwrap(); + }); + + // Spawn the daemon event loop in the background + tokio::spawn(async move { + daemon.run_event_loop().await.unwrap(); + }); + + // Give the server a moment to bind. + tokio::time::sleep(Duration::from_millis(50)).await; + + let client = HistoryClient::new(socket_path.to_string_lossy().to_string()) + .await + .unwrap(); + + (client, handle, tmp) + } + + #[tokio::test] + async fn test_status() { + let (mut client, _handle, _tmp) = start_test_daemon().await; + + let status = client.status().await.unwrap(); + assert!(status.healthy); + assert_eq!(status.version, env!("CARGO_PKG_VERSION")); + assert_eq!(status.protocol, 1); + assert!(status.pid > 0); + } + + #[tokio::test] + async fn test_start_end_history() { + use atuin_client::history::History; + + let (mut client, _handle, _tmp) = start_test_daemon().await; + + let history = History::daemon() + .timestamp(time::OffsetDateTime::now_utc()) + .command("echo hello".to_string()) + .cwd("/tmp".to_string()) + .session("test-session".to_string()) + .hostname("test-host".to_string()) + .build() + .into(); + + let start_reply = client.start_history(history).await.unwrap(); + assert!(!start_reply.id.is_empty()); + + let end_reply = client + .end_history(start_reply.id, 1_000_000, 0) + .await + .unwrap(); + assert!(!end_reply.id.is_empty()); + } + + #[tokio::test] + async fn test_end_unknown_history_fails() { + let (mut client, _handle, _tmp) = start_test_daemon().await; + + let result = client + .end_history("nonexistent-id".to_string(), 1000, 0) + .await; + assert!(result.is_err()); + } + + #[tokio::test] + async fn test_shutdown() { + let (mut client, _handle, _tmp) = start_test_daemon().await; + + let accepted = client.shutdown().await.unwrap(); + assert!(accepted); + + // Give server time to shut down. + tokio::time::sleep(Duration::from_millis(100)).await; + + // Subsequent calls should fail since the server is gone. + let result = client.status().await; + assert!(result.is_err()); + } +} diff --git a/crates/atuin-dotfiles/Cargo.toml b/crates/atuin-dotfiles/Cargo.toml new file mode 100644 index 00000000..3f3ae2d0 --- /dev/null +++ b/crates/atuin-dotfiles/Cargo.toml @@ -0,0 +1,25 @@ +[package] +name = "atuin-dotfiles" +description = "The dotfiles crate for Atuin" +edition = "2024" +version = { workspace = true } + +authors.workspace = true +rust-version.workspace = true +license.workspace = true +homepage.workspace = true +repository.workspace = true +readme.workspace = true + +# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html + +[dependencies] +atuin-common = { path = "../atuin-common", version = "18.13.2" } +atuin-client = { path = "../atuin-client", version = "18.13.2" } + +eyre = { workspace = true } +tokio = { workspace = true } +rmp = { version = "0.8.14" } +rand = { workspace = true } +serde = { workspace = true } +crypto_secretbox = "0.1.1" diff --git a/crates/atuin-dotfiles/src/lib.rs b/crates/atuin-dotfiles/src/lib.rs new file mode 100644 index 00000000..74daf8ef --- /dev/null +++ b/crates/atuin-dotfiles/src/lib.rs @@ -0,0 +1,2 @@ +pub mod shell; +pub mod store; diff --git a/crates/atuin-dotfiles/src/shell.rs b/crates/atuin-dotfiles/src/shell.rs new file mode 100644 index 00000000..73a9ce8c --- /dev/null +++ b/crates/atuin-dotfiles/src/shell.rs @@ -0,0 +1,241 @@ +use eyre::{Result, ensure, eyre}; +use rmp::{decode, encode}; +use serde::Serialize; + +use atuin_common::shell::{Shell, ShellError}; + +use crate::store::AliasStore; + +pub mod bash; +pub mod fish; +pub mod powershell; +pub mod xonsh; +pub mod zsh; + +#[derive(Debug, Clone, PartialEq, Eq, Serialize)] +pub struct Alias { + pub name: String, + pub value: String, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize)] +pub struct Var { + pub name: String, + pub value: String, + + // False? This is a _shell var_ + // True? This is an _env var_ + pub export: bool, +} + +impl Var { + /// Serialize into the given vec + /// This is intended to be called by the store + pub fn serialize(&self, output: &mut Vec<u8>) -> Result<()> { + encode::write_array_len(output, 3)?; // 3 fields + + encode::write_str(output, self.name.as_str())?; + encode::write_str(output, self.value.as_str())?; + encode::write_bool(output, self.export)?; + + Ok(()) + } + + pub fn deserialize(bytes: &mut decode::Bytes) -> Result<Self> { + fn error_report<E: std::fmt::Debug>(err: E) -> eyre::Report { + eyre!("{err:?}") + } + + let nfields = decode::read_array_len(bytes).map_err(error_report)?; + + ensure!( + nfields == 3, + "too many entries in v0 dotfiles env create record, got {}, expected {}", + nfields, + 3 + ); + + let bytes = bytes.remaining_slice(); + + let (key, bytes) = decode::read_str_from_slice(bytes).map_err(error_report)?; + let (value, bytes) = decode::read_str_from_slice(bytes).map_err(error_report)?; + + let mut bytes = decode::Bytes::new(bytes); + let export = decode::read_bool(&mut bytes).map_err(error_report)?; + + ensure!( + bytes.remaining_slice().is_empty(), + "trailing bytes in encoded dotfiles env record, malformed" + ); + + Ok(Var { + name: key.to_owned(), + value: value.to_owned(), + export, + }) + } +} + +pub fn parse_alias(line: &str) -> Option<Alias> { + // consider the fact we might be importing a fish alias + // 'alias' output + // fish: alias foo bar + // posix: foo=bar + + let is_fish = line.split(' ').next().unwrap_or("") == "alias"; + + let parts: Vec<&str> = if is_fish { + line.split(' ') + .enumerate() + .filter_map(|(n, i)| if n == 0 { None } else { Some(i) }) + .collect() + } else { + line.split('=').collect() + }; + + if parts.len() <= 1 { + return None; + } + + let mut parts = parts.iter().map(|s| s.to_string()); + + let name = parts.next().unwrap(); + + let remaining = if is_fish { + parts.collect::<Vec<String>>().join(" ") + } else { + parts.collect::<Vec<String>>().join("=") + }; + + Some(Alias { + name, + value: remaining.trim().to_string(), + }) +} + +pub fn existing_aliases(shell: Option<Shell>) -> Result<Vec<Alias>, ShellError> { + let shell = if let Some(shell) = shell { + shell + } else { + Shell::current() + }; + + // this only supports posix-y shells atm + if !shell.is_posixish() { + return Err(ShellError::NotSupported); + } + + // This will return a list of aliases, each on its own line + // They will be in the form foo=bar + let aliases = shell.run_interactive(["alias"])?; + + let aliases: Vec<Alias> = aliases.lines().filter_map(parse_alias).collect(); + + Ok(aliases) +} + +/// Import aliases from the current shell +/// This will not import aliases already in the store +/// Returns aliases that were set +pub async fn import_aliases(store: &AliasStore) -> Result<Vec<Alias>> { + let shell_aliases = existing_aliases(None)?; + let store_aliases = store.aliases().await?; + + let mut res = Vec::new(); + + for alias in shell_aliases { + // O(n), but n is small, and imports infrequent + // can always make a map + if store_aliases.contains(&alias) { + continue; + } + + res.push(alias.clone()); + store.set(&alias.name, &alias.value).await?; + } + + Ok(res) +} + +#[cfg(test)] +mod tests { + use crate::shell::{Alias, parse_alias}; + + #[test] + fn test_parse_simple_alias() { + let alias = super::parse_alias("foo=bar").expect("failed to parse alias"); + assert_eq!(alias.name, "foo"); + assert_eq!(alias.value, "bar"); + } + + #[test] + fn test_parse_quoted_alias() { + let alias = super::parse_alias("emacs='TERM=xterm-24bits emacs -nw'") + .expect("failed to parse alias"); + + assert_eq!(alias.name, "emacs"); + assert_eq!(alias.value, "'TERM=xterm-24bits emacs -nw'"); + + let git_alias = super::parse_alias("gwip='git add -A; git rm $(git ls-files --deleted) 2> /dev/null; git commit --no-verify --no-gpg-sign --message \"--wip-- [skip ci]\"'").expect("failed to parse alias"); + assert_eq!(git_alias.name, "gwip"); + assert_eq!( + git_alias.value, + "'git add -A; git rm $(git ls-files --deleted) 2> /dev/null; git commit --no-verify --no-gpg-sign --message \"--wip-- [skip ci]\"'" + ); + } + + #[test] + fn test_parse_quoted_alias_equals() { + let alias = super::parse_alias("emacs='TERM=xterm-24bits emacs -nw --foo=bar'") + .expect("failed to parse alias"); + assert_eq!(alias.name, "emacs"); + assert_eq!(alias.value, "'TERM=xterm-24bits emacs -nw --foo=bar'"); + } + + #[test] + fn test_parse_fish() { + let alias = super::parse_alias("alias foo bar").expect("failed to parse alias"); + assert_eq!(alias.name, "foo"); + assert_eq!(alias.value, "bar"); + + let alias = + super::parse_alias("alias x 'exa --icons --git --classify --group-directories-first'") + .expect("failed to parse alias"); + + assert_eq!(alias.name, "x"); + assert_eq!( + alias.value, + "'exa --icons --git --classify --group-directories-first'" + ); + } + + #[test] + fn test_parse_with_fortune() { + // Because we run the alias command in an interactive subshell + // there may be other output. + // Ensure that the parser can handle it + // Annoyingly not all aliases are picked up all the time if we use + // a non-interactive subshell. Boo. + let shell = " +/ In a consumer society there are \\ +| inevitably two kinds of slaves: the | +| prisoners of addiction and the | +\\ prisoners of envy. / + ------------------------------------- + \\ ^__^ + \\ (oo)\\_______ + (__)\\ )\\/\\ + ||----w | + || || +emacs='TERM=xterm-24bits emacs -nw --foo=bar' +k=kubectl +"; + + let aliases: Vec<Alias> = shell.lines().filter_map(parse_alias).collect(); + assert_eq!(aliases[0].name, "emacs"); + assert_eq!(aliases[0].value, "'TERM=xterm-24bits emacs -nw --foo=bar'"); + + assert_eq!(aliases[1].name, "k"); + assert_eq!(aliases[1].value, "kubectl"); + } +} diff --git a/crates/atuin-dotfiles/src/shell/bash.rs b/crates/atuin-dotfiles/src/shell/bash.rs new file mode 100644 index 00000000..2b9b4c88 --- /dev/null +++ b/crates/atuin-dotfiles/src/shell/bash.rs @@ -0,0 +1,68 @@ +use std::path::PathBuf; + +use crate::store::{AliasStore, var::VarStore}; + +async fn cached_aliases(path: PathBuf, store: &AliasStore) -> String { + match tokio::fs::read_to_string(path).await { + Ok(aliases) => aliases, + Err(r) => { + // we failed to read the file for some reason, but the file does exist + // fallback to generating new aliases on the fly + + store.posix().await.unwrap_or_else(|e| { + format!("echo 'Atuin: failed to read and generate aliases: \n{r}\n{e}'",) + }) + } + } +} + +async fn cached_vars(path: PathBuf, store: &VarStore) -> String { + match tokio::fs::read_to_string(path).await { + Ok(vars) => vars, + Err(r) => { + // we failed to read the file for some reason, but the file does exist + // fallback to generating new vars on the fly + + store.posix().await.unwrap_or_else(|e| { + format!("echo 'Atuin: failed to read and generate vars: \n{r}\n{e}'",) + }) + } + } +} + +/// Return bash dotfile config +/// +/// Do not return an error. We should not prevent the shell from starting. +/// +/// In the worst case, Atuin should not function but the shell should start correctly. +/// +/// While currently this only returns aliases, it will be extended to also return other synced dotfiles +pub async fn alias_config(store: &AliasStore) -> String { + // First try to read the cached config + let aliases = atuin_common::utils::dotfiles_cache_dir().join("aliases.bash"); + + if aliases.exists() { + return cached_aliases(aliases, store).await; + } + + if let Err(e) = store.build().await { + return format!("echo 'Atuin: failed to generate aliases: {e}'"); + } + + cached_aliases(aliases, store).await +} + +pub async fn var_config(store: &VarStore) -> String { + // First try to read the cached config + let vars = atuin_common::utils::dotfiles_cache_dir().join("vars.bash"); + + if vars.exists() { + return cached_vars(vars, store).await; + } + + if let Err(e) = store.build().await { + return format!("echo 'Atuin: failed to generate vars: {e}'"); + } + + cached_vars(vars, store).await +} diff --git a/crates/atuin-dotfiles/src/shell/fish.rs b/crates/atuin-dotfiles/src/shell/fish.rs new file mode 100644 index 00000000..6d472f67 --- /dev/null +++ b/crates/atuin-dotfiles/src/shell/fish.rs @@ -0,0 +1,69 @@ +// Configuration for fish +use std::path::PathBuf; + +use crate::store::{AliasStore, var::VarStore}; + +async fn cached_aliases(path: PathBuf, store: &AliasStore) -> String { + match tokio::fs::read_to_string(path).await { + Ok(aliases) => aliases, + Err(r) => { + // we failed to read the file for some reason, but the file does exist + // fallback to generating new aliases on the fly + + store.posix().await.unwrap_or_else(|e| { + format!("echo 'Atuin: failed to read and generate aliases: \n{r}\n{e}'",) + }) + } + } +} + +async fn cached_vars(path: PathBuf, store: &VarStore) -> String { + match tokio::fs::read_to_string(path).await { + Ok(vars) => vars, + Err(r) => { + // we failed to read the file for some reason, but the file does exist + // fallback to generating new vars on the fly + + store.posix().await.unwrap_or_else(|e| { + format!("echo 'Atuin: failed to read and generate vars: \n{r}\n{e}'",) + }) + } + } +} + +/// Return fish dotfile config +/// +/// Do not return an error. We should not prevent the shell from starting. +/// +/// In the worst case, Atuin should not function but the shell should start correctly. +/// +/// While currently this only returns aliases, it will be extended to also return other synced dotfiles +pub async fn alias_config(store: &AliasStore) -> String { + // First try to read the cached config + let aliases = atuin_common::utils::dotfiles_cache_dir().join("aliases.fish"); + + if aliases.exists() { + return cached_aliases(aliases, store).await; + } + + if let Err(e) = store.build().await { + return format!("echo 'Atuin: failed to generate aliases: {e}'"); + } + + cached_aliases(aliases, store).await +} + +pub async fn var_config(store: &VarStore) -> String { + // First try to read the cached config + let vars = atuin_common::utils::dotfiles_cache_dir().join("vars.fish"); + + if vars.exists() { + return cached_vars(vars, store).await; + } + + if let Err(e) = store.build().await { + return format!("echo 'Atuin: failed to generate vars: {e}'"); + } + + cached_vars(vars, store).await +} diff --git a/crates/atuin-dotfiles/src/shell/powershell.rs b/crates/atuin-dotfiles/src/shell/powershell.rs new file mode 100644 index 00000000..1daee28b --- /dev/null +++ b/crates/atuin-dotfiles/src/shell/powershell.rs @@ -0,0 +1,169 @@ +use crate::shell::{Alias, Var}; +use crate::store::{AliasStore, var::VarStore}; +use std::path::PathBuf; + +async fn cached_aliases(path: PathBuf, store: &AliasStore) -> String { + match tokio::fs::read_to_string(path).await { + Ok(aliases) => aliases, + Err(r) => { + // we failed to read the file for some reason, but the file does exist + // fallback to generating new aliases on the fly + + store.powershell().await.unwrap_or_else(|e| { + format!("echo 'Atuin: failed to read and generate aliases: \n{r}\n{e}'",) + }) + } + } +} + +async fn cached_vars(path: PathBuf, store: &VarStore) -> String { + match tokio::fs::read_to_string(path).await { + Ok(vars) => vars, + Err(r) => { + // we failed to read the file for some reason, but the file does exist + // fallback to generating new vars on the fly + + store.powershell().await.unwrap_or_else(|e| { + format!("echo 'Atuin: failed to read and generate vars: \n{r}\n{e}'",) + }) + } + } +} + +/// Return powershell dotfile config +/// +/// Do not return an error. We should not prevent the shell from starting. +/// +/// In the worst case, Atuin should not function but the shell should start correctly. +/// +/// While currently this only returns aliases, it will be extended to also return other synced dotfiles +pub async fn alias_config(store: &AliasStore) -> String { + // First try to read the cached config + let aliases = atuin_common::utils::dotfiles_cache_dir().join("aliases.ps1"); + + if aliases.exists() { + return cached_aliases(aliases, store).await; + } + + if let Err(e) = store.build().await { + return format!("echo 'Atuin: failed to generate aliases: {e}'"); + } + + cached_aliases(aliases, store).await +} + +pub async fn var_config(store: &VarStore) -> String { + // First try to read the cached config + let vars = atuin_common::utils::dotfiles_cache_dir().join("vars.ps1"); + + if vars.exists() { + return cached_vars(vars, store).await; + } + + if let Err(e) = store.build().await { + return format!("echo 'Atuin: failed to generate vars: {e}'"); + } + + cached_vars(vars, store).await +} + +pub fn format_alias(alias: &Alias) -> String { + // Set-Alias doesn't support adding implicit arguments, so use a function. + // See https://github.com/PowerShell/PowerShell/issues/12962 + + let mut result = secure_command(&format!( + "function {} {{\n {}{} @args\n}}", + alias.name, + if alias.value.starts_with(['"', '\'']) { + "& " + } else { + "" + }, + alias.value + )); + + // This makes the file layout prettier + result.insert(0, '\n'); + result +} + +pub fn format_var(var: &Var) -> String { + secure_command(&format!( + "${}{} = '{}'", + if var.export { "env:" } else { "" }, + var.name, + var.value.replace("'", "''") + )) +} + +/// Wraps the given command in an Invoke-Expression to ensure the outer script is not halted +/// if the inner command contains a syntax error. +fn secure_command(command: &str) -> String { + format!( + "Invoke-Expression -ErrorAction Continue -Command '{}'\n", + command.replace("'", "''") + ) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn aliases() { + assert_eq!( + format_alias(&Alias { + name: "gp".to_string(), + value: "git push".to_string(), + }), + "\n".to_string() + + &secure_command( + "function gp { + git push @args +}" + ) + ); + + assert_eq!( + format_alias(&Alias { + name: "spc".to_string(), + value: "\"path with spaces\" arg".to_string(), + }), + "\n".to_string() + + &secure_command( + "function spc { + & \"path with spaces\" arg @args +}" + ) + ); + } + + #[test] + fn vars() { + assert_eq!( + format_var(&Var { + name: "FOO".to_owned(), + value: "bar 'baz'".to_owned(), + export: true, + }), + secure_command("$env:FOO = 'bar ''baz'''") + ); + + assert_eq!( + format_var(&Var { + name: "TEST".to_owned(), + value: "1".to_owned(), + export: false, + }), + secure_command("$TEST = '1'") + ); + } + + #[test] + fn invoke_expression() { + assert_eq!( + secure_command("echo 'foo'"), + "Invoke-Expression -ErrorAction Continue -Command 'echo ''foo'''\n" + ) + } +} diff --git a/crates/atuin-dotfiles/src/shell/xonsh.rs b/crates/atuin-dotfiles/src/shell/xonsh.rs new file mode 100644 index 00000000..1e56fc1d --- /dev/null +++ b/crates/atuin-dotfiles/src/shell/xonsh.rs @@ -0,0 +1,68 @@ +use std::path::PathBuf; + +use crate::store::{AliasStore, var::VarStore}; + +async fn cached_aliases(path: PathBuf, store: &AliasStore) -> String { + match tokio::fs::read_to_string(path).await { + Ok(aliases) => aliases, + Err(r) => { + // we failed to read the file for some reason, but the file does exist + // fallback to generating new aliases on the fly + + store.xonsh().await.unwrap_or_else(|e| { + format!("echo 'Atuin: failed to read and generate aliases: \n{r}\n{e}'",) + }) + } + } +} + +async fn cached_vars(path: PathBuf, store: &VarStore) -> String { + match tokio::fs::read_to_string(path).await { + Ok(vars) => vars, + Err(r) => { + // we failed to read the file for some reason, but the file does exist + // fallback to generating new vars on the fly + + store.xonsh().await.unwrap_or_else(|e| { + format!("echo 'Atuin: failed to read and generate vars: \n{r}\n{e}'",) + }) + } + } +} + +/// Return xonsh dotfile config +/// +/// Do not return an error. We should not prevent the shell from starting. +/// +/// In the worst case, Atuin should not function but the shell should start correctly. +/// +/// While currently this only returns aliases, it will be extended to also return other synced dotfiles +pub async fn alias_config(store: &AliasStore) -> String { + // First try to read the cached config + let aliases = atuin_common::utils::dotfiles_cache_dir().join("aliases.xsh"); + + if aliases.exists() { + return cached_aliases(aliases, store).await; + } + + if let Err(e) = store.build().await { + return format!("echo 'Atuin: failed to generate aliases: {e}'"); + } + + cached_aliases(aliases, store).await +} + +pub async fn var_config(store: &VarStore) -> String { + // First try to read the cached config + let vars = atuin_common::utils::dotfiles_cache_dir().join("vars.xsh"); + + if vars.exists() { + return cached_vars(vars, store).await; + } + + if let Err(e) = store.build().await { + return format!("echo 'Atuin: failed to generate vars: {e}'"); + } + + cached_vars(vars, store).await +} diff --git a/crates/atuin-dotfiles/src/shell/zsh.rs b/crates/atuin-dotfiles/src/shell/zsh.rs new file mode 100644 index 00000000..117e9403 --- /dev/null +++ b/crates/atuin-dotfiles/src/shell/zsh.rs @@ -0,0 +1,68 @@ +use std::path::PathBuf; + +use crate::store::{AliasStore, var::VarStore}; + +async fn cached_aliases(path: PathBuf, store: &AliasStore) -> String { + match tokio::fs::read_to_string(path).await { + Ok(aliases) => aliases, + Err(r) => { + // we failed to read the file for some reason, but the file does exist + // fallback to generating new aliases on the fly + + store.posix().await.unwrap_or_else(|e| { + format!("echo 'Atuin: failed to read and generate aliases: \n{r}\n{e}'",) + }) + } + } +} + +async fn cached_vars(path: PathBuf, store: &VarStore) -> String { + match tokio::fs::read_to_string(path).await { + Ok(aliases) => aliases, + Err(r) => { + // we failed to read the file for some reason, but the file does exist + // fallback to generating new vars on the fly + + store.posix().await.unwrap_or_else(|e| { + format!("echo 'Atuin: failed to read and generate aliases: \n{r}\n{e}'",) + }) + } + } +} + +/// Return zsh dotfile config +/// +/// Do not return an error. We should not prevent the shell from starting. +/// +/// In the worst case, Atuin should not function but the shell should start correctly. +/// +/// While currently this only returns aliases, it will be extended to also return other synced dotfiles +pub async fn alias_config(store: &AliasStore) -> String { + // First try to read the cached config + let aliases = atuin_common::utils::dotfiles_cache_dir().join("aliases.zsh"); + + if aliases.exists() { + return cached_aliases(aliases, store).await; + } + + if let Err(e) = store.build().await { + return format!("echo 'Atuin: failed to generate aliases: {e}'"); + } + + cached_aliases(aliases, store).await +} + +pub async fn var_config(store: &VarStore) -> String { + // First try to read the cached config + let vars = atuin_common::utils::dotfiles_cache_dir().join("vars.zsh"); + + if vars.exists() { + return cached_vars(vars, store).await; + } + + if let Err(e) = store.build().await { + return format!("echo 'Atuin: failed to generate aliases: {e}'"); + } + + cached_vars(vars, store).await +} diff --git a/crates/atuin-dotfiles/src/store.rs b/crates/atuin-dotfiles/src/store.rs new file mode 100644 index 00000000..17597065 --- /dev/null +++ b/crates/atuin-dotfiles/src/store.rs @@ -0,0 +1,421 @@ +use std::collections::BTreeMap; + +use atuin_client::record::sqlite_store::SqliteStore; +// Sync aliases +// This will be noticeable similar to the kv store, though I expect the two shall diverge +// While we will support a range of shell config, I'd rather have a larger number of small records +// + stores, rather than one mega config store. +use atuin_common::record::{DecryptedData, Host, HostId}; +use atuin_common::utils::unquote; +use eyre::{Result, bail, ensure, eyre}; + +use atuin_client::record::encryption::PASETO_V4; +use atuin_client::record::store::Store; + +use crate::shell::Alias; + +const CONFIG_SHELL_ALIAS_VERSION: &str = "v0"; +const CONFIG_SHELL_ALIAS_TAG: &str = "config-shell-alias"; +const CONFIG_SHELL_ALIAS_FIELD_MAX_LEN: usize = 20000; // 20kb max total len, way more than should be needed. + +mod alias; +pub mod var; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum AliasRecord { + Create(Alias), // create a full record + Delete(String), // delete by name +} + +impl AliasRecord { + pub fn serialize(&self) -> Result<DecryptedData> { + use rmp::encode; + + let mut output = vec![]; + + match self { + AliasRecord::Create(alias) => { + encode::write_u8(&mut output, 0)?; // create + encode::write_array_len(&mut output, 2)?; // 2 fields + + encode::write_str(&mut output, alias.name.as_str())?; + encode::write_str(&mut output, alias.value.as_str())?; + } + AliasRecord::Delete(name) => { + encode::write_u8(&mut output, 1)?; // delete + encode::write_array_len(&mut output, 1)?; // 1 field + + encode::write_str(&mut output, name.as_str())?; + } + } + + Ok(DecryptedData(output)) + } + + pub fn deserialize(data: &DecryptedData, version: &str) -> Result<Self> { + use rmp::decode; + + fn error_report<E: std::fmt::Debug>(err: E) -> eyre::Report { + eyre!("{err:?}") + } + + match version { + CONFIG_SHELL_ALIAS_VERSION => { + let mut bytes = decode::Bytes::new(&data.0); + + let record_type = decode::read_u8(&mut bytes).map_err(error_report)?; + + match record_type { + // create + 0 => { + let nfields = decode::read_array_len(&mut bytes).map_err(error_report)?; + ensure!( + nfields == 2, + "too many entries in v0 shell alias create record" + ); + + let bytes = bytes.remaining_slice(); + + let (key, bytes) = + decode::read_str_from_slice(bytes).map_err(error_report)?; + let (value, bytes) = + decode::read_str_from_slice(bytes).map_err(error_report)?; + + if !bytes.is_empty() { + bail!("trailing bytes in encoded shell alias record. malformed") + } + + Ok(AliasRecord::Create(Alias { + name: key.to_owned(), + value: value.to_owned(), + })) + } + + // delete + 1 => { + let nfields = decode::read_array_len(&mut bytes).map_err(error_report)?; + ensure!( + nfields == 1, + "too many entries in v0 shell alias delete record" + ); + + let bytes = bytes.remaining_slice(); + + let (key, bytes) = + decode::read_str_from_slice(bytes).map_err(error_report)?; + + if !bytes.is_empty() { + bail!("trailing bytes in encoded shell alias record. malformed") + } + + Ok(AliasRecord::Delete(key.to_owned())) + } + + n => { + bail!("unknown AliasRecord type {n}") + } + } + } + _ => { + bail!("unknown version {version:?}") + } + } + } +} + +#[derive(Debug, Clone)] +pub struct AliasStore { + pub store: SqliteStore, + pub host_id: HostId, + pub encryption_key: [u8; 32], +} + +impl AliasStore { + // will want to init the actual kv store when that is done + pub fn new(store: SqliteStore, host_id: HostId, encryption_key: [u8; 32]) -> AliasStore { + AliasStore { + store, + host_id, + encryption_key, + } + } + + pub async fn posix(&self) -> Result<String> { + let aliases = self.aliases().await?; + Ok(Self::format_posix(&aliases)) + } + + pub async fn xonsh(&self) -> Result<String> { + let aliases = self.aliases().await?; + Ok(Self::format_xonsh(&aliases)) + } + + pub async fn powershell(&self) -> Result<String> { + let aliases = self.aliases().await?; + Ok(Self::format_powershell(&aliases)) + } + + fn format_posix(aliases: &[Alias]) -> String { + let mut config = String::new(); + + for alias in aliases { + // If it's quoted, remove the quotes. If it's not quoted, do nothing. + let value = unquote(alias.value.as_str()).unwrap_or(alias.value.clone()); + + // we're about to quote it ourselves anyway! + config.push_str(&format!("alias {}='{}'\n", alias.name, value)); + } + + config + } + + fn format_xonsh(aliases: &[Alias]) -> String { + let mut config = String::new(); + + for alias in aliases { + config.push_str(&format!("aliases['{}'] ='{}'\n", alias.name, alias.value)); + } + + config + } + + fn format_powershell(aliases: &[Alias]) -> String { + let mut config = String::new(); + + for alias in aliases { + config.push_str(&crate::shell::powershell::format_alias(alias)); + } + + config + } + + pub async fn build(&self) -> Result<()> { + let dir = atuin_common::utils::dotfiles_cache_dir(); + tokio::fs::create_dir_all(dir.clone()).await?; + + let aliases = self.aliases().await?; + + // Build for all supported shells + let posix = Self::format_posix(&aliases); + let xonsh = Self::format_xonsh(&aliases); + let powershell = Self::format_powershell(&aliases); + + // All the same contents, maybe optimize in the future or perhaps there will be quirks + // per-shell + // I'd prefer separation atm + let zsh = dir.join("aliases.zsh"); + let bash = dir.join("aliases.bash"); + let fish = dir.join("aliases.fish"); + let xsh = dir.join("aliases.xsh"); + let ps1 = dir.join("aliases.ps1"); + + tokio::fs::write(zsh, &posix).await?; + tokio::fs::write(bash, &posix).await?; + tokio::fs::write(fish, &posix).await?; + tokio::fs::write(xsh, &xonsh).await?; + tokio::fs::write(ps1, &powershell).await?; + + Ok(()) + } + + pub async fn set(&self, name: &str, value: &str) -> Result<()> { + if name.len() + value.len() > CONFIG_SHELL_ALIAS_FIELD_MAX_LEN { + return Err(eyre!( + "alias record too large: max len {} bytes", + CONFIG_SHELL_ALIAS_FIELD_MAX_LEN + )); + } + + let record = AliasRecord::Create(Alias { + name: name.to_string(), + value: value.to_string(), + }); + + let bytes = record.serialize()?; + + let idx = self + .store + .last(self.host_id, CONFIG_SHELL_ALIAS_TAG) + .await? + .map_or(0, |entry| entry.idx + 1); + + let record = atuin_common::record::Record::builder() + .host(Host::new(self.host_id)) + .version(CONFIG_SHELL_ALIAS_VERSION.to_string()) + .tag(CONFIG_SHELL_ALIAS_TAG.to_string()) + .idx(idx) + .data(bytes) + .build(); + + self.store + .push(&record.encrypt::<PASETO_V4>(&self.encryption_key)) + .await?; + + // set mutates shell config, so build again + self.build().await?; + + Ok(()) + } + + pub async fn delete(&self, name: &str) -> Result<()> { + if name.len() > CONFIG_SHELL_ALIAS_FIELD_MAX_LEN { + return Err(eyre!( + "alias record too large: max len {} bytes", + CONFIG_SHELL_ALIAS_FIELD_MAX_LEN + )); + } + + let record = AliasRecord::Delete(name.to_string()); + + let bytes = record.serialize()?; + + let idx = self + .store + .last(self.host_id, CONFIG_SHELL_ALIAS_TAG) + .await? + .map_or(0, |entry| entry.idx + 1); + + let record = atuin_common::record::Record::builder() + .host(Host::new(self.host_id)) + .version(CONFIG_SHELL_ALIAS_VERSION.to_string()) + .tag(CONFIG_SHELL_ALIAS_TAG.to_string()) + .idx(idx) + .data(bytes) + .build(); + + self.store + .push(&record.encrypt::<PASETO_V4>(&self.encryption_key)) + .await?; + + // delete mutates shell config, so build again + self.build().await?; + + Ok(()) + } + + pub async fn aliases(&self) -> Result<Vec<Alias>> { + let mut build = BTreeMap::new(); + + // this is sorted, oldest to newest + let tagged = self.store.all_tagged(CONFIG_SHELL_ALIAS_TAG).await?; + + for record in tagged { + let version = record.version.clone(); + + let decrypted = match version.as_str() { + CONFIG_SHELL_ALIAS_VERSION => record.decrypt::<PASETO_V4>(&self.encryption_key)?, + version => bail!("unknown version {version:?}"), + }; + + let ar = AliasRecord::deserialize(&decrypted.data, version.as_str())?; + + match ar { + AliasRecord::Create(a) => { + build.insert(a.name.clone(), a); + } + AliasRecord::Delete(d) => { + build.remove(&d); + } + } + } + + Ok(build.into_values().collect()) + } +} + +#[cfg(test)] +pub(crate) fn test_local_timeout() -> f64 { + std::env::var("ATUIN_TEST_LOCAL_TIMEOUT") + .ok() + .and_then(|x| x.parse().ok()) + // this hardcoded value should be replaced by a simple way to get the + // default local_timeout of Settings if possible + .unwrap_or(2.0) +} + +#[cfg(test)] +mod tests { + use rand::rngs::OsRng; + + use atuin_client::record::sqlite_store::SqliteStore; + + use crate::shell::Alias; + + use super::{AliasRecord, AliasStore, CONFIG_SHELL_ALIAS_VERSION, test_local_timeout}; + use crypto_secretbox::{KeyInit, XSalsa20Poly1305}; + + #[test] + fn encode_decode() { + let record = Alias { + name: "k".to_owned(), + value: "kubectl".to_owned(), + }; + let record = AliasRecord::Create(record); + + let snapshot = [204, 0, 146, 161, 107, 167, 107, 117, 98, 101, 99, 116, 108]; + + let encoded = record.serialize().unwrap(); + let decoded = AliasRecord::deserialize(&encoded, CONFIG_SHELL_ALIAS_VERSION).unwrap(); + + assert_eq!(encoded.0, &snapshot); + assert_eq!(decoded, record); + } + + #[tokio::test] + async fn build_aliases() { + let store = SqliteStore::new(":memory:", test_local_timeout()) + .await + .unwrap(); + let key: [u8; 32] = XSalsa20Poly1305::generate_key(&mut OsRng).into(); + let host_id = atuin_common::record::HostId(atuin_common::utils::uuid_v7()); + + let alias = AliasStore::new(store, host_id, key); + + alias.set("k", "kubectl").await.unwrap(); + alias.set("gp", "git push").await.unwrap(); + alias + .set("kgap", "'kubectl get pods --all-namespaces'") + .await + .unwrap(); + + let mut aliases = alias.aliases().await.unwrap(); + + aliases.sort_by_key(|a| a.name.clone()); + + assert_eq!(aliases.len(), 3); + + assert_eq!( + aliases[0], + Alias { + name: String::from("gp"), + value: String::from("git push") + } + ); + + assert_eq!( + aliases[1], + Alias { + name: String::from("k"), + value: String::from("kubectl") + } + ); + + assert_eq!( + aliases[2], + Alias { + name: String::from("kgap"), + value: String::from("'kubectl get pods --all-namespaces'") + } + ); + + let build = alias.posix().await.expect("failed to build aliases"); + + assert_eq!( + build, + "alias gp='git push' +alias k='kubectl' +alias kgap='kubectl get pods --all-namespaces' +" + ) + } +} diff --git a/crates/atuin-dotfiles/src/store/alias.rs b/crates/atuin-dotfiles/src/store/alias.rs new file mode 100644 index 00000000..8b137891 --- /dev/null +++ b/crates/atuin-dotfiles/src/store/alias.rs @@ -0,0 +1 @@ + diff --git a/crates/atuin-dotfiles/src/store/var.rs b/crates/atuin-dotfiles/src/store/var.rs new file mode 100644 index 00000000..9d25b85d --- /dev/null +++ b/crates/atuin-dotfiles/src/store/var.rs @@ -0,0 +1,542 @@ +/// Store for shell vars +/// I should abstract this and reuse code between the alias/env stores +/// This is easier for now +/// Once I have two implementations, building a common base is much easier. +use std::collections::BTreeMap; + +use atuin_client::record::sqlite_store::SqliteStore; +use atuin_common::record::{DecryptedData, Host, HostId}; +use eyre::{Result, bail, ensure, eyre}; + +use atuin_client::record::encryption::PASETO_V4; +use atuin_client::record::store::Store; + +use crate::shell::Var; + +const DOTFILES_VAR_VERSION: &str = "v0"; +const DOTFILES_VAR_TAG: &str = "dotfiles-var"; +const DOTFILES_VAR_LEN: usize = 20000; // 20kb max total len, way more than should be needed. + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum VarRecord { + Create(Var), // create a full record + Delete(String), // delete by name +} + +impl VarRecord { + pub fn serialize(&self) -> Result<DecryptedData> { + use rmp::encode; + + let mut output = vec![]; + + match self { + VarRecord::Create(env) => { + encode::write_u8(&mut output, 0)?; // create + + env.serialize(&mut output)?; + } + VarRecord::Delete(env) => { + encode::write_u8(&mut output, 1)?; // delete + encode::write_array_len(&mut output, 1)?; // 1 field + + encode::write_str(&mut output, env.as_str())?; + } + } + + Ok(DecryptedData(output)) + } + + pub fn deserialize(data: &DecryptedData, version: &str) -> Result<Self> { + use rmp::decode; + + fn error_report<E: std::fmt::Debug>(err: E) -> eyre::Report { + eyre!("{err:?}") + } + + match version { + DOTFILES_VAR_VERSION => { + let mut bytes = decode::Bytes::new(&data.0); + + let record_type = decode::read_u8(&mut bytes).map_err(error_report)?; + + match record_type { + // create + 0 => { + let env = Var::deserialize(&mut bytes)?; + Ok(VarRecord::Create(env)) + } + + // delete + 1 => { + let nfields = decode::read_array_len(&mut bytes).map_err(error_report)?; + ensure!( + nfields == 1, + "too many entries in v0 dotfiles var delete record" + ); + + let bytes = bytes.remaining_slice(); + + let (key, bytes) = + decode::read_str_from_slice(bytes).map_err(error_report)?; + + if !bytes.is_empty() { + bail!("trailing bytes in encoded dotfiles var record. malformed") + } + + Ok(VarRecord::Delete(key.to_owned())) + } + + n => { + bail!("unknown Dotfiles var record type {n}") + } + } + } + _ => { + bail!("unknown version {version:?}") + } + } + } +} + +#[derive(Debug, Clone)] +pub struct VarStore { + pub store: SqliteStore, + pub host_id: HostId, + pub encryption_key: [u8; 32], +} + +impl VarStore { + // will want to init the actual kv store when that is done + pub fn new(store: SqliteStore, host_id: HostId, encryption_key: [u8; 32]) -> VarStore { + VarStore { + store, + host_id, + encryption_key, + } + } + + /// Escape a value for use in POSIX shells (bash, zsh) + /// This adds double quotes around the value and escapes any embedded double quotes + fn escape_posix_value(value: &str) -> String { + // If the value contains no special characters, we can use it unquoted + if value + .chars() + .all(|c| c.is_alphanumeric() || c == '_' || c == '-' || c == '/' || c == '.') + { + value.to_string() + } else { + // Otherwise, wrap in double quotes and escape any special characters + format!( + "\"{}\"", + value + .replace('\\', "\\\\") + .replace('"', "\\\"") + .replace('$', "\\$") + .replace('`', "\\`") + ) + } + } + + /// Escape a value for use in fish shell + /// Fish uses single quotes for literal strings, but we need to handle embedded single quotes + fn escape_fish_value(value: &str) -> String { + // If the value contains no special characters, we can use it unquoted + if value + .chars() + .all(|c| c.is_alphanumeric() || c == '_' || c == '-' || c == '/' || c == '.') + { + value.to_string() + } else { + // Use single quotes and escape any embedded single quotes + format!("'{}'", value.replace('\'', "\\'")) + } + } + + /// Escape a value for use in xonsh + /// Xonsh uses Python-style string literals + fn escape_xonsh_value(value: &str) -> String { + // If the value contains no special characters, we can use it unquoted + if value + .chars() + .all(|c| c.is_alphanumeric() || c == '_' || c == '-' || c == '/' || c == '.') + { + value.to_string() + } else { + // Use double quotes and escape appropriately for Python strings + format!("\"{}\"", value.replace('\\', "\\\\").replace('"', "\\\"")) + } + } + + pub async fn xonsh(&self) -> Result<String> { + let env = self.vars().await?; + Ok(Self::format_xonsh(&env)) + } + + pub async fn fish(&self) -> Result<String> { + let env = self.vars().await?; + Ok(Self::format_fish(&env)) + } + + pub async fn posix(&self) -> Result<String> { + let env = self.vars().await?; + Ok(Self::format_posix(&env)) + } + + pub async fn powershell(&self) -> Result<String> { + let env = self.vars().await?; + Ok(Self::format_powershell(&env)) + } + + fn format_xonsh(env: &[Var]) -> String { + let mut config = String::new(); + + for env in env { + let escaped_value = Self::escape_xonsh_value(&env.value); + config.push_str(&format!("${}={}\n", env.name, escaped_value)); + } + + config + } + + fn format_fish(env: &[Var]) -> String { + let mut config = String::new(); + + for env in env { + let escaped_value = Self::escape_fish_value(&env.value); + config.push_str(&format!("set -gx {} {}\n", env.name, escaped_value)); + } + + config + } + + fn format_posix(env: &[Var]) -> String { + let mut config = String::new(); + + for env in env { + let escaped_value = Self::escape_posix_value(&env.value); + if env.export { + config.push_str(&format!("export {}={}\n", env.name, escaped_value)); + } else { + config.push_str(&format!("{}={}\n", env.name, escaped_value)); + } + } + + config + } + + fn format_powershell(env: &[Var]) -> String { + let mut config = String::new(); + + for var in env { + config.push_str(&crate::shell::powershell::format_var(var)); + } + + config + } + + pub async fn build(&self) -> Result<()> { + let dir = atuin_common::utils::dotfiles_cache_dir(); + tokio::fs::create_dir_all(dir.clone()).await?; + + let env = self.vars().await?; + + // Build for all supported shells + let posix = Self::format_posix(&env); + let xonsh = Self::format_xonsh(&env); + let fsh = Self::format_fish(&env); + let powershell = Self::format_powershell(&env); + + // All the same contents, maybe optimize in the future or perhaps there will be quirks + // per-shell + // I'd prefer separation atm + let zsh = dir.join("vars.zsh"); + let bash = dir.join("vars.bash"); + let fish = dir.join("vars.fish"); + let xsh = dir.join("vars.xsh"); + let ps1 = dir.join("vars.ps1"); + + tokio::fs::write(zsh, &posix).await?; + tokio::fs::write(bash, &posix).await?; + tokio::fs::write(fish, &fsh).await?; + tokio::fs::write(xsh, &xonsh).await?; + tokio::fs::write(ps1, &powershell).await?; + + Ok(()) + } + + pub async fn set(&self, name: &str, value: &str, export: bool) -> Result<()> { + if name.len() + value.len() > DOTFILES_VAR_LEN { + return Err(eyre!( + "var record too large: max len {} bytes", + DOTFILES_VAR_LEN + )); + } + + let record = VarRecord::Create(Var { + name: name.to_string(), + value: value.to_string(), + export, + }); + + let bytes = record.serialize()?; + + let idx = self + .store + .last(self.host_id, DOTFILES_VAR_TAG) + .await? + .map_or(0, |entry| entry.idx + 1); + + let record = atuin_common::record::Record::builder() + .host(Host::new(self.host_id)) + .version(DOTFILES_VAR_VERSION.to_string()) + .tag(DOTFILES_VAR_TAG.to_string()) + .idx(idx) + .data(bytes) + .build(); + + self.store + .push(&record.encrypt::<PASETO_V4>(&self.encryption_key)) + .await?; + + // set mutates shell config, so build again + self.build().await?; + + Ok(()) + } + + pub async fn delete(&self, name: &str) -> Result<()> { + if name.len() > DOTFILES_VAR_LEN { + return Err(eyre!( + "var record too large: max len {} bytes", + DOTFILES_VAR_LEN, + )); + } + + let record = VarRecord::Delete(name.to_string()); + + let bytes = record.serialize()?; + + let idx = self + .store + .last(self.host_id, DOTFILES_VAR_TAG) + .await? + .map_or(0, |entry| entry.idx + 1); + + let record = atuin_common::record::Record::builder() + .host(Host::new(self.host_id)) + .version(DOTFILES_VAR_VERSION.to_string()) + .tag(DOTFILES_VAR_TAG.to_string()) + .idx(idx) + .data(bytes) + .build(); + + self.store + .push(&record.encrypt::<PASETO_V4>(&self.encryption_key)) + .await?; + + // delete mutates shell config, so build again + self.build().await?; + + Ok(()) + } + + pub async fn vars(&self) -> Result<Vec<Var>> { + let mut build = BTreeMap::new(); + + // this is sorted, oldest to newest + let tagged = self.store.all_tagged(DOTFILES_VAR_TAG).await?; + + for record in tagged { + let version = record.version.clone(); + + let decrypted = match version.as_str() { + DOTFILES_VAR_VERSION => record.decrypt::<PASETO_V4>(&self.encryption_key)?, + version => bail!("unknown version {version:?}"), + }; + + let ar = VarRecord::deserialize(&decrypted.data, version.as_str())?; + + match ar { + VarRecord::Create(a) => { + build.insert(a.name.clone(), a); + } + VarRecord::Delete(d) => { + build.remove(&d); + } + } + } + + Ok(build.into_values().collect()) + } +} + +#[cfg(test)] +mod tests { + use rand::rngs::OsRng; + + use atuin_client::record::sqlite_store::SqliteStore; + + use crate::{shell::Var, store::test_local_timeout}; + + use super::{DOTFILES_VAR_VERSION, VarRecord, VarStore}; + use crypto_secretbox::{KeyInit, XSalsa20Poly1305}; + + #[test] + fn encode_decode() { + let record = Var { + name: "BEEP".to_owned(), + value: "boop".to_owned(), + export: false, + }; + let record = VarRecord::Create(record); + + let snapshot = [ + 204, 0, 147, 164, 66, 69, 69, 80, 164, 98, 111, 111, 112, 194, + ]; + + let encoded = record.serialize().unwrap(); + let decoded = VarRecord::deserialize(&encoded, DOTFILES_VAR_VERSION).unwrap(); + + assert_eq!(encoded.0, &snapshot); + assert_eq!(decoded, record); + } + + #[test] + fn test_escape_posix_value() { + // Simple values should not be quoted + assert_eq!(VarStore::escape_posix_value("simple"), "simple"); + assert_eq!(VarStore::escape_posix_value("path/to/file"), "path/to/file"); + assert_eq!( + VarStore::escape_posix_value("value_with_underscores"), + "value_with_underscores" + ); + + // Values with spaces should be quoted + assert_eq!( + VarStore::escape_posix_value("hello world"), + "\"hello world\"" + ); + assert_eq!(VarStore::escape_posix_value("bar baz"), "\"bar baz\""); + + // Values with special characters should be quoted and escaped + assert_eq!( + VarStore::escape_posix_value("say \"hello\""), + "\"say \\\"hello\\\"\"" + ); + assert_eq!( + VarStore::escape_posix_value("path\\with\\backslashes"), + "\"path\\\\with\\\\backslashes\"" + ); + assert_eq!( + VarStore::escape_posix_value("say $hello"), + "\"say \\$hello\"" + ); + assert_eq!( + VarStore::escape_posix_value("see `example.md`"), + "\"see \\`example.md\\`\"" + ); + } + + #[test] + fn test_escape_fish_value() { + // Simple values should not be quoted + assert_eq!(VarStore::escape_fish_value("simple"), "simple"); + assert_eq!(VarStore::escape_fish_value("path/to/file"), "path/to/file"); + + // Values with spaces should be single-quoted + assert_eq!(VarStore::escape_fish_value("hello world"), "'hello world'"); + assert_eq!(VarStore::escape_fish_value("bar baz"), "'bar baz'"); + + // Values with single quotes should be escaped + assert_eq!(VarStore::escape_fish_value("don't"), "'don\\'t'"); + } + + #[test] + fn test_escape_xonsh_value() { + // Simple values should not be quoted + assert_eq!(VarStore::escape_xonsh_value("simple"), "simple"); + assert_eq!(VarStore::escape_xonsh_value("path/to/file"), "path/to/file"); + + // Values with spaces should be quoted + assert_eq!( + VarStore::escape_xonsh_value("hello world"), + "\"hello world\"" + ); + assert_eq!(VarStore::escape_xonsh_value("bar baz"), "\"bar baz\""); + + // Values with special characters should be quoted and escaped + assert_eq!( + VarStore::escape_xonsh_value("say \"hello\""), + "\"say \\\"hello\\\"\"" + ); + assert_eq!( + VarStore::escape_xonsh_value("path\\with\\backslashes"), + "\"path\\\\with\\\\backslashes\"" + ); + } + + #[tokio::test] + async fn build_vars() { + let store = SqliteStore::new(":memory:", test_local_timeout()) + .await + .unwrap(); + let key: [u8; 32] = XSalsa20Poly1305::generate_key(&mut OsRng).into(); + let host_id = atuin_common::record::HostId(atuin_common::utils::uuid_v7()); + + let env = VarStore::new(store, host_id, key); + + env.set("BEEP", "boop", false).await.unwrap(); + env.set("HOMEBREW_NO_AUTO_UPDATE", "1", true).await.unwrap(); + + let mut env_vars = env.vars().await.unwrap(); + + env_vars.sort_by_key(|a| a.name.clone()); + + assert_eq!(env_vars.len(), 2); + + assert_eq!( + env_vars[0], + Var { + name: String::from("BEEP"), + value: String::from("boop"), + export: false, + } + ); + + assert_eq!( + env_vars[1], + Var { + name: String::from("HOMEBREW_NO_AUTO_UPDATE"), + value: String::from("1"), + export: true, + } + ); + } + + #[tokio::test] + async fn test_var_generation_with_spaces() { + let store = SqliteStore::new(":memory:", test_local_timeout()) + .await + .unwrap(); + let key: [u8; 32] = XSalsa20Poly1305::generate_key(&mut OsRng).into(); + let host_id = atuin_common::record::HostId(atuin_common::utils::uuid_v7()); + + let env = VarStore::new(store, host_id, key); + + // Test the exact scenario from the bug report + env.set("FOO", "bar baz", true).await.unwrap(); + + let posix_output = env.posix().await.unwrap(); + let fish_output = env.fish().await.unwrap(); + let xonsh_output = env.xonsh().await.unwrap(); + + // POSIX should quote the value + assert_eq!(posix_output, "export FOO=\"bar baz\"\n"); + + // Fish should quote the value + assert_eq!(fish_output, "set -gx FOO 'bar baz'\n"); + + // Xonsh should quote the value + assert_eq!(xonsh_output, "$FOO=\"bar baz\"\n"); + } +} diff --git a/crates/atuin-hex/Cargo.toml b/crates/atuin-hex/Cargo.toml new file mode 100644 index 00000000..8a574a55 --- /dev/null +++ b/crates/atuin-hex/Cargo.toml @@ -0,0 +1,21 @@ +[package] +name = "atuin-hex" +edition = "2024" +description = "a terminal emulator for atuin" + +version = { workspace = true } +authors = { workspace = true } +rust-version = { workspace = true } +license = { workspace = true } +homepage = { workspace = true } +repository = { workspace = true } + +[dependencies] +clap = { workspace = true } + +[target.'cfg(all(unix, not(target_os = "illumos")))'.dependencies] +crossterm = { workspace = true } +eyre = { workspace = true } +portable-pty = "0.8" +signal-hook = "0.3" +vt100 = "0.15" diff --git a/crates/atuin-hex/src/lib.rs b/crates/atuin-hex/src/lib.rs new file mode 100644 index 00000000..ff37cfe3 --- /dev/null +++ b/crates/atuin-hex/src/lib.rs @@ -0,0 +1,465 @@ +pub mod osc133; + +use clap::{Args, Subcommand, ValueEnum}; + +#[derive(Subcommand, Debug)] +pub enum Cmd { + /// Print shell code to initialize atuin-hex on shell startup + Init(Init), +} + +#[derive(Args, Debug)] +pub struct Init { + /// Shell to generate init for. If omitted, attempt auto-detection + #[arg(value_enum)] + shell: Option<Shell>, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq, ValueEnum)] +#[value(rename_all = "lower")] +#[allow(clippy::enum_variant_names, clippy::doc_markdown)] +enum Shell { + /// Zsh setup + Zsh, + /// Bash setup + Bash, + /// Fish setup + Fish, +} + +impl Shell { + fn as_str(self) -> &'static str { + match self { + Self::Bash => "bash", + Self::Zsh => "zsh", + Self::Fish => "fish", + } + } +} + +impl Init { + fn run(self) -> Result<(), String> { + let shell = detect_shell(self.shell)?; + let script = render_init(shell); + print!("{script}"); + Ok(()) + } +} + +pub fn run(cmd: Option<Cmd>) { + match cmd { + Some(Cmd::Init(init)) => { + if let Err(err) = init.run() { + eprintln!("atuin hex: {err}"); + std::process::exit(1); + } + } + None => app::main(), + } +} + +fn detect_shell(cli_shell: Option<Shell>) -> Result<Shell, String> { + if let Some(shell) = cli_shell { + return Ok(shell); + } + + if let Ok(shell) = std::env::var("ATUIN_SHELL") + && let Some(shell) = shell_from_name(&shell) + { + return Ok(shell); + } + + if let Ok(shell) = std::env::var("SHELL") + && let Some(shell) = shell_from_name(&shell) + { + return Ok(shell); + } + + Err( + "could not detect a supported shell. Please specify one explicitly: bash, zsh, or fish" + .to_string(), + ) +} + +fn shell_from_name(name: &str) -> Option<Shell> { + let shell = name + .trim() + .rsplit('/') + .next() + .unwrap_or(name) + .trim_start_matches('-') + .to_ascii_lowercase(); + + match shell.as_str() { + "bash" => Some(Shell::Bash), + "zsh" => Some(Shell::Zsh), + "fish" => Some(Shell::Fish), + _ => None, + } +} + +fn init_command(shell: Shell) -> String { + format!("atuin init {}", shell.as_str()) +} + +fn render_init(shell: Shell) -> String { + let init_command = init_command(shell); + + match shell { + Shell::Bash | Shell::Zsh => format!( + r#"if [[ "$-" == *i* ]] && [[ -t 0 ]] && [[ -t 1 ]]; then + _atuin_hex_tmux_current="${{TMUX:-}}" + _atuin_hex_tmux_previous="${{ATUIN_HEX_TMUX:-}}" + + if [[ -z "${{ATUIN_HEX_ACTIVE:-}}" ]] || [[ "$_atuin_hex_tmux_current" != "$_atuin_hex_tmux_previous" ]]; then + export ATUIN_HEX_ACTIVE=1 + export ATUIN_HEX_TMUX="$_atuin_hex_tmux_current" + exec atuin hex + fi + + unset _atuin_hex_tmux_current _atuin_hex_tmux_previous +fi + +eval "$({init_command})" +"# + ), + Shell::Fish => format!( + r#"if status is-interactive; and test -t 0; and test -t 1 + set -l _atuin_hex_tmux_current "" + if set -q TMUX + set _atuin_hex_tmux_current "$TMUX" + end + + set -l _atuin_hex_tmux_previous "" + if set -q ATUIN_HEX_TMUX + set _atuin_hex_tmux_previous "$ATUIN_HEX_TMUX" + end + + if not set -q ATUIN_HEX_ACTIVE + set -gx ATUIN_HEX_ACTIVE 1 + set -gx ATUIN_HEX_TMUX "$_atuin_hex_tmux_current" + exec atuin hex + else if test "$_atuin_hex_tmux_current" != "$_atuin_hex_tmux_previous" + set -gx ATUIN_HEX_ACTIVE 1 + set -gx ATUIN_HEX_TMUX "$_atuin_hex_tmux_current" + exec atuin hex + end +end + +{init_command} | source +"# + ), + } +} + +#[cfg(any(not(unix), target_os = "illumos"))] +mod app { + pub(crate) fn main() { + eprintln!("atuin hex currently supports unix platforms excluding illumos"); + std::process::exit(1); + } +} + +#[cfg(all(unix, not(target_os = "illumos")))] +mod app { + use std::io::{Read, Write}; + use std::os::unix::net::UnixListener; + use std::sync::mpsc; + + use crossterm::terminal; + use portable_pty::{CommandBuilder, PtySize, native_pty_system}; + + enum ParserMsg { + Data(Vec<u8>), + Resize { rows: u16, cols: u16 }, + ScreenRequest(mpsc::Sender<Vec<u8>>), + } + + pub(crate) fn main() { + if let Err(e) = run() { + let _ = terminal::disable_raw_mode(); + eprintln!("atuin hex: {e:#}"); + std::process::exit(1); + } + } + + fn socket_path() -> std::path::PathBuf { + let dir = std::env::temp_dir(); + dir.join(format!("atuin-hex-{}.sock", std::process::id())) + } + + /// Wire format written to the Unix socket: + /// + /// ```text + /// [rows: u16 BE][cols: u16 BE][cursor_row: u16 BE][cursor_col: u16 BE] + /// [row_0_len: u32 BE][row_0_bytes...] + /// [row_1_len: u32 BE][row_1_bytes...] + /// ... + /// ``` + /// + /// Each row's bytes come from `screen.rows_formatted(0, cols)` and contain + /// pre-built ANSI escape sequences. The client can write them directly to + /// stdout without needing its own vt100 parser. + fn encode_screen(parser: &vt100::Parser) -> Vec<u8> { + let screen = parser.screen(); + let (rows, cols) = screen.size(); + let (cursor_row, cursor_col) = screen.cursor_position(); + + let mut buf: Vec<u8> = Vec::with_capacity(256 + (rows as usize * cols as usize)); + buf.extend_from_slice(&rows.to_be_bytes()); + buf.extend_from_slice(&cols.to_be_bytes()); + buf.extend_from_slice(&cursor_row.to_be_bytes()); + buf.extend_from_slice(&cursor_col.to_be_bytes()); + + for row_bytes in screen.rows_formatted(0, cols) { + let len = row_bytes.len() as u32; + buf.extend_from_slice(&len.to_be_bytes()); + buf.extend_from_slice(&row_bytes); + } + + buf + } + + fn handle_parser_msg(parser: &mut vt100::Parser, msg: ParserMsg) { + match msg { + ParserMsg::Data(data) => parser.process(&data), + ParserMsg::Resize { rows, cols } => parser.set_size(rows, cols), + ParserMsg::ScreenRequest(reply_tx) => { + let _ = reply_tx.send(encode_screen(parser)); + } + } + } + + fn run() -> eyre::Result<()> { + let (cols, rows) = terminal::size()?; + + let pty_system = native_pty_system(); + let pair = pty_system + .openpty(PtySize { + rows, + cols, + pixel_width: 0, + pixel_height: 0, + }) + .map_err(|e| eyre::eyre!("{e:#}"))?; + + // Set up socket path and expose it to child processes + let sock_path = socket_path(); + // Clean up any stale socket from a previous crash + let _ = std::fs::remove_file(&sock_path); + + let mut cmd = CommandBuilder::new_default_prog(); + cmd.cwd(std::env::current_dir()?); + cmd.env("ATUIN_HEX_SOCKET", sock_path.as_os_str()); + + let mut child = pair + .slave + .spawn_command(cmd) + .map_err(|e| eyre::eyre!("{e:#}"))?; + + // Close slave side in parent process + drop(pair.slave); + + let mut pty_reader = pair + .master + .try_clone_reader() + .map_err(|e| eyre::eyre!("{e:#}"))?; + let mut pty_writer = pair + .master + .take_writer() + .map_err(|e| eyre::eyre!("{e:#}"))?; + + // Channel: stdout/sigwinch/socket threads -> parser thread (bounded, non-blocking send) + let (msg_tx, msg_rx) = mpsc::sync_channel::<ParserMsg>(64); + + // --- Parser thread --- + // Maintains a persistent vt100::Parser fed bytes as they arrive. + // On screen request: reads current state directly (no replay). + std::thread::spawn(move || { + let mut parser = vt100::Parser::new(rows, cols, 0); + + loop { + // Block until at least one message arrives + let first = match msg_rx.recv() { + Ok(msg) => msg, + Err(_) => break, + }; + + handle_parser_msg(&mut parser, first); + + // Drain all remaining pending messages so the parser stays + // caught up during high-throughput bursts (e.g. `cat bigfile`). + // The channel holds at most 64 items, so this is bounded. + while let Ok(msg) = msg_rx.try_recv() { + handle_parser_msg(&mut parser, msg); + } + } + }); + + // --- Socket server thread --- + // Listens on Unix socket; on connection, requests screen state from parser thread. + { + let sock_path_clone = sock_path.clone(); + let screen_tx = msg_tx.clone(); + std::thread::spawn(move || { + let listener = match UnixListener::bind(&sock_path_clone) { + Ok(l) => l, + Err(e) => { + eprintln!("atuin hex: failed to bind socket: {e}"); + return; + } + }; + + for stream in listener.incoming() { + let mut stream = match stream { + Ok(s) => s, + Err(_) => break, + }; + + let (reply_tx, reply_rx) = mpsc::channel(); + if screen_tx.send(ParserMsg::ScreenRequest(reply_tx)).is_err() { + break; + } + if let Ok(data) = reply_rx.recv() { + let _ = stream.write_all(&data); + let _ = stream.flush(); + } + } + }); + } + + // Handle terminal resize via SIGWINCH + { + use signal_hook::consts::SIGWINCH; + use signal_hook::iterator::Signals; + + let master = pair.master; + let resize_tx = msg_tx.clone(); + let mut signals = Signals::new([SIGWINCH])?; + + std::thread::spawn(move || { + for _ in signals.forever() { + if let Ok((cols, rows)) = terminal::size() { + let _ = master.resize(PtySize { + rows, + cols, + pixel_width: 0, + pixel_height: 0, + }); + let _ = resize_tx.try_send(ParserMsg::Resize { rows, cols }); + } + } + }); + } + + terminal::enable_raw_mode()?; + + // PTY -> stdout (with OSC 133 parsing + buffer feed) + let stdout_thread = std::thread::spawn(move || { + let mut stdout = std::io::stdout(); + let mut parser = crate::osc133::Parser::new(); + let mut buf = [0u8; 8192]; + loop { + match pty_reader.read(&mut buf) { + Ok(0) | Err(_) => break, + Ok(n) => { + parser.push(&buf[..n], |_event| { + // Zone transitions are tracked inside the parser. + // Callers can query parser.zone() after push. + }); + + // Feed bytes to the shadow parser. Drops on backpressure — + // the screen snapshot may be stale during bursts, but + // self-corrects once output settles. + let _ = msg_tx.try_send(ParserMsg::Data(buf[..n].to_vec())); + + if stdout.write_all(&buf[..n]).is_err() { + break; + } + let _ = stdout.flush(); + } + } + } + }); + + // stdin -> PTY + std::thread::spawn(move || { + let mut stdin = std::io::stdin(); + let mut buf = [0u8; 8192]; + loop { + match stdin.read(&mut buf) { + Ok(0) | Err(_) => break, + Ok(n) => { + if pty_writer.write_all(&buf[..n]).is_err() { + break; + } + } + } + } + }); + + let status = child.wait()?; + let _ = stdout_thread.join(); + + let _ = terminal::disable_raw_mode(); + + // Clean up socket file + let _ = std::fs::remove_file(&sock_path); + + std::process::exit(process_exit_code(status.exit_code())); + } + + fn process_exit_code(code: u32) -> i32 { + i32::try_from(code).unwrap_or(1) + } + + #[cfg(test)] + mod tests { + use super::process_exit_code; + + #[test] + fn process_exit_code_preserves_valid_values() { + assert_eq!(process_exit_code(0), 0); + assert_eq!(process_exit_code(127), 127); + assert_eq!(process_exit_code(i32::MAX as u32), i32::MAX); + } + + #[test] + fn process_exit_code_defaults_when_out_of_range() { + assert_eq!(process_exit_code(i32::MAX as u32 + 1), 1); + } + } +} + +#[cfg(test)] +mod tests { + use super::{Shell, init_command, render_init, shell_from_name}; + + #[test] + fn shell_from_name_handles_paths() { + assert_eq!(shell_from_name("/bin/zsh"), Some(Shell::Zsh)); + assert_eq!(shell_from_name("/usr/local/bin/bash"), Some(Shell::Bash)); + assert_eq!(shell_from_name("fish"), Some(Shell::Fish)); + } + + #[test] + fn init_command_is_bootstrap_only() { + let command = init_command(Shell::Zsh); + assert_eq!(command, "atuin init zsh"); + } + + #[test] + fn posix_init_uses_exec_and_tmux_guard() { + let script = render_init(Shell::Bash); + assert!(script.contains("exec atuin hex")); + assert!(script.contains("ATUIN_HEX_TMUX")); + assert!(script.contains("eval \"$(atuin init bash)\"")); + } + + #[test] + fn fish_init_uses_source() { + let script = render_init(Shell::Fish); + assert!(script.contains("exec atuin hex")); + assert!(script.contains("atuin init fish | source")); + } +} diff --git a/crates/atuin-hex/src/osc133.rs b/crates/atuin-hex/src/osc133.rs new file mode 100644 index 00000000..d6ee1220 --- /dev/null +++ b/crates/atuin-hex/src/osc133.rs @@ -0,0 +1,657 @@ +//! Streaming parser for OSC 133 (FinalTerm semantic prompt) escape sequences. +//! +//! OSC 133 marks four regions of a shell interaction: +//! +//! | Marker | Meaning | +//! |--------|--------------------------------------| +//! | A | Prompt is about to be printed | +//! | B | Prompt ended — command input begins | +//! | C | Command submitted — output begins | +//! | D[;n] | Command finished with exit code *n* | +//! +//! The wire format is `ESC ] 133 ; <cmd> [; <params>] ST` where ST is either +//! BEL (0x07) or ESC \ (0x1B 0x5C). +//! +//! # Design goals +//! +//! * **Zero-copy** — the parser observes the byte stream without buffering or +//! modifying it. +//! * **Zero-alloc** — after construction no heap allocation occurs. +//! * **Non-blocking** — [`Parser::push`] processes whatever bytes are available +//! and returns immediately. +//! * **Transparent** — the caller is responsible for forwarding bytes to their +//! destination; the parser only emits [`Event`]s through a callback. + +/// Events emitted when an OSC 133 marker is detected. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum Event { + /// `ESC ] 133 ; A ST` — the shell is about to display its prompt. + PromptStart, + /// `ESC ] 133 ; B ST` — the prompt has ended; the user may type a command. + CommandStart, + /// `ESC ] 133 ; C ST` — the command has been submitted for execution. + CommandExecuted, + /// `ESC ] 133 ; D [; <exit_code>] ST` — command output is complete. + CommandFinished { + /// The exit code reported after the `;`, if present and valid. + exit_code: Option<i32>, + }, +} + +/// The current semantic zone as determined by the most recent OSC 133 marker. +#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)] +#[allow(dead_code)] +pub enum Zone { + /// No marker seen yet, or after a `D` marker (between commands). + #[default] + Unknown, + /// Between `A` and `B` — the shell is rendering its prompt. + Prompt, + /// Between `B` and `C` — the user is editing a command line. + Input, + /// Between `C` and `D` — command output is being produced. + Output, +} + +// --------------------------------------------------------------------------- +// Internal constants +// --------------------------------------------------------------------------- + +const ESC: u8 = 0x1B; +const BEL: u8 = 0x07; +const BACKSLASH: u8 = b'\\'; +const RIGHT_BRACKET: u8 = b']'; + +/// Maximum bytes we'll buffer for the OSC parameter string. 32 bytes is far +/// more than any valid OSC 133 payload needs (e.g. `133;D;127` is 9 bytes). +/// Longer (non-133) OSC sequences simply stop accumulating once the buffer is +/// full — the dispatch logic will harmlessly ignore them. +const PARAM_BUF_CAP: usize = 32; + +// --------------------------------------------------------------------------- +// State machine +// --------------------------------------------------------------------------- + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum State { + /// Normal pass-through. + Ground, + /// Saw ESC (0x1B). + Esc, + /// Inside an OSC sequence (`ESC ]`), accumulating parameter bytes. + OscParam, + /// Inside an OSC sequence, saw ESC — next byte decides if this is `ESC \` + /// (string terminator) or something else. + OscEsc, +} + +/// A streaming, zero-allocation parser for OSC 133 escape sequences. +/// +/// Feed arbitrary byte slices into [`Parser::push`]. The parser detects +/// OSC 133 markers and reports [`Event`]s through a caller-supplied callback +/// without modifying the data. It can sit transparently between a PTY reader +/// and stdout. +pub struct Parser { + state: State, + zone: Zone, + param_buf: [u8; PARAM_BUF_CAP], + param_len: usize, +} + +impl Default for Parser { + fn default() -> Self { + Self::new() + } +} + +impl Parser { + /// Create a new parser in the initial (ground / unknown-zone) state. + #[inline] + pub fn new() -> Self { + Self { + state: State::Ground, + zone: Zone::Unknown, + param_buf: [0u8; PARAM_BUF_CAP], + param_len: 0, + } + } + + /// The current semantic zone based on markers seen so far. + #[inline] + #[allow(dead_code)] + pub fn zone(&self) -> Zone { + self.zone + } + + /// Process a chunk of bytes, calling `on_event` for every OSC 133 marker + /// found. + /// + /// All bytes in `data` should still be forwarded to the terminal by the + /// caller — this method only *observes* the stream. + #[inline] + pub fn push(&mut self, data: &[u8], mut on_event: impl FnMut(Event)) { + for &byte in data { + match self.state { + State::Ground => { + if byte == ESC { + self.state = State::Esc; + } + } + State::Esc => { + if byte == RIGHT_BRACKET { + self.state = State::OscParam; + self.param_len = 0; + } else { + self.state = State::Ground; + } + } + State::OscParam => { + if byte == BEL { + self.dispatch(&mut on_event); + self.state = State::Ground; + } else if byte == ESC { + self.state = State::OscEsc; + } else if self.param_len < PARAM_BUF_CAP { + self.param_buf[self.param_len] = byte; + self.param_len += 1; + } + // If param_len == PARAM_BUF_CAP we silently stop + // accumulating — dispatch will ignore non-133 sequences. + } + State::OscEsc => { + if byte == BACKSLASH { + self.dispatch(&mut on_event); + } + // Whether we got a valid ST or not, return to ground. + // (A new ESC ] would restart accumulation via the Ground + // -> Esc -> OscParam path on the *next* byte.) + self.state = State::Ground; + } + } + } + } + + /// Inspect the accumulated parameter buffer. If it holds an OSC 133 + /// payload, emit the corresponding [`Event`] and update the zone. + #[inline] + fn dispatch(&mut self, on_event: &mut impl FnMut(Event)) { + let params = &self.param_buf[..self.param_len]; + + // Must start with "133;" + if params.len() < 5 || ¶ms[..4] != b"133;" { + return; + } + + let cmd = params[4]; + let event = match cmd { + b'A' => { + self.zone = Zone::Prompt; + Event::PromptStart + } + b'B' => { + self.zone = Zone::Input; + Event::CommandStart + } + b'C' => { + self.zone = Zone::Output; + Event::CommandExecuted + } + b'D' => { + let exit_code = if params.len() > 6 && params[5] == b';' { + std::str::from_utf8(¶ms[6..]) + .ok() + .and_then(|s| s.parse::<i32>().ok()) + } else { + None + }; + self.zone = Zone::Unknown; + Event::CommandFinished { exit_code } + } + _ => return, + }; + + on_event(event); + } +} + +// --------------------------------------------------------------------------- +// Tests +// --------------------------------------------------------------------------- + +#[cfg(test)] +mod tests { + use super::*; + + /// Collect all events from a single `push` call. + fn parse_events(data: &[u8]) -> Vec<Event> { + let mut parser = Parser::new(); + let mut events = Vec::new(); + parser.push(data, |e| events.push(e)); + events + } + + // -- Basic event detection ------------------------------------------------ + + #[test] + fn detect_prompt_start_bel() { + let data = b"\x1b]133;A\x07"; + assert_eq!(parse_events(data), vec![Event::PromptStart]); + } + + #[test] + fn detect_prompt_start_st() { + let data = b"\x1b]133;A\x1b\\"; + assert_eq!(parse_events(data), vec![Event::PromptStart]); + } + + #[test] + fn detect_command_start_bel() { + let data = b"\x1b]133;B\x07"; + assert_eq!(parse_events(data), vec![Event::CommandStart]); + } + + #[test] + fn detect_command_start_st() { + let data = b"\x1b]133;B\x1b\\"; + assert_eq!(parse_events(data), vec![Event::CommandStart]); + } + + #[test] + fn detect_command_executed_bel() { + let data = b"\x1b]133;C\x07"; + assert_eq!(parse_events(data), vec![Event::CommandExecuted]); + } + + #[test] + fn detect_command_executed_st() { + let data = b"\x1b]133;C\x1b\\"; + assert_eq!(parse_events(data), vec![Event::CommandExecuted]); + } + + #[test] + fn detect_command_finished_no_exit_code() { + let data = b"\x1b]133;D\x07"; + assert_eq!( + parse_events(data), + vec![Event::CommandFinished { exit_code: None }] + ); + } + + #[test] + fn detect_command_finished_exit_zero() { + let data = b"\x1b]133;D;0\x07"; + assert_eq!( + parse_events(data), + vec![Event::CommandFinished { exit_code: Some(0) }] + ); + } + + #[test] + fn detect_command_finished_exit_nonzero() { + let data = b"\x1b]133;D;127\x07"; + assert_eq!( + parse_events(data), + vec![Event::CommandFinished { + exit_code: Some(127) + }] + ); + } + + #[test] + fn detect_command_finished_negative_exit_code() { + let data = b"\x1b]133;D;-1\x07"; + assert_eq!( + parse_events(data), + vec![Event::CommandFinished { + exit_code: Some(-1) + }] + ); + } + + #[test] + fn detect_command_finished_exit_code_st() { + let data = b"\x1b]133;D;42\x1b\\"; + assert_eq!( + parse_events(data), + vec![Event::CommandFinished { + exit_code: Some(42) + }] + ); + } + + #[test] + fn invalid_exit_code_yields_none() { + let data = b"\x1b]133;D;abc\x07"; + assert_eq!( + parse_events(data), + vec![Event::CommandFinished { exit_code: None }] + ); + } + + // -- Zone tracking -------------------------------------------------------- + + #[test] + fn zone_starts_unknown() { + let parser = Parser::new(); + assert_eq!(parser.zone(), Zone::Unknown); + } + + #[test] + fn full_zone_cycle() { + let mut parser = Parser::new(); + let mut events = Vec::new(); + + parser.push(b"\x1b]133;A\x07", |e| events.push(e)); + assert_eq!(parser.zone(), Zone::Prompt); + + parser.push(b"\x1b]133;B\x07", |e| events.push(e)); + assert_eq!(parser.zone(), Zone::Input); + + parser.push(b"\x1b]133;C\x07", |e| events.push(e)); + assert_eq!(parser.zone(), Zone::Output); + + parser.push(b"\x1b]133;D;0\x07", |e| events.push(e)); + assert_eq!(parser.zone(), Zone::Unknown); + + assert_eq!( + events, + vec![ + Event::PromptStart, + Event::CommandStart, + Event::CommandExecuted, + Event::CommandFinished { exit_code: Some(0) }, + ] + ); + } + + // -- Multiple events in one push ------------------------------------------ + + #[test] + fn multiple_events_single_push() { + let data = b"\x1b]133;A\x07$ \x1b]133;B\x07ls\n\x1b]133;C\x07file.txt\n\x1b]133;D;0\x07"; + let events = parse_events(data); + assert_eq!( + events, + vec![ + Event::PromptStart, + Event::CommandStart, + Event::CommandExecuted, + Event::CommandFinished { exit_code: Some(0) }, + ] + ); + } + + // -- Split across push boundaries ----------------------------------------- + + #[test] + fn split_esc_and_bracket() { + let mut parser = Parser::new(); + let mut events = Vec::new(); + + parser.push(b"\x1b", |e| events.push(e)); + assert!(events.is_empty()); + + parser.push(b"]133;A\x07", |e| events.push(e)); + assert_eq!(events, vec![Event::PromptStart]); + } + + #[test] + fn split_mid_param() { + let mut parser = Parser::new(); + let mut events = Vec::new(); + + parser.push(b"\x1b]13", |e| events.push(e)); + assert!(events.is_empty()); + + parser.push(b"3;D;42\x07", |e| events.push(e)); + assert_eq!( + events, + vec![Event::CommandFinished { + exit_code: Some(42) + }] + ); + } + + #[test] + fn split_before_terminator() { + let mut parser = Parser::new(); + let mut events = Vec::new(); + + parser.push(b"\x1b]133;B", |e| events.push(e)); + assert!(events.is_empty()); + + parser.push(b"\x07", |e| events.push(e)); + assert_eq!(events, vec![Event::CommandStart]); + } + + #[test] + fn split_esc_backslash_terminator() { + let mut parser = Parser::new(); + let mut events = Vec::new(); + + parser.push(b"\x1b]133;C\x1b", |e| events.push(e)); + assert!(events.is_empty()); + + parser.push(b"\\", |e| events.push(e)); + assert_eq!(events, vec![Event::CommandExecuted]); + } + + // -- Interleaved normal text ---------------------------------------------- + + #[test] + fn normal_text_before_and_after() { + let data = b"hello world\x1b]133;A\x07prompt text\x1b]133;B\x07command"; + let events = parse_events(data); + assert_eq!(events, vec![Event::PromptStart, Event::CommandStart]); + } + + // -- Non-133 OSC sequences (should be ignored) ---------------------------- + + #[test] + fn non_133_osc_ignored() { + let data = b"\x1b]0;window title\x07\x1b]133;A\x07"; + let events = parse_events(data); + assert_eq!(events, vec![Event::PromptStart]); + } + + #[test] + fn osc_7_ignored() { + let data = b"\x1b]7;file:///home/user\x07"; + assert!(parse_events(data).is_empty()); + } + + // -- Unknown command letter ----------------------------------------------- + + #[test] + fn unknown_command_ignored() { + let data = b"\x1b]133;Z\x07"; + assert!(parse_events(data).is_empty()); + } + + // -- Malformed sequences -------------------------------------------------- + + #[test] + fn esc_followed_by_non_bracket() { + let data = b"\x1b[31m\x1b]133;A\x07"; + let events = parse_events(data); + assert_eq!(events, vec![Event::PromptStart]); + } + + #[test] + fn lone_esc_at_end_of_chunk() { + let mut parser = Parser::new(); + let mut events = Vec::new(); + + parser.push(b"\x1b", |e| events.push(e)); + assert!(events.is_empty()); + + // Feed non-bracket to abort the escape, then a real sequence. + parser.push(b"x\x1b]133;A\x07", |e| events.push(e)); + assert_eq!(events, vec![Event::PromptStart]); + } + + #[test] + fn truncated_133_prefix() { + // "13" followed by terminator — not "133;" so no event. + let data = b"\x1b]13\x07"; + assert!(parse_events(data).is_empty()); + } + + #[test] + fn empty_osc() { + let data = b"\x1b]\x07"; + assert!(parse_events(data).is_empty()); + } + + // -- Buffer overflow (very long non-133 OSC) ------------------------------ + + #[test] + fn very_long_osc_does_not_panic() { + let mut data = Vec::new(); + data.extend_from_slice(b"\x1b]"); + data.extend(std::iter::repeat(b'x').take(1000)); + data.push(BEL); + // Should not panic and should produce no event. + assert!(parse_events(&data).is_empty()); + } + + // -- Empty input ---------------------------------------------------------- + + #[test] + fn empty_input() { + assert!(parse_events(b"").is_empty()); + } + + #[test] + fn only_normal_text() { + let data = b"just some regular terminal output\r\n"; + assert!(parse_events(data).is_empty()); + } + + // -- Repeated prompts (empty command) ------------------------------------ + + #[test] + fn repeated_prompt_cycle() { + let mut parser = Parser::new(); + let mut events = Vec::new(); + + // User hits enter on an empty prompt twice. + let data = b"\x1b]133;A\x07$ \x1b]133;B\x07\x1b]133;D\x07\x1b]133;A\x07$ \x1b]133;B\x07"; + parser.push(data, |e| events.push(e)); + + assert_eq!( + events, + vec![ + Event::PromptStart, + Event::CommandStart, + Event::CommandFinished { exit_code: None }, + Event::PromptStart, + Event::CommandStart, + ] + ); + assert_eq!(parser.zone(), Zone::Input); + } + + // -- Byte-at-a-time feeding ----------------------------------------------- + + #[test] + fn byte_at_a_time() { + let data = b"\x1b]133;D;99\x07"; + let mut parser = Parser::new(); + let mut events = Vec::new(); + + for &byte in data { + parser.push(&[byte], |e| events.push(e)); + } + + assert_eq!( + events, + vec![Event::CommandFinished { + exit_code: Some(99) + }] + ); + } + + // -- Mixed terminators ---------------------------------------------------- + + #[test] + fn mixed_bel_and_st_terminators() { + let data = b"\x1b]133;A\x07\x1b]133;B\x1b\\\x1b]133;C\x07\x1b]133;D;1\x1b\\"; + let events = parse_events(data); + assert_eq!( + events, + vec![ + Event::PromptStart, + Event::CommandStart, + Event::CommandExecuted, + Event::CommandFinished { exit_code: Some(1) }, + ] + ); + } + + // -- Default trait -------------------------------------------------------- + + #[test] + fn parser_default() { + let parser = Parser::default(); + assert_eq!(parser.zone(), Zone::Unknown); + } + + #[test] + fn zone_default() { + assert_eq!(Zone::default(), Zone::Unknown); + } + + // -- D with empty exit code field ----------------------------------------- + + #[test] + fn d_with_semicolon_but_empty_code() { + // "133;D;" — semicolon present but no digits. + let data = b"\x1b]133;D;\x07"; + assert_eq!( + parse_events(data), + vec![Event::CommandFinished { exit_code: None }] + ); + } + + // -- Consecutive OSC sequences without gap -------------------------------- + + #[test] + fn back_to_back_osc_no_gap() { + let data = b"\x1b]133;A\x07\x1b]133;B\x07"; + let events = parse_events(data); + assert_eq!(events, vec![Event::PromptStart, Event::CommandStart]); + } + + // -- CSI sequences interleaved (should not confuse parser) ---------------- + + #[test] + fn csi_sequences_ignored() { + // CSI (ESC [) color codes mixed with OSC 133. + let data = b"\x1b[32m\x1b]133;A\x07\x1b[0m$ \x1b]133;B\x07"; + let events = parse_events(data); + assert_eq!(events, vec![Event::PromptStart, Event::CommandStart]); + } + + // -- Large exit codes ----------------------------------------------------- + + #[test] + fn large_exit_code() { + let data = b"\x1b]133;D;2147483647\x07"; + assert_eq!( + parse_events(data), + vec![Event::CommandFinished { + exit_code: Some(i32::MAX) + }] + ); + } + + #[test] + fn overflow_exit_code_yields_none() { + let data = b"\x1b]133;D;9999999999999\x07"; + assert_eq!( + parse_events(data), + vec![Event::CommandFinished { exit_code: None }] + ); + } +} diff --git a/crates/atuin-history/Cargo.toml b/crates/atuin-history/Cargo.toml new file mode 100644 index 00000000..3bd3a100 --- /dev/null +++ b/crates/atuin-history/Cargo.toml @@ -0,0 +1,30 @@ +[package] +name = "atuin-history" +description = "The history crate for Atuin" +edition = "2024" +version = { workspace = true } + +authors.workspace = true +rust-version.workspace = true +license.workspace = true +homepage.workspace = true +repository.workspace = true +readme.workspace = true + +# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html + +[dependencies] +atuin-client = { path = "../atuin-client", version = "18.13.2" } + +time = { workspace = true } +serde = { workspace = true } +crossterm = { workspace = true, features = ["use-dev-tty"] } +unicode-segmentation = "1.11.0" + +[dev-dependencies] +divan = "0.1.14" +rand = { workspace = true } + +[[bench]] +name = "smart_sort" +harness = false diff --git a/crates/atuin-history/benches/smart_sort.rs b/crates/atuin-history/benches/smart_sort.rs new file mode 100644 index 00000000..a78064de --- /dev/null +++ b/crates/atuin-history/benches/smart_sort.rs @@ -0,0 +1,35 @@ +use atuin_client::history::History; +use atuin_history::sort::sort; + +use rand::Rng; + +fn main() { + // Run registered benchmarks. + divan::main(); +} + +// Smart sort usually runs on 200 entries, test on a few sizes +#[divan::bench(args=[100, 200, 400, 800, 1600, 10000])] +fn smart_sort(lines: usize) { + // benchmark a few different sizes of "history" + // first we need to generate some history. This will use a whole bunch of memory, sorry + let mut rng = rand::thread_rng(); + let now = time::OffsetDateTime::now_utc().unix_timestamp(); + + let possible_commands = ["echo", "ls", "cd", "grep", "atuin", "curl"]; + let mut commands = Vec::<History>::with_capacity(lines); + + for _ in 0..lines { + let command = possible_commands[rng.gen_range(0..possible_commands.len())]; + + let command = History::import() + .command(command) + .timestamp(time::OffsetDateTime::from_unix_timestamp(rng.gen_range(0..now)).unwrap()) + .build() + .into(); + + commands.push(command); + } + + let _ = sort("curl", commands); +} diff --git a/crates/atuin-history/src/lib.rs b/crates/atuin-history/src/lib.rs new file mode 100644 index 00000000..e7b33916 --- /dev/null +++ b/crates/atuin-history/src/lib.rs @@ -0,0 +1,2 @@ +pub mod sort; +pub mod stats; diff --git a/crates/atuin-history/src/sort.rs b/crates/atuin-history/src/sort.rs new file mode 100644 index 00000000..4465a142 --- /dev/null +++ b/crates/atuin-history/src/sort.rs @@ -0,0 +1,46 @@ +use atuin_client::history::History; + +type ScoredHistory = (f64, History); + +// Fuzzy search already comes sorted by minspan +// This sorting should be applicable to all search modes, and solve the more "obvious" issues +// first. +// Later on, we can pass in context and do some boosts there too. +pub fn sort(query: &str, input: Vec<History>) -> Vec<History> { + // This can totally be extended. We need to be _careful_ that it's not slow. + // We also need to balance sorting db-side with sorting here. SQLite can do a lot, + // but some things are just much easier/more doable in Rust. + + let mut scored = input + .into_iter() + .map(|h| { + // If history is _prefixed_ with the query, score it more highly + let score = if h.command.starts_with(query) { + 2.0 + } else if h.command.contains(query) { + 1.75 + } else { + 1.0 + }; + + // calculate how long ago the history was, in seconds + let now = time::OffsetDateTime::now_utc().unix_timestamp(); + let time = h.timestamp.unix_timestamp(); + let diff = std::cmp::max(1, now - time); // no /0 please + + // prefer newer history, but not hugely so as to offset the other scoring + // the numbers will get super small over time, but I don't want time to overpower other + // scoring + #[allow(clippy::cast_precision_loss)] + let time_score = 1.0 + (1.0 / diff as f64); + let score = score * time_score; + + (score, h) + }) + .collect::<Vec<ScoredHistory>>(); + + scored.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap().reverse()); + + // Remove the scores and return the history + scored.into_iter().map(|(_, h)| h).collect::<Vec<History>>() +} diff --git a/crates/atuin-history/src/stats.rs b/crates/atuin-history/src/stats.rs new file mode 100644 index 00000000..8bf03e42 --- /dev/null +++ b/crates/atuin-history/src/stats.rs @@ -0,0 +1,554 @@ +use std::collections::{HashMap, HashSet}; + +use crossterm::style::{Color, ResetColor, SetAttribute, SetForegroundColor}; +use serde::{Deserialize, Serialize}; +use unicode_segmentation::UnicodeSegmentation; + +use atuin_client::{history::History, settings::Settings, theme::Meaning, theme::Theme}; + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct Stats { + pub total_commands: usize, + pub unique_commands: usize, + pub top: Vec<(Vec<String>, usize)>, +} + +fn first_non_whitespace(s: &str) -> Option<usize> { + s.char_indices() + // find the first non whitespace char + .find(|(_, c)| !c.is_ascii_whitespace()) + // return the index of that char + .map(|(i, _)| i) +} + +fn first_whitespace(s: &str) -> usize { + s.char_indices() + // find the first whitespace char + .find(|(_, c)| c.is_ascii_whitespace()) + // return the index of that char, (or the max length of the string) + .map_or(s.len(), |(i, _)| i) +} + +fn interesting_command<'a>(settings: &Settings, mut command: &'a str) -> &'a str { + // Sort by length so that we match the longest prefix first + let mut common_prefix = settings.stats.common_prefix.clone(); + common_prefix.sort_by_key(|b| std::cmp::Reverse(b.len())); + + // Trim off the common prefix, if it exists + for p in &common_prefix { + if command.starts_with(p) { + let i = p.len(); + let prefix = &command[..i]; + command = command[i..].trim_start(); + if command.is_empty() { + // no commands following, just use the prefix + return prefix; + } + break; + } + } + + // Sort the common_subcommands by length so that we match the longest subcommand first + let mut common_subcommands = settings.stats.common_subcommands.clone(); + common_subcommands.sort_by_key(|b| std::cmp::Reverse(b.len())); + + // Check for a common subcommand + for p in &common_subcommands { + if command.starts_with(p) { + // if the subcommand is the same length as the command, then we just use the subcommand + if p.len() == command.len() { + return command; + } + // otherwise we need to use the subcommand + the next word + let non_whitespace = first_non_whitespace(&command[p.len()..]).unwrap_or(0); + let j = + p.len() + non_whitespace + first_whitespace(&command[p.len() + non_whitespace..]); + return &command[..j]; + } + } + // Return the first word if there is no subcommand + &command[..first_whitespace(command)] +} + +fn split_at_pipe(command: &str) -> Vec<&str> { + let mut result = vec![]; + let mut quoted = false; + let mut start = 0; + let mut graphemes = UnicodeSegmentation::grapheme_indices(command, true); + + while let Some((i, c)) = graphemes.next() { + let current = i; + match c { + "\"" => { + if command[start..current] != *"\"" { + quoted = !quoted; + } + } + "'" => { + if command[start..current] != *"'" { + quoted = !quoted; + } + } + "\\" => if graphemes.next().is_some() {}, + "|" => { + if !quoted { + if current > start && command[start..].starts_with('|') { + start += 1; + } + result.push(&command[start..current]); + start = current; + } + } + _ => {} + } + } + if command[start..].starts_with('|') { + start += 1; + } + result.push(&command[start..]); + result +} + +fn strip_leading_env_vars(command: &str) -> &str { + // fast path: no equals sign, no environment variable + if !command.contains('=') { + return command; + } + + let mut in_token = false; + let mut token_start_pos = 0; + let mut in_single_quotes = false; + let mut in_double_quotes = false; + let mut escape_next = false; + let mut has_equals_outside_quotes = false; + + for (i, g) in UnicodeSegmentation::grapheme_indices(command, true) { + if escape_next { + escape_next = false; + continue; + } + + if !in_token { + token_start_pos = i; + } + + match g { + "\\" => { + escape_next = true; + in_token = true; + } + "'" if !in_double_quotes => { + in_single_quotes = !in_single_quotes; + in_token = true; + } + "\"" if !in_single_quotes => { + in_double_quotes = !in_double_quotes; + in_token = true; + } + "=" if !in_single_quotes && !in_double_quotes => { + has_equals_outside_quotes = true; + in_token = true; + } + " " | "\t" if !in_single_quotes && !in_double_quotes => { + if in_token { + if !has_equals_outside_quotes { + // if we're not in an env var, we can break early + break; + } + in_token = false; + has_equals_outside_quotes = false; + } + } + _ => { + in_token = true; + } + } + } + + command[token_start_pos..].trim() +} + +pub fn pretty_print(stats: Stats, ngram_size: usize, theme: &Theme) { + let max = stats.top.iter().map(|x| x.1).max().unwrap(); + let num_pad = max.ilog10() as usize + 1; + + // Find the length of the longest command name for each column + let column_widths = stats + .top + .iter() + .map(|(commands, _)| commands.iter().map(|c| c.len()).collect::<Vec<usize>>()) + .fold(vec![0; ngram_size], |acc, item| { + acc.iter() + .zip(item.iter()) + .map(|(a, i)| *std::cmp::max(a, i)) + .collect() + }); + + for (command, count) in stats.top { + let gray = SetForegroundColor(match theme.as_style(Meaning::Muted).foreground_color { + Some(color) => color, + None => Color::Grey, + }); + let bold = SetAttribute(crossterm::style::Attribute::Bold); + + let in_ten = 10 * count / max; + + print!("["); + print!( + "{}", + SetForegroundColor(match theme.get_error().foreground_color { + Some(color) => color, + None => Color::Red, + }) + ); + + for i in 0..in_ten { + if i == 2 { + print!( + "{}", + SetForegroundColor(match theme.get_warning().foreground_color { + Some(color) => color, + None => Color::Yellow, + }) + ); + } + + if i == 5 { + print!( + "{}", + SetForegroundColor(match theme.get_info().foreground_color { + Some(color) => color, + None => Color::Green, + }) + ); + } + + print!("▮"); + } + + for _ in in_ten..10 { + print!(" "); + } + + let formatted_command = command + .iter() + .zip(column_widths.iter()) + .map(|(cmd, width)| format!("{cmd:width$}")) + .collect::<Vec<_>>() + .join(" | "); + + println!( + "{ResetColor}] {gray}{count:num_pad$}{ResetColor} {bold}{formatted_command}{ResetColor}" + ); + } + println!("Total commands: {}", stats.total_commands); + println!("Unique commands: {}", stats.unique_commands); +} + +pub fn compute( + settings: &Settings, + history: &[History], + count: usize, + ngram_size: usize, +) -> Option<Stats> { + let mut commands = HashSet::<&str>::with_capacity(history.len()); + let mut total_unignored = 0; + let mut prefixes = HashMap::<Vec<&str>, usize>::with_capacity(history.len()); + + for i in history { + // just in case it somehow has a leading tab or space or something (legacy atuin didn't ignore space prefixes) + let command = strip_leading_env_vars(i.command.trim()); + let prefix = interesting_command(settings, command); + + if settings.stats.ignored_commands.iter().any(|c| c == prefix) { + continue; + } + + total_unignored += 1; + commands.insert(command); + + split_at_pipe(command) + .iter() + .map(|l| { + let command = l.trim(); + commands.insert(command); + command + }) + .collect::<Vec<_>>() + .windows(ngram_size) + .for_each(|w| { + *prefixes + .entry(w.iter().map(|c| interesting_command(settings, c)).collect()) + .or_default() += 1; + }); + } + + let unique = commands.len(); + let mut top = prefixes.into_iter().collect::<Vec<_>>(); + + top.sort_unstable_by_key(|x| std::cmp::Reverse(x.1)); + top.truncate(count); + + if top.is_empty() { + return None; + } + + Some(Stats { + unique_commands: unique, + total_commands: total_unignored, + top: top + .into_iter() + .map(|t| (t.0.into_iter().map(|s| s.to_string()).collect(), t.1)) + .collect(), + }) +} + +#[cfg(test)] +mod tests { + use atuin_client::history::History; + use atuin_client::settings::Settings; + use time::OffsetDateTime; + + use super::compute; + use super::{interesting_command, split_at_pipe, strip_leading_env_vars}; + + #[test] + fn ignored_env_vars() { + let settings = Settings::utc(); + + let history: History = History::capture() + .timestamp(time::OffsetDateTime::now_utc()) + .command("FOO='BAR=🚀' echo foo") + .cwd("/") + .build() + .into(); + + let stats = compute(&settings, &[history], 10, 1).expect("failed to compute stats"); + assert_eq!(stats.top.first().unwrap().0, vec!["echo"]); + } + + #[test] + fn ignored_commands() { + let mut settings = Settings::utc(); + settings.stats.ignored_commands.push("cd".to_string()); + + let history = [ + History::import() + .timestamp(OffsetDateTime::now_utc()) + .command("cd foo") + .build() + .into(), + History::import() + .timestamp(OffsetDateTime::now_utc()) + .command("cargo build stuff") + .build() + .into(), + ]; + + let stats = compute(&settings, &history, 10, 1).expect("failed to compute stats"); + assert_eq!(stats.total_commands, 1); + assert_eq!(stats.unique_commands, 1); + } + + #[test] + fn interesting_commands() { + let settings = Settings::utc(); + + assert_eq!(interesting_command(&settings, "cargo"), "cargo"); + assert_eq!( + interesting_command(&settings, "cargo build foo bar"), + "cargo build" + ); + assert_eq!( + interesting_command(&settings, "sudo cargo build foo bar"), + "cargo build" + ); + assert_eq!(interesting_command(&settings, "sudo"), "sudo"); + } + + // Test with spaces in the common_prefix + #[test] + fn interesting_commands_spaces() { + let mut settings = Settings::utc(); + settings.stats.common_prefix.push("sudo test".to_string()); + + assert_eq!(interesting_command(&settings, "sudo test"), "sudo test"); + assert_eq!(interesting_command(&settings, "sudo test "), "sudo test"); + assert_eq!(interesting_command(&settings, "sudo test foo bar"), "foo"); + assert_eq!( + interesting_command(&settings, "sudo test foo bar"), + "foo" + ); + + // Works with a common_subcommand as well + assert_eq!( + interesting_command(&settings, "sudo test cargo build foo bar"), + "cargo build" + ); + + // We still match on just the sudo prefix + assert_eq!(interesting_command(&settings, "sudo"), "sudo"); + assert_eq!(interesting_command(&settings, "sudo foo"), "foo"); + } + + // Test with spaces in the common_subcommand + #[test] + fn interesting_commands_spaces_subcommand() { + let mut settings = Settings::utc(); + settings + .stats + .common_subcommands + .push("cargo build".to_string()); + + assert_eq!(interesting_command(&settings, "cargo build"), "cargo build"); + assert_eq!( + interesting_command(&settings, "cargo build "), + "cargo build" + ); + assert_eq!( + interesting_command(&settings, "cargo build foo bar"), + "cargo build foo" + ); + + // Works with a common_prefix as well + assert_eq!( + interesting_command(&settings, "sudo cargo build foo bar"), + "cargo build foo" + ); + + // We still match on just cargo as a subcommand + assert_eq!(interesting_command(&settings, "cargo"), "cargo"); + assert_eq!(interesting_command(&settings, "cargo foo"), "cargo foo"); + } + + // Test with spaces in the common_prefix and common_subcommand + #[test] + fn interesting_commands_spaces_both() { + let mut settings = Settings::utc(); + settings.stats.common_prefix.push("sudo test".to_string()); + settings + .stats + .common_subcommands + .push("cargo build".to_string()); + + assert_eq!( + interesting_command(&settings, "sudo test cargo build"), + "cargo build" + ); + assert_eq!( + interesting_command(&settings, "sudo test cargo build"), + "cargo build" + ); + assert_eq!( + interesting_command(&settings, "sudo test cargo build "), + "cargo build" + ); + assert_eq!( + interesting_command(&settings, "sudo test cargo build foo bar"), + "cargo build foo" + ); + } + + #[test] + fn split_simple() { + assert_eq!(split_at_pipe("fd | rg"), ["fd ", " rg"]); + } + + #[test] + fn split_multi() { + assert_eq!( + split_at_pipe("kubectl | jq | rg"), + ["kubectl ", " jq ", " rg"] + ); + } + + #[test] + fn split_simple_quoted() { + assert_eq!( + split_at_pipe("foo | bar 'baz {} | quux' | xyzzy"), + ["foo ", " bar 'baz {} | quux' ", " xyzzy"] + ); + } + + #[test] + fn split_multi_quoted() { + assert_eq!( + split_at_pipe("foo | bar 'baz \"{}\" | quux' | xyzzy"), + ["foo ", " bar 'baz \"{}\" | quux' ", " xyzzy"] + ); + } + + #[test] + fn escaped_pipes() { + assert_eq!( + split_at_pipe("foo | bar baz \\| quux"), + ["foo ", " bar baz \\| quux"] + ); + } + + #[test] + fn emoji() { + assert_eq!( + split_at_pipe("git commit -m \"🚀\""), + ["git commit -m \"🚀\""] + ); + } + + #[test] + fn starts_with_pipe() { + assert_eq!( + split_at_pipe("| sed 's/[0-9a-f]//g'"), + ["", " sed 's/[0-9a-f]//g'"] + ); + } + + #[test] + fn starts_with_spaces_and_pipe() { + assert_eq!( + split_at_pipe(" | sed 's/[0-9a-f]//g'"), + [" ", " sed 's/[0-9a-f]//g'"] + ); + } + + #[test] + fn strip_leading_env_vars_simple() { + assert_eq!( + strip_leading_env_vars("FOO=bar BAZ=quux echo foo"), + "echo foo" + ); + } + + #[test] + fn strip_leading_env_vars_quoted_single() { + assert_eq!(strip_leading_env_vars("FOO='BAR=baz' echo foo"), "echo foo"); + } + + #[test] + fn strip_leading_env_vars_quoted_double() { + assert_eq!( + strip_leading_env_vars("FOO=\"BAR=baz\" echo foo"), + "echo foo" + ); + } + + #[test] + fn strip_leading_env_vars_quoted_single_and_double() { + assert_eq!( + strip_leading_env_vars("FOO='BAR=\"baz\"' echo foo \"BAR=quux\""), + "echo foo \"BAR=quux\"" + ); + } + + #[test] + fn strip_leading_env_vars_emojis() { + assert_eq!( + strip_leading_env_vars("FOO='BAR=🚀' echo foo \"BAR=quux\" foo"), + "echo foo \"BAR=quux\" foo" + ); + } + + #[test] + fn strip_leading_env_vars_name_same_as_command() { + assert_eq!(strip_leading_env_vars("FOO='bar' bar baz"), "bar baz"); + } +} diff --git a/crates/atuin-kv/Cargo.toml b/crates/atuin-kv/Cargo.toml new file mode 100644 index 00000000..d72a3343 --- /dev/null +++ b/crates/atuin-kv/Cargo.toml @@ -0,0 +1,27 @@ +[package] +name = "atuin-kv" +edition = "2024" +version = { workspace = true } +description = "The kv crate for Atuin" + +authors.workspace = true +rust-version.workspace = true +license.workspace = true +homepage.workspace = true +repository.workspace = true +readme.workspace = true + +# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html + +[dependencies] +atuin-client = { path = "../atuin-client", version = "18.13.2" } +atuin-common = { path = "../atuin-common", version = "18.13.2" } + +tracing = { workspace = true } +tracing-subscriber = { workspace = true } +rmp = { version = "0.8.14" } +eyre = { workspace = true } +tokio = { workspace = true } +typed-builder = { workspace = true } +pretty_assertions = { workspace = true } +sqlx = { workspace = true } diff --git a/crates/atuin-kv/migrations/20250501160746_create_kv_db.down.sql b/crates/atuin-kv/migrations/20250501160746_create_kv_db.down.sql new file mode 100644 index 00000000..bce8dfd3 --- /dev/null +++ b/crates/atuin-kv/migrations/20250501160746_create_kv_db.down.sql @@ -0,0 +1,2 @@ +-- Add down migration script here +DROP TABLE kv; diff --git a/crates/atuin-kv/migrations/20250501160746_create_kv_db.up.sql b/crates/atuin-kv/migrations/20250501160746_create_kv_db.up.sql new file mode 100644 index 00000000..77384044 --- /dev/null +++ b/crates/atuin-kv/migrations/20250501160746_create_kv_db.up.sql @@ -0,0 +1,12 @@ +-- Add up migration script here +CREATE TABLE + kv ( + namespace TEXT NOT NULL, + key TEXT NOT NULL, + value TEXT NOT NULL, + inserted_at INTEGER NOT NULL DEFAULT (strftime ('%s', 'now')) + ); + +CREATE INDEX idx_kv_namespace ON kv (namespace); + +CREATE UNIQUE INDEX idx_kv ON kv (namespace, key); diff --git a/crates/atuin-kv/src/database.rs b/crates/atuin-kv/src/database.rs new file mode 100644 index 00000000..446aa7af --- /dev/null +++ b/crates/atuin-kv/src/database.rs @@ -0,0 +1,229 @@ +use std::{path::Path, str::FromStr, time::Duration}; + +use atuin_common::utils; +use sqlx::{ + Result, Row, + sqlite::{ + SqliteConnectOptions, SqliteJournalMode, SqlitePool, SqlitePoolOptions, SqliteRow, + SqliteSynchronous, + }, +}; +use tokio::fs; +use tracing::debug; + +use crate::store::entry::KvEntry; + +#[derive(Debug, Clone)] +pub struct Database { + pub pool: SqlitePool, +} + +impl Database { + pub async fn new(path: impl AsRef<Path>, timeout: f64) -> Result<Self> { + let path = path.as_ref(); + debug!("opening KV sqlite database at {:?}", path); + + if utils::broken_symlink(path) { + eprintln!( + "Atuin: KV sqlite db path ({path:?}) is a broken symlink. Unable to read or create replacement." + ); + std::process::exit(1); + } + + if !path.exists() + && let Some(dir) = path.parent() + { + fs::create_dir_all(dir).await?; + } + + let opts = SqliteConnectOptions::from_str(path.as_os_str().to_str().unwrap())? + .journal_mode(SqliteJournalMode::Wal) + .optimize_on_close(true, None) + .synchronous(SqliteSynchronous::Normal) + .with_regexp() + .foreign_keys(true) + .create_if_missing(true); + + let pool = SqlitePoolOptions::new() + .acquire_timeout(Duration::from_secs_f64(timeout)) + .connect_with(opts) + .await?; + + Self::setup_db(&pool).await?; + Ok(Self { pool }) + } + + pub async fn sqlite_version(&self) -> Result<String> { + sqlx::query_scalar("SELECT sqlite_version()") + .fetch_one(&self.pool) + .await + } + + async fn setup_db(pool: &SqlitePool) -> Result<()> { + debug!("running sqlite database setup"); + + sqlx::migrate!("./migrations").run(pool).await?; + + Ok(()) + } + + async fn save_raw(tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, e: &KvEntry) -> Result<()> { + sqlx::query( + "insert into kv(namespace, key, value) + values(?1, ?2, ?3) + on conflict(namespace, key) do update set + namespace = excluded.namespace, + key = excluded.key, + value = excluded.value", + ) + .bind(e.namespace.as_str()) + .bind(e.key.as_str()) + .bind(e.value.as_str()) + .execute(&mut **tx) + .await?; + + Ok(()) + } + + async fn delete_raw( + tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, + namespace: &str, + key: &str, + ) -> Result<()> { + sqlx::query("delete from kv where namespace = ?1 and key = ?2") + .bind(namespace) + .bind(key) + .execute(&mut **tx) + .await?; + Ok(()) + } + + pub async fn save(&self, e: &KvEntry) -> Result<()> { + debug!("saving kv entry to sqlite"); + let mut tx = self.pool.begin().await?; + Self::save_raw(&mut tx, e).await?; + tx.commit().await?; + + Ok(()) + } + + pub async fn delete(&self, namespace: &str, key: &str) -> Result<()> { + debug!("deleting kv entry {namespace}/{key}"); + + let mut tx = self.pool.begin().await?; + Self::delete_raw(&mut tx, namespace, key).await?; + tx.commit().await?; + + Ok(()) + } + + fn query_kv_entry(row: SqliteRow) -> KvEntry { + let namespace = row.get("namespace"); + let key = row.get("key"); + let value = row.get("value"); + + KvEntry::builder() + .namespace(namespace) + .key(key) + .value(value) + .build() + } + + pub async fn load(&self, namespace: &str, key: &str) -> Result<Option<KvEntry>> { + debug!("loading kv entry {namespace}.{key}"); + + let res = sqlx::query("select * from kv where namespace = ?1 and key = ?2") + .bind(namespace) + .bind(key) + .map(Self::query_kv_entry) + .fetch_optional(&self.pool) + .await?; + + Ok(res) + } + + pub async fn list(&self, namespace: Option<&str>) -> Result<Vec<KvEntry>> { + debug!("listing kv entries"); + + let res = if let Some(namespace) = namespace { + sqlx::query("select * from kv where namespace = ?1 order by key asc") + .bind(namespace) + .map(Self::query_kv_entry) + .fetch_all(&self.pool) + .await? + } else { + sqlx::query("select * from kv order by namespace, key asc") + .map(Self::query_kv_entry) + .fetch_all(&self.pool) + .await? + }; + + Ok(res) + } +} + +#[cfg(test)] +mod test { + use super::*; + + #[tokio::test] + async fn test_list() { + let db = Database::new("sqlite::memory:", 1.0).await.unwrap(); + let scripts = db.list(None).await.unwrap(); + assert_eq!(scripts.len(), 0); + + let entry = KvEntry::builder() + .namespace("test".to_string()) + .key("test".to_string()) + .value("test".to_string()) + .build(); + + db.save(&entry).await.unwrap(); + + let entries = db.list(None).await.unwrap(); + assert_eq!(entries.len(), 1); + assert_eq!(entries[0].namespace, "test"); + assert_eq!(entries[0].key, "test"); + assert_eq!(entries[0].value, "test"); + } + + #[tokio::test] + async fn test_save_load() { + let db = Database::new("sqlite::memory:", 1.0).await.unwrap(); + + let entry = KvEntry::builder() + .namespace("test".to_string()) + .key("test".to_string()) + .value("test".to_string()) + .build(); + + db.save(&entry).await.unwrap(); + + let loaded = db + .load(&entry.namespace, &entry.key) + .await + .unwrap() + .unwrap(); + + assert_eq!(loaded, entry); + } + + #[tokio::test] + async fn test_delete() { + let db = Database::new("sqlite::memory:", 1.0).await.unwrap(); + + let entry = KvEntry::builder() + .namespace("test".to_string()) + .key("test".to_string()) + .value("test".to_string()) + .build(); + + db.save(&entry).await.unwrap(); + + assert_eq!(db.list(None).await.unwrap().len(), 1); + db.delete(&entry.namespace, &entry.key).await.unwrap(); + + let loaded = db.list(None).await.unwrap(); + assert_eq!(loaded.len(), 0); + } +} diff --git a/crates/atuin-kv/src/lib.rs b/crates/atuin-kv/src/lib.rs new file mode 100644 index 00000000..ad57b6ac --- /dev/null +++ b/crates/atuin-kv/src/lib.rs @@ -0,0 +1,2 @@ +pub mod database; +pub mod store; diff --git a/crates/atuin-kv/src/store.rs b/crates/atuin-kv/src/store.rs new file mode 100644 index 00000000..23fd7934 --- /dev/null +++ b/crates/atuin-kv/src/store.rs @@ -0,0 +1,211 @@ +use std::collections::HashSet; + +use eyre::{Result, bail}; + +use atuin_client::record::sqlite_store::SqliteStore; +use atuin_client::record::{encryption::PASETO_V4, store::Store}; +use atuin_common::record::{Host, HostId, Record, RecordId, RecordIdx}; +use entry::KvEntry; +use record::{KV_TAG, KV_VERSION, KvRecord}; + +use crate::database::Database; + +pub mod entry; +pub mod record; + +#[derive(Debug, Clone)] +pub struct KvStore { + pub record_store: SqliteStore, + pub kv_db: Database, + pub host_id: HostId, + pub encryption_key: [u8; 32], +} + +impl KvStore { + pub fn new( + record_store: SqliteStore, + kv_db: Database, + host_id: HostId, + encryption_key: [u8; 32], + ) -> Self { + KvStore { + record_store, + kv_db, + host_id, + encryption_key, + } + } + + pub async fn set(&self, namespace: &str, key: &str, value: &str) -> Result<()> { + let kv_record = KvRecord::builder() + .namespace(namespace.to_string()) + .key(key.to_string()) + .value(Some(value.to_string())) + .build(); + + self.push_record(kv_record).await?; + + let kv = KvEntry::builder() + .namespace(namespace.to_string()) + .key(key.to_string()) + .value(value.to_string()) + .build(); + + self.kv_db.save(&kv).await?; + + Ok(()) + } + + pub async fn get(&self, namespace: &str, key: &str) -> Result<Option<String>> { + let kv = self.kv_db.load(namespace, key).await?; + Ok(kv.map(|kv| kv.value)) + } + + pub async fn delete(&self, namespace: &str, keys: &[String]) -> Result<()> { + for key in keys { + let record = KvRecord::builder() + .namespace(namespace.to_string()) + .key(key.to_string()) + .value(None) + .build(); + + self.push_record(record).await?; + self.kv_db.delete(namespace, key).await?; + } + + Ok(()) + } + + pub async fn list(&self, namespace: Option<&str>) -> Result<Vec<KvEntry>> { + let entries = self.kv_db.list(namespace).await?; + + Ok(entries) + } + + async fn push_record(&self, record: KvRecord) -> Result<(RecordId, RecordIdx)> { + let bytes = record.serialize()?; + let idx = self + .record_store + .last(self.host_id, KV_TAG) + .await? + .map_or(0, |p| p.idx + 1); + + let record = Record::builder() + .host(Host::new(self.host_id)) + .version(KV_VERSION.to_string()) + .tag(KV_TAG.to_string()) + .idx(idx) + .data(bytes) + .build(); + + let id = record.id; + + self.record_store + .push(&record.encrypt::<PASETO_V4>(&self.encryption_key)) + .await?; + + Ok((id, idx)) + } + + pub async fn build(&self) -> Result<()> { + let mut tagged = self.record_store.all_tagged(KV_TAG).await?; + tagged.reverse(); + + let cached = self.kv_db.list(None).await?; + + let mut visited = HashSet::new(); + + // Iterate through all KV records from newest to oldest; + // only visit each KV once, inserting or deleting based on the first time we see it + for record in tagged { + let decrypted = match record.version.as_str() { + "v0" | KV_VERSION => record.decrypt::<PASETO_V4>(&self.encryption_key)?, + version => bail!("unknown version {version:?}"), + }; + + let kv = KvRecord::deserialize(&decrypted.data, &decrypted.version)?; + let uniq_id = format!("{}.{}", kv.namespace, kv.key); + + if visited.insert(uniq_id) { + match kv.value { + Some(value) => { + self.kv_db + .save( + &KvEntry::builder() + .namespace(kv.namespace.clone()) + .key(kv.key.clone()) + .value(value) + .build(), + ) + .await?; + } + None => { + self.kv_db + .delete(kv.namespace.as_str(), kv.key.as_str()) + .await?; + } + } + } + } + + // Any KVs that were in the cache but not in the tagged list should be deleted; + // this should never happen in practice since the cache is always built from the tagged list, + // but just in case because ** S O F T W A R E ** + for kv in cached { + if !visited.contains(&format!("{}.{}", kv.namespace, kv.key)) { + self.kv_db + .delete(kv.namespace.as_str(), kv.key.as_str()) + .await?; + } + } + + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + async fn setup() -> Result<KvStore> { + let record_store = SqliteStore::new("sqlite::memory:", 1.0).await.unwrap(); + let kv_db = Database::new("sqlite::memory:", 1.0).await.unwrap(); + let host_id = atuin_common::record::HostId(atuin_common::utils::uuid_v7()); + let encryption_key = [0; 32]; + Ok(KvStore::new(record_store, kv_db, host_id, encryption_key)) + } + + #[tokio::test] + async fn test_kv_store() -> Result<()> { + let store = setup().await?; + + store.set("test", "key", "value").await.unwrap(); + let value = store.get("test", "key").await.unwrap(); + assert_eq!(value, Some("value".to_string())); + + let records = store.record_store.all_tagged(KV_TAG).await?; + assert_eq!(records.len(), 1); + + let list = store.list(Some("test")).await.unwrap(); + let expected = vec![ + KvEntry::builder() + .namespace("test".to_string()) + .key("key".to_string()) + .value("value".to_string()) + .build(), + ]; + assert_eq!(list, expected); + + let ns_list = store.list(None).await.unwrap(); + assert_eq!(ns_list, expected); + + store.delete("test", &["key".to_string()]).await.unwrap(); + let value = store.get("test", "key").await.unwrap(); + assert_eq!(value, None); + + let records = store.record_store.all_tagged(KV_TAG).await?; + assert_eq!(records.len(), 2); + + Ok(()) + } +} diff --git a/crates/atuin-kv/src/store/entry.rs b/crates/atuin-kv/src/store/entry.rs new file mode 100644 index 00000000..1d6a1ef8 --- /dev/null +++ b/crates/atuin-kv/src/store/entry.rs @@ -0,0 +1,8 @@ +use typed_builder::TypedBuilder; + +#[derive(Debug, Clone, PartialEq, Eq, TypedBuilder)] +pub struct KvEntry { + pub namespace: String, + pub key: String, + pub value: String, +} diff --git a/crates/atuin-kv/src/store/record.rs b/crates/atuin-kv/src/store/record.rs new file mode 100644 index 00000000..37254176 --- /dev/null +++ b/crates/atuin-kv/src/store/record.rs @@ -0,0 +1,159 @@ +use atuin_common::record::DecryptedData; +use eyre::{Result, bail, ensure, eyre}; +use typed_builder::TypedBuilder; + +pub const KV_VERSION: &str = "v1"; +pub const KV_TAG: &str = "kv"; +pub const KV_VAL_MAX_LEN: usize = 100 * 1024; + +#[derive(Debug, Clone, PartialEq, Eq, TypedBuilder)] +pub struct KvRecord { + pub namespace: String, + pub key: String, + pub value: Option<String>, +} + +impl KvRecord { + pub fn serialize(&self) -> Result<DecryptedData> { + use rmp::encode; + + let mut output = vec![]; + + // INFO: ensure this is updated when adding new fields + encode::write_array_len(&mut output, 4)?; + + encode::write_str(&mut output, &self.namespace)?; + encode::write_str(&mut output, &self.key)?; + encode::write_bool(&mut output, self.value.is_some())?; + + if let Some(value) = &self.value { + encode::write_str(&mut output, value)?; + } + + Ok(DecryptedData(output)) + } + + pub fn deserialize(data: &DecryptedData, version: &str) -> Result<Self> { + use rmp::decode; + + fn error_report<E: std::fmt::Debug>(err: E) -> eyre::Report { + eyre!("{err:?}") + } + + match version { + "v0" => { + let mut bytes = decode::Bytes::new(&data.0); + + let nfields = decode::read_array_len(&mut bytes).map_err(error_report)?; + ensure!(nfields == 3, "too many entries in v0 kv record"); + + let bytes = bytes.remaining_slice(); + + let (namespace, bytes) = + decode::read_str_from_slice(bytes).map_err(error_report)?; + let (key, bytes) = decode::read_str_from_slice(bytes).map_err(error_report)?; + let (value, bytes) = decode::read_str_from_slice(bytes).map_err(error_report)?; + + if !bytes.is_empty() { + bail!("trailing bytes in encoded kvrecord. malformed") + } + + Ok(KvRecord { + namespace: namespace.to_owned(), + key: key.to_owned(), + value: Some(value.to_owned()), + }) + } + KV_VERSION => { + let mut bytes = decode::Bytes::new(&data.0); + + let nfields = decode::read_array_len(&mut bytes).map_err(error_report)?; + ensure!(nfields == 4, "too many entries in v1 kv record"); + + let bytes = bytes.remaining_slice(); + + let (namespace, bytes) = + decode::read_str_from_slice(bytes).map_err(error_report)?; + let (key, mut bytes) = decode::read_str_from_slice(bytes).map_err(error_report)?; + let has_value = decode::read_bool(&mut bytes).map_err(error_report)?; + + let (value, bytes) = if has_value { + let (value, bytes) = + decode::read_str_from_slice(bytes).map_err(error_report)?; + (Some(value.to_owned()), bytes) + } else { + (None, bytes) + }; + + if !bytes.is_empty() { + bail!("trailing bytes in encoded kvrecord. malformed") + } + + Ok(KvRecord { + namespace: namespace.to_owned(), + key: key.to_owned(), + value, + }) + } + _ => { + bail!("unknown version {version:?}") + } + } + } +} + +#[cfg(test)] +mod tests { + use super::{DecryptedData, KV_VERSION, KvRecord}; + + #[test] + fn encode_decode_some() { + let kv = KvRecord { + namespace: "foo".to_owned(), + key: "bar".to_owned(), + value: Some("baz".to_owned()), + }; + let snapshot = [ + 0x94, 0xa3, b'f', b'o', b'o', 0xa3, b'b', b'a', b'r', 0xc3, 0xa3, b'b', b'a', b'z', + ]; + + let encoded = kv.serialize().unwrap(); + let decoded = KvRecord::deserialize(&encoded, KV_VERSION).unwrap(); + + assert_eq!(encoded.0, &snapshot); + assert_eq!(decoded, kv); + } + + #[test] + fn encode_decode_none() { + let kv = KvRecord { + namespace: "foo".to_owned(), + key: "bar".to_owned(), + value: None, + }; + let snapshot = [0x94, 0xa3, b'f', b'o', b'o', 0xa3, b'b', b'a', b'r', 0xc2]; + + let encoded = kv.serialize().unwrap(); + let decoded = KvRecord::deserialize(&encoded, KV_VERSION).unwrap(); + + assert_eq!(encoded.0, &snapshot); + assert_eq!(decoded, kv); + } + + #[test] + fn decode_v0() { + let kv = KvRecord { + namespace: "foo".to_owned(), + key: "bar".to_owned(), + value: Some("baz".to_owned()), + }; + + let snapshot = vec![ + 0x93, 0xa3, b'f', b'o', b'o', 0xa3, b'b', b'a', b'r', 0xa3, b'b', b'a', b'z', + ]; + + let decoded = KvRecord::deserialize(&DecryptedData(snapshot), "v0").unwrap(); + + assert_eq!(decoded, kv); + } +} diff --git a/crates/atuin-nucleo/.github/workflows/ci.yml b/crates/atuin-nucleo/.github/workflows/ci.yml new file mode 100644 index 00000000..e478b6ae --- /dev/null +++ b/crates/atuin-nucleo/.github/workflows/ci.yml @@ -0,0 +1,83 @@ +name: CI +on: + pull_request: + push: + branches: + - master + +jobs: + check-msrv: + name: Check + strategy: + matrix: + toolchain: + - "1.65" + - stable + runs-on: ubuntu-latest + steps: + - name: Checkout sources + uses: actions/checkout@v4 + + - name: Install toolchain + uses: dtolnay/rust-toolchain@master + with: + toolchain: ${{ matrix.toolchain}} + + - uses: Swatinem/rust-cache@v2 + + - name: Run cargo check + run: cargo check + - name: Run cargo check withoult default features + run: cargo check --no-default-features + + test: + name: Test + runs-on: ubuntu-latest + steps: + - name: Checkout sources + uses: actions/checkout@v4 + + - name: Install stable toolchain + uses: dtolnay/rust-toolchain@stable + + - uses: Swatinem/rust-cache@v2 + + - name: Run cargo test + run: cargo test --workspace + + lints: + name: Lints + runs-on: ubuntu-latest + steps: + - name: Checkout sources + uses: actions/checkout@v4 + + - name: Install stable toolchain + uses: dtolnay/rust-toolchain@stable + with: + components: rustfmt, clippy + + - uses: Swatinem/rust-cache@v2 + + - name: Run cargo fmt + run: cargo fmt --all --check + + - name: Run cargo clippy + run: cargo clippy --workspace --all-targets -- -D warnings + - name: Run cargo clippy withoult default features + run: cargo clippy --workspace --all-targets --no-default-features -- -D warnings + + - name: Run cargo doc + run: cargo doc --no-deps --workspace --document-private-items + env: + RUSTDOCFLAGS: -D warnings + + typos: + name: Typos + runs-on: ubuntu-latest + steps: + - name: Checkout sources + uses: actions/checkout@v4 + + - name: Run typos + uses: crate-ci/typos@v1.16.11 diff --git a/crates/atuin-nucleo/.gitignore b/crates/atuin-nucleo/.gitignore new file mode 100644 index 00000000..73fab072 --- /dev/null +++ b/crates/atuin-nucleo/.gitignore @@ -0,0 +1,10 @@ +# Generated by Cargo +# will have compiled files and executables +debug/ +target/ + +# These are backup files generated by rustfmt +**/*.rs.bk + +# MSVC Windows builds of rustc generate these, which store debugging information +*.pdb diff --git a/crates/atuin-nucleo/CHANGELOG.md b/crates/atuin-nucleo/CHANGELOG.md new file mode 100644 index 00000000..4ea54b2c --- /dev/null +++ b/crates/atuin-nucleo/CHANGELOG.md @@ -0,0 +1,54 @@ +# Changelog + +# [0.5.0] - 2024-4-2 + +## **Breaking Changes** + +* `Injector::push` now passes a reference to the push value to the closure generating the columns + + +# [0.4.1] - 2024-3-11 + +## Bugfixes + +* crash when restarting picker with fast active stream + +# [0.4.0] - 2024-2-20 + +## Added + +* `active_injectors()` to retrieve the number of injectors that can potentially add new items to the matcher in the future. + +## Bugfixes + +* fix Unicode substring matcher expecting an exact match (rejecting trailing characters) +* fix crashes and false positives in unicode substring matcher + +# [0.3.0] - 2023-12-22 + +## **Breaking Changes** + +* Pattern API method now requires a Unicode `Normalization` strategy in addition to a `CaseMatching` strategy. + +## Bugfixes + +* avoid incorrect matches when searching for ASCII needles in a Unicode haystack +* correctly handle Unicode normalization when there are normalizable characters in the pattern, for example characters with umlauts +* when the needle is composed of a single char, return the score and index + of the best position instead of always returning the first matched character + in the haystack + +# [0.2.1] - 2023-09-02 + +## Bugfixes + +* ensure matcher runs on first call to `tick` + +# [0.2.0] - 2023-09-01 + +*initial public release* + + +[0.3.0]: https://github.com/helix-editor/nucleo/releases/tag/nucleo-v0.3.0 +[0.2.1]: https://github.com/helix-editor/nucleo/releases/tag/nucleo-v0.2.1 +[0.2.0]: https://github.com/helix-editor/nucleo/releases/tag/nucleo-v0.2.0 diff --git a/crates/atuin-nucleo/Cargo.lock b/crates/atuin-nucleo/Cargo.lock new file mode 100644 index 00000000..d31c11c0 --- /dev/null +++ b/crates/atuin-nucleo/Cargo.lock @@ -0,0 +1,319 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 3 + +[[package]] +name = "autocfg" +version = "1.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ace50bade8e6234aa140d9a2f552bbee1db4d353f69b8217bc503490fc1a9f26" + +[[package]] +name = "benches" +version = "0.1.0" +dependencies = [ + "brunch", + "fuzzy-matcher", + "nucleo", + "walkdir", +] + +[[package]] +name = "bitflags" +version = "2.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b048fb63fd8b5923fc5aa7b340d8e156aec7ec02f0c78fa8a6ddc2613f6f71de" + +[[package]] +name = "brunch" +version = "0.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3932d710d985d35c7b08e7e439a6ac8607aa8f619d373eb1f808578cd3cd56e5" +dependencies = [ + "dactyl", + "unicode-width", +] + +[[package]] +name = "cfg-if" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd" + +[[package]] +name = "crossbeam-deque" +version = "0.8.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "613f8cc01fe9cf1a3eb3d7f488fd2fa8388403e97039e2f73692932e291a770d" +dependencies = [ + "crossbeam-epoch", + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-epoch" +version = "0.9.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b82ac4a3c2ca9c3460964f020e1402edd5753411d7737aa39c3714ad1b5420e" +dependencies = [ + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-utils" +version = "0.8.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22ec99545bb0ed0ea7bb9b8e1e9122ea386ff8a48c0922e43f36d45ab09e0e80" + +[[package]] +name = "dactyl" +version = "0.7.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7ecad1ab94b1336be6cff409436ad9ceedb0afd52a85d54132189c2c3babb049" + +[[package]] +name = "either" +version = "1.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "60b1af1c220855b6ceac025d3f6ecdd2b7c4894bfe9cd9bda4fbb4bc7c0d4cf0" + +[[package]] +name = "fuzzy-matcher" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "54614a3312934d066701a80f20f15fa3b56d67ac7722b39eea5b4c9dd1d66c94" +dependencies = [ + "thread_local", +] + +[[package]] +name = "libc" +version = "0.2.167" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09d6582e104315a817dff97f75133544b2e094ee22447d2acf4a74e189ba06fc" + +[[package]] +name = "lock_api" +version = "0.4.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "07af8b9cdd281b7915f413fa73f29ebd5d55d0d3f0155584dade1ff18cea1b17" +dependencies = [ + "autocfg", + "scopeguard", +] + +[[package]] +name = "memchr" +version = "2.7.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "78ca9ab1a0babb1e7d5695e3530886289c18cf2f87ec19a575a0abdce112e3a3" + +[[package]] +name = "nucleo" +version = "0.5.0" +dependencies = [ + "nucleo-matcher", + "parking_lot", + "rayon", +] + +[[package]] +name = "nucleo-matcher" +version = "0.3.1" +dependencies = [ + "memchr", + "unicode-segmentation", +] + +[[package]] +name = "once_cell" +version = "1.20.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1261fe7e33c73b354eab43b1273a57c8f967d0391e80353e51f764ac02cf6775" + +[[package]] +name = "parking_lot" +version = "0.12.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f1bf18183cf54e8d6059647fc3063646a1801cf30896933ec2311622cc4b9a27" +dependencies = [ + "lock_api", + "parking_lot_core", +] + +[[package]] +name = "parking_lot_core" +version = "0.9.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e401f977ab385c9e4e3ab30627d6f26d00e2c73eef317493c4ec6d468726cf8" +dependencies = [ + "cfg-if", + "libc", + "redox_syscall", + "smallvec", + "windows-targets", +] + +[[package]] +name = "rayon" +version = "1.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b418a60154510ca1a002a752ca9714984e21e4241e804d32555251faf8b78ffa" +dependencies = [ + "either", + "rayon-core", +] + +[[package]] +name = "rayon-core" +version = "1.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1465873a3dfdaa8ae7cb14b4383657caab0b3e8a0aa9ae8e04b044854c8dfce2" +dependencies = [ + "crossbeam-deque", + "crossbeam-utils", +] + +[[package]] +name = "redox_syscall" +version = "0.5.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b6dfecf2c74bce2466cabf93f6664d6998a69eb21e39f4207930065b27b771f" +dependencies = [ + "bitflags", +] + +[[package]] +name = "same-file" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93fc1dc3aaa9bfed95e02e6eadabb4baf7e3078b0bd1b4d7b6b0b68378900502" +dependencies = [ + "winapi-util", +] + +[[package]] +name = "scopeguard" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" + +[[package]] +name = "smallvec" +version = "1.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3c5e1a9a646d36c3599cd173a41282daf47c44583ad367b8e6837255952e5c67" + +[[package]] +name = "thread_local" +version = "1.1.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b9ef9bad013ada3808854ceac7b46812a6465ba368859a37e2100283d2d719c" +dependencies = [ + "cfg-if", + "once_cell", +] + +[[package]] +name = "unicode-segmentation" +version = "1.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f6ccf251212114b54433ec949fd6a7841275f9ada20dddd2f29e9ceea4501493" + +[[package]] +name = "unicode-width" +version = "0.1.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7dd6e30e90baa6f72411720665d41d89b9a3d039dc45b8faea1ddd07f617f6af" + +[[package]] +name = "walkdir" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29790946404f91d9c5d06f9874efddea1dc06c5efe94541a7d6863108e3a5e4b" +dependencies = [ + "same-file", + "winapi-util", +] + +[[package]] +name = "winapi-util" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf221c93e13a30d793f7645a0e7762c55d169dbb0a49671918a2319d289b10bb" +dependencies = [ + "windows-sys", +] + +[[package]] +name = "windows-sys" +version = "0.59.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b" +dependencies = [ + "windows-targets", +] + +[[package]] +name = "windows-targets" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" +dependencies = [ + "windows_aarch64_gnullvm", + "windows_aarch64_msvc", + "windows_i686_gnu", + "windows_i686_gnullvm", + "windows_i686_msvc", + "windows_x86_64_gnu", + "windows_x86_64_gnullvm", + "windows_x86_64_msvc", +] + +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" + +[[package]] +name = "windows_aarch64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" + +[[package]] +name = "windows_i686_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b" + +[[package]] +name = "windows_i686_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" + +[[package]] +name = "windows_i686_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" + +[[package]] +name = "windows_x86_64_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" + +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" + +[[package]] +name = "windows_x86_64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" diff --git a/crates/atuin-nucleo/Cargo.toml b/crates/atuin-nucleo/Cargo.toml new file mode 100644 index 00000000..eb6315fb --- /dev/null +++ b/crates/atuin-nucleo/Cargo.toml @@ -0,0 +1,20 @@ +[package] +name = "nucleo" +description = "plug and play high performance fuzzy matcher" +authors = ["Pascal Kuthe <pascalkuthe@pm.me>"] +version = "0.5.0" +edition = "2021" +license = "MPL-2.0" +repository = "https://github.com/helix-editor/nucleo" +readme = "README.md" +exclude = ["/typos.toml", "/tarpaulin.toml"] + +[lib] + +[dependencies] +nucleo-matcher = { version = "0.3.1", path = "matcher" } +parking_lot = { version = "0.12.1", features = ["send_guard", "arc_lock"] } +rayon = "1.7.0" + +[workspace] +members = ["matcher", "bench"] diff --git a/crates/atuin-nucleo/LICENSE b/crates/atuin-nucleo/LICENSE new file mode 100644 index 00000000..a612ad98 --- /dev/null +++ b/crates/atuin-nucleo/LICENSE @@ -0,0 +1,373 @@ +Mozilla Public License Version 2.0 +================================== + +1. Definitions +-------------- + +1.1. "Contributor" + means each individual or legal entity that creates, contributes to + the creation of, or owns Covered Software. + +1.2. "Contributor Version" + means the combination of the Contributions of others (if any) used + by a Contributor and that particular Contributor's Contribution. + +1.3. "Contribution" + means Covered Software of a particular Contributor. + +1.4. "Covered Software" + means Source Code Form to which the initial Contributor has attached + the notice in Exhibit A, the Executable Form of such Source Code + Form, and Modifications of such Source Code Form, in each case + including portions thereof. + +1.5. "Incompatible With Secondary Licenses" + means + + (a) that the initial Contributor has attached the notice described + in Exhibit B to the Covered Software; or + + (b) that the Covered Software was made available under the terms of + version 1.1 or earlier of the License, but not also under the + terms of a Secondary License. + +1.6. "Executable Form" + means any form of the work other than Source Code Form. + +1.7. "Larger Work" + means a work that combines Covered Software with other material, in + a separate file or files, that is not Covered Software. + +1.8. "License" + means this document. + +1.9. "Licensable" + means having the right to grant, to the maximum extent possible, + whether at the time of the initial grant or subsequently, any and + all of the rights conveyed by this License. + +1.10. "Modifications" + means any of the following: + + (a) any file in Source Code Form that results from an addition to, + deletion from, or modification of the contents of Covered + Software; or + + (b) any new file in Source Code Form that contains any Covered + Software. + +1.11. "Patent Claims" of a Contributor + means any patent claim(s), including without limitation, method, + process, and apparatus claims, in any patent Licensable by such + Contributor that would be infringed, but for the grant of the + License, by the making, using, selling, offering for sale, having + made, import, or transfer of either its Contributions or its + Contributor Version. + +1.12. "Secondary License" + means either the GNU General Public License, Version 2.0, the GNU + Lesser General Public License, Version 2.1, the GNU Affero General + Public License, Version 3.0, or any later versions of those + licenses. + +1.13. "Source Code Form" + means the form of the work preferred for making modifications. + +1.14. "You" (or "Your") + means an individual or a legal entity exercising rights under this + License. For legal entities, "You" includes any entity that + controls, is controlled by, or is under common control with You. For + purposes of this definition, "control" means (a) the power, direct + or indirect, to cause the direction or management of such entity, + whether by contract or otherwise, or (b) ownership of more than + fifty percent (50%) of the outstanding shares or beneficial + ownership of such entity. + +2. License Grants and Conditions +-------------------------------- + +2.1. Grants + +Each Contributor hereby grants You a world-wide, royalty-free, +non-exclusive license: + +(a) under intellectual property rights (other than patent or trademark) + Licensable by such Contributor to use, reproduce, make available, + modify, display, perform, distribute, and otherwise exploit its + Contributions, either on an unmodified basis, with Modifications, or + as part of a Larger Work; and + +(b) under Patent Claims of such Contributor to make, use, sell, offer + for sale, have made, import, and otherwise transfer either its + Contributions or its Contributor Version. + +2.2. Effective Date + +The licenses granted in Section 2.1 with respect to any Contribution +become effective for each Contribution on the date the Contributor first +distributes such Contribution. + +2.3. Limitations on Grant Scope + +The licenses granted in this Section 2 are the only rights granted under +this License. No additional rights or licenses will be implied from the +distribution or licensing of Covered Software under this License. +Notwithstanding Section 2.1(b) above, no patent license is granted by a +Contributor: + +(a) for any code that a Contributor has removed from Covered Software; + or + +(b) for infringements caused by: (i) Your and any other third party's + modifications of Covered Software, or (ii) the combination of its + Contributions with other software (except as part of its Contributor + Version); or + +(c) under Patent Claims infringed by Covered Software in the absence of + its Contributions. + +This License does not grant any rights in the trademarks, service marks, +or logos of any Contributor (except as may be necessary to comply with +the notice requirements in Section 3.4). + +2.4. Subsequent Licenses + +No Contributor makes additional grants as a result of Your choice to +distribute the Covered Software under a subsequent version of this +License (see Section 10.2) or under the terms of a Secondary License (if +permitted under the terms of Section 3.3). + +2.5. Representation + +Each Contributor represents that the Contributor believes its +Contributions are its original creation(s) or it has sufficient rights +to grant the rights to its Contributions conveyed by this License. + +2.6. Fair Use + +This License is not intended to limit any rights You have under +applicable copyright doctrines of fair use, fair dealing, or other +equivalents. + +2.7. Conditions + +Sections 3.1, 3.2, 3.3, and 3.4 are conditions of the licenses granted +in Section 2.1. + +3. Responsibilities +------------------- + +3.1. Distribution of Source Form + +All distribution of Covered Software in Source Code Form, including any +Modifications that You create or to which You contribute, must be under +the terms of this License. You must inform recipients that the Source +Code Form of the Covered Software is governed by the terms of this +License, and how they can obtain a copy of this License. You may not +attempt to alter or restrict the recipients' rights in the Source Code +Form. + +3.2. Distribution of Executable Form + +If You distribute Covered Software in Executable Form then: + +(a) such Covered Software must also be made available in Source Code + Form, as described in Section 3.1, and You must inform recipients of + the Executable Form how they can obtain a copy of such Source Code + Form by reasonable means in a timely manner, at a charge no more + than the cost of distribution to the recipient; and + +(b) You may distribute such Executable Form under the terms of this + License, or sublicense it under different terms, provided that the + license for the Executable Form does not attempt to limit or alter + the recipients' rights in the Source Code Form under this License. + +3.3. Distribution of a Larger Work + +You may create and distribute a Larger Work under terms of Your choice, +provided that You also comply with the requirements of this License for +the Covered Software. If the Larger Work is a combination of Covered +Software with a work governed by one or more Secondary Licenses, and the +Covered Software is not Incompatible With Secondary Licenses, this +License permits You to additionally distribute such Covered Software +under the terms of such Secondary License(s), so that the recipient of +the Larger Work may, at their option, further distribute the Covered +Software under the terms of either this License or such Secondary +License(s). + +3.4. Notices + +You may not remove or alter the substance of any license notices +(including copyright notices, patent notices, disclaimers of warranty, +or limitations of liability) contained within the Source Code Form of +the Covered Software, except that You may alter any license notices to +the extent required to remedy known factual inaccuracies. + +3.5. Application of Additional Terms + +You may choose to offer, and to charge a fee for, warranty, support, +indemnity or liability obligations to one or more recipients of Covered +Software. However, You may do so only on Your own behalf, and not on +behalf of any Contributor. You must make it absolutely clear that any +such warranty, support, indemnity, or liability obligation is offered by +You alone, and You hereby agree to indemnify every Contributor for any +liability incurred by such Contributor as a result of warranty, support, +indemnity or liability terms You offer. You may include additional +disclaimers of warranty and limitations of liability specific to any +jurisdiction. + +4. Inability to Comply Due to Statute or Regulation +--------------------------------------------------- + +If it is impossible for You to comply with any of the terms of this +License with respect to some or all of the Covered Software due to +statute, judicial order, or regulation then You must: (a) comply with +the terms of this License to the maximum extent possible; and (b) +describe the limitations and the code they affect. Such description must +be placed in a text file included with all distributions of the Covered +Software under this License. Except to the extent prohibited by statute +or regulation, such description must be sufficiently detailed for a +recipient of ordinary skill to be able to understand it. + +5. Termination +-------------- + +5.1. The rights granted under this License will terminate automatically +if You fail to comply with any of its terms. However, if You become +compliant, then the rights granted under this License from a particular +Contributor are reinstated (a) provisionally, unless and until such +Contributor explicitly and finally terminates Your grants, and (b) on an +ongoing basis, if such Contributor fails to notify You of the +non-compliance by some reasonable means prior to 60 days after You have +come back into compliance. Moreover, Your grants from a particular +Contributor are reinstated on an ongoing basis if such Contributor +notifies You of the non-compliance by some reasonable means, this is the +first time You have received notice of non-compliance with this License +from such Contributor, and You become compliant prior to 30 days after +Your receipt of the notice. + +5.2. If You initiate litigation against any entity by asserting a patent +infringement claim (excluding declaratory judgment actions, +counter-claims, and cross-claims) alleging that a Contributor Version +directly or indirectly infringes any patent, then the rights granted to +You by any and all Contributors for the Covered Software under Section +2.1 of this License shall terminate. + +5.3. In the event of termination under Sections 5.1 or 5.2 above, all +end user license agreements (excluding distributors and resellers) which +have been validly granted by You or Your distributors under this License +prior to termination shall survive termination. + +************************************************************************ +* * +* 6. Disclaimer of Warranty * +* ------------------------- * +* * +* Covered Software is provided under this License on an "as is" * +* basis, without warranty of any kind, either expressed, implied, or * +* statutory, including, without limitation, warranties that the * +* Covered Software is free of defects, merchantable, fit for a * +* particular purpose or non-infringing. The entire risk as to the * +* quality and performance of the Covered Software is with You. * +* Should any Covered Software prove defective in any respect, You * +* (not any Contributor) assume the cost of any necessary servicing, * +* repair, or correction. This disclaimer of warranty constitutes an * +* essential part of this License. No use of any Covered Software is * +* authorized under this License except under this disclaimer. * +* * +************************************************************************ + +************************************************************************ +* * +* 7. Limitation of Liability * +* -------------------------- * +* * +* Under no circumstances and under no legal theory, whether tort * +* (including negligence), contract, or otherwise, shall any * +* Contributor, or anyone who distributes Covered Software as * +* permitted above, be liable to You for any direct, indirect, * +* special, incidental, or consequential damages of any character * +* including, without limitation, damages for lost profits, loss of * +* goodwill, work stoppage, computer failure or malfunction, or any * +* and all other commercial damages or losses, even if such party * +* shall have been informed of the possibility of such damages. This * +* limitation of liability shall not apply to liability for death or * +* personal injury resulting from such party's negligence to the * +* extent applicable law prohibits such limitation. Some * +* jurisdictions do not allow the exclusion or limitation of * +* incidental or consequential damages, so this exclusion and * +* limitation may not apply to You. * +* * +************************************************************************ + +8. Litigation +------------- + +Any litigation relating to this License may be brought only in the +courts of a jurisdiction where the defendant maintains its principal +place of business and such litigation shall be governed by laws of that +jurisdiction, without reference to its conflict-of-law provisions. +Nothing in this Section shall prevent a party's ability to bring +cross-claims or counter-claims. + +9. Miscellaneous +---------------- + +This License represents the complete agreement concerning the subject +matter hereof. If any provision of this License is held to be +unenforceable, such provision shall be reformed only to the extent +necessary to make it enforceable. Any law or regulation which provides +that the language of a contract shall be construed against the drafter +shall not be used to construe this License against a Contributor. + +10. Versions of the License +--------------------------- + +10.1. New Versions + +Mozilla Foundation is the license steward. Except as provided in Section +10.3, no one other than the license steward has the right to modify or +publish new versions of this License. Each version will be given a +distinguishing version number. + +10.2. Effect of New Versions + +You may distribute the Covered Software under the terms of the version +of the License under which You originally received the Covered Software, +or under the terms of any subsequent version published by the license +steward. + +10.3. Modified Versions + +If you create software not governed by this License, and you want to +create a new license for such software, you may create and use a +modified version of this License if you rename the license and remove +any references to the name of the license steward (except to note that +such modified license differs from this License). + +10.4. Distributing Source Code Form that is Incompatible With Secondary +Licenses + +If You choose to distribute Source Code Form that is Incompatible With +Secondary Licenses under the terms of this version of the License, the +notice described in Exhibit B of this License must be attached. + +Exhibit A - Source Code Form License Notice +------------------------------------------- + + This Source Code Form is subject to the terms of the Mozilla Public + License, v. 2.0. If a copy of the MPL was not distributed with this + file, You can obtain one at http://mozilla.org/MPL/2.0/. + +If it is not possible or desirable to put the notice in a particular +file, then You may include the notice in a location (such as a LICENSE +file in a relevant directory) where a recipient would be likely to look +for such a notice. + +You may add additional accurate notices of copyright ownership. + +Exhibit B - "Incompatible With Secondary Licenses" Notice +--------------------------------------------------------- + + This Source Code Form is "Incompatible With Secondary Licenses", as + defined by the Mozilla Public License, v. 2.0. diff --git a/crates/atuin-nucleo/README.md b/crates/atuin-nucleo/README.md new file mode 100644 index 00000000..e617cfa0 --- /dev/null +++ b/crates/atuin-nucleo/README.md @@ -0,0 +1,114 @@ +# Nucleo + + +`nucleo` is a highly performant fuzzy matcher written in Rust. It aims to fill the same use case as `fzf` and `skim`. Compared to `fzf` `nucleo` has a significantly faster matching algorithm. This mainly makes a difference when matching patterns with low selectivity on many items. An (unscientific) comparison is shown in the benchmark section below. + +> Note: If you are looking for a replacement of the `fuzzy-matcher` crate and not a fully managed fuzzy picker, you should use the [`nucleo-matcher`](https://crates.io/crates/nucleo-matcher) crate. + +`nucleo` uses the exact **same scoring system as fzf**. That means you should get the same ranking quality (or better) as you are used to from fzf. However, `nucleo` has a more faithful implementation of the Smith-Waterman algorithm which is normally used in DNA sequence alignment (see https://www.cs.cmu.edu/~ckingsf/bioinfo-lectures/gaps.pdf) with two separate matrices (instead of one like fzf). This means that `nucleo` finds the optimal match more often. For example if you match `foo` in `xf foo` `nucleo` will match `x__foo` but `fzf` will match `xf_oo` (you can increase the word length the result will stay the same). The former is the more intuitive match and has a higher score according to the ranking system that both `nucleo` and fzf. + +**Compared to `skim`** (and the `fuzzy-matcher` crate) `nucleo` has an even larger performance advantage and is often around **six times faster** (see benchmarks below). Furthermore, the bonus system used by nucleo and fzf is (in my opinion) more consistent/superior. `nucleo` also handles non-ascii text much better. (`skim`s bonus system and even case insensitivity only work for ASCII). + +Nucleo also handles Unicode graphemes more correctly. `Fzf` and `skim` both operate on Unicode code points (chars). That means that multi codepoint graphemes can have weird effects (match multiple times, weirdly change the score, ...). `nucleo` will always use the first codepoint of the grapheme for matching instead (and reports grapheme indices, so they can be highlighted correctly). + +## Status + +Nucleo is used in the helix-editor and therefore has a large user base with lots of real world testing. The core matcher implementation is considered complete and is unlikely to see major changes. The `nucleo-matcher` crate is finished and ready for widespread use, breaking changes should be very rare (a 1.0 release should not be far away). + +While the high level `nucleo` crate also works well (and is also used in helix), there are still additional features that will be added in the future. The high level crate also need better documentation and will likely see a few API changes in the future. + +## Benchmarks + +> WIP currently more of a demonstration than a comprehensive benchmark suit +> most notably scientific comparisons with `fzf` are missing (a pain because it can't be called as a library) + + +### Matcher micro benchmarks + +Benchmark comparing the runtime of various patterns matched against all files in the source of the linux kernel. Repeat on your system with `BENCHMARK_DIR=<path_to_linux> cargo run -p benches --release` (you can specify an empty directory and the kernel is cloned automatically). + +Method | Mean | Samples +-----------------------|-----------|----------- +nucleo "never_matches" | 2.30 ms |2,493/2,500 +skim "never_matches" | 17.44 ms | 574/574 +nucleo "copying" | 2.12 ms |2,496/2,500 +skim "copying" | 16.85 ms | 593/594 +nucleo "/doc/kernel" | 2.59 ms |2,499/2,500 +skim "/doc/kernel" | 18.32 ms | 546/546 +nucleo "//.h" | 9.53 ms |1,049/1,049 +skim "//.h" | 35.46 ms | 282/282 + + +### Comparison with fzf + +For example in the following two screencasts the pattern `///.` is pasted into `fzf` and `nucleo` (both with about 3 million items open). + +`fzf` takes a while to filter the text (about 1 second) while `nucleo` has barely any noticeable delay (a single frame in the screencast so about 1/30 seconds). This comparison was made on a very beefy CPU (Ryzen 5950x) so on slower systems the difference may be larger: + +[](https://asciinema.org/a/600517) +[](https://asciinema.org/a/600516) + + + +# Future Work + +* [x] merge integration into helix +* [ ] build a standalone CLI application + * [ ] reach feature parity with `fzf` (mostly `--no-sort` and `--tac`) + * [ ] add a way to allow columnar matching +* [ ] expose C API so both the high level API and the matching algorithm itself can be used in other applications (like various nvim plugins) + +# Naming + +The name `nucleo` plays on the fact that the `Smith-Waterman` algorithm (that it's based on) was originally developed for matching DNA/RNA sequences. The elements of DNA/RNA that are matched are called *nucleotides* which was shortened to `nucleo` here. + +The name also indicates its close relationship with the *helix* editor (sticking with the DNA theme). + +# Implementation Details + +> This is only intended for those interested and will not be relevant to most people. I plan to turn this into a blog post when I have more time + +<!-- Nucleo matching algorithm has `O(N-M)` space complexity while ranking/filtering (and not computing indices) compared to the `O(MN)` space complexity of fzf. --> + +<!-- Furthermore, `nucleo` also features fully lock-free multithreaded streaming so if used as a library its possible to performantly scale streaming to a practically unlimited number of producer threads (for example running `ignore` or `jwalk` across all cores) without any buffering or other additional logic. --> + + +The fuzzy matching algorithm is based on the `Smith-Waterman` (with affine gaps) as described in https://www.cs.cmu.edu/~ckingsf/bioinfo-lectures/gaps.pdf (TODO: explain). `Nucleo` faithfully implements this algorithm and therefore has two separate matrices. However, by precomputing the next `m-matrix` row we can avoid storing the p-matrix at all and instead just store the value in a variable as we iterate the row. + +Nucleo also never really stores the `m-matrix` instead we only ever store the current row (which simultaneously serves as the next row). During index calculation a full matrix is however required to backtrack which indices were actually matched. We only store two bools here (to indicate where we came from in the matrix). + +By comparison `skim` stores the full p and m matrix in that case. `fzf` always allocates a full `mn` matrix (even during matching!). + +`nucleo`s' matrix is only width `n-m+1` instead of width `n`. This comes from the observation that the `p` char requires `p-1` chars before it and `m-p` chars after it, so there are always `p-1 + m-p = m+1` chars that can never match the current char. This works especially well with only using a single row because the first relevant char is always at the same position even though it's technically further to the right. This is particularly nice because we precalculate the m-matrix row. The m-matrix is computed from diagonal elements, so the precalculated values stay in the same matrix cell. + +Compared to `skim` nucleo does couple simpler (but arguably even more impactful) optimizations: +* *Presegment Unicode*: Unicode segmentation is somewhat slow and matcher will filter the same elements quite often so only doing it once is nice. It also prevents a very common source of bugs (mixing of char indices which we use here and utf8 indices) and makes the code a lot simpler as a result. Fzf does the same. +* *Aggressive prefiltering*: Especially for ASCII this works very well, but we also do this for Unicode to a lesser extent. This ensures we reject non-matching haystacks as fast as possible. Usually most haystacks will not match when fuzzy matching large lists so having fast path for that case is a huge win. +* *Special-case ASCII*: 90% of practical text is ASCII. ASCII can be stored as bytes instead of `chars`, so cache locality is improved a lot, and we can use `memchar` for superfast prefilters (even case-insensitive prefilter are possible that way) +* *Fallback for very long matches*: We fall back to greedy matcher which runs in `O(N)` (and `O(1)` space complexity) to avoid the `O(mn)` blowup for large matches. This is fzfs old algorithm and yields decent (but not great) results. + + + +<!-- There is a misunderstanding in both `skim` and fzf. Basically what they do is give a bonus to each character (like word boundaries). That makes senes and is reasonable, but the problem is that they use the **maximum bonus** when multiple chars match in sequence. That means that the bonus of a character depends on which characters exactly matched around it. But the fundamental assumption of this algorithm (and why it doesn't require backtracking) is that the score of each character is independent of what other chars matched (this is the difference between the affine gap and the generic gap case shown in the paper too). During fuzzing I found many cases where this mechanism leads to a non-optimal match being reported (so the sort order and fuzzy indices would be wrong). In my testing removing this mechanism and slightly tweaking the bonus calculation results in similar match quality but made sure the algorithm always worked correctly (and removed a bunch of weird edges cases). --> + <!-- * [ ] it seems this makes us overemphasize word boundaries for small search strings, this is likely okay as the consecutive bonus wins fairly quickly. Maybe we just do a greedy search for the first 2 chars to reduce visual noise? --> +<!-- * [x] substring/prefix/postfix/exact matcher --> +<!-- * [ ] case mismatch penalty. This doesn't seem like a good idea to me. `FZF` doesn't do this (only skin), smart case should cover most cases. .would be nice for fully case-insensitive matching without smart case like in autocompletion tough. Realistically there won't be more than 3 items that are identical with different casing tough, so I don't think it matters too much. It is a bit annoying to implement since you can no longer pre-normalize queries(or need two queries) :/ --> +<!-- * [ ] high level API (worker thread, query parsing, sorting), in progress --> + <!-- * apparently sorting is superfast (at most 5% of match time for `nucleo` matcher with a highly selective query, otherwise its completely negligible compared to fuzzy matching). All the bending over backwards `fzf` does (and `skim` copied but way worse) seems a little silly. I think `fzf` does it because go doesn't have a good parallel sort. `Fzf` divides the matches into a couple fairly large chunks and sorts those on each worker thread and then lazily merges the result. That makes the sorting without the merging `Nlog(N/M)` which is basically equivalent for large `N` and small `M` as is the case here. At least its parallel tough. In rust we have a great pattern defeating parallel quicksort tough (rayon) which is way easier. --> + <!-- * [x] basic implementation (workers, streaming, invalidation) --> + <!-- * [x] verify it actually works --> + <!-- * [x] query paring --> + <!-- * [x] hook up to helix --> + <!-- * [x] currently I simply use a tick system (called on every redraw), together with a redraw/tick nofication (ideally debounced) is that enough? yes works nicely --> + <!-- * [x] for streaming callers should buffer their data. Can we provide a better API for that beyond what is currently there? yes lock-free stream --> + <!-- * [ ] cleanup code, improve API --> + <!-- * [ ] write docs --> + +<!-- * tests --> + <!-- * [x] fuzz the fuzzy matcher --> + <!-- * [x] port the full `fzf` test suite for fuzzy matching --> + <!-- * [ ] port the full `skim` test suite for fuzzy matching --> + <!-- * [ ] highlevel API --> + <!-- * [~] test substring/exact/prefix/postfix match --> + <!-- * [ ] coverage report (fuzzy matcher was at 86%) --> + diff --git a/crates/atuin-nucleo/bench/Cargo.toml b/crates/atuin-nucleo/bench/Cargo.toml new file mode 100644 index 00000000..0dfb81d5 --- /dev/null +++ b/crates/atuin-nucleo/bench/Cargo.toml @@ -0,0 +1,12 @@ +[package] +name = "benches" +version = "0.1.0" +edition = "2021" + +# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html + +[dependencies] +nucleo = { version = "*", path = "../" } +brunch = "0.5.0" +fuzzy-matcher = "0.3.7" +walkdir = "2" diff --git a/crates/atuin-nucleo/bench/src/main.rs b/crates/atuin-nucleo/bench/src/main.rs new file mode 100644 index 00000000..bc77b03d --- /dev/null +++ b/crates/atuin-nucleo/bench/src/main.rs @@ -0,0 +1,76 @@ +use std::hint::black_box; +use std::path::PathBuf; +use std::process::Command; + +use brunch::{Bench, Benches}; +use fuzzy_matcher::FuzzyMatcher; +use nucleo::{Utf32Str, Utf32String}; + +fn bench_dir() -> PathBuf { + std::env::var_os("BENCHMARK_DIR") + .expect("the BENCHMARK_DIR must be set to the directory to traverse for the benchmark") + .into() +} + +fn checkout_linux_if_needed() { + let linux_dir = bench_dir(); + if !linux_dir.exists() { + println!("will git clone linux..."); + let output = Command::new("git") + .arg("clone") + .arg("https://github.com/BurntSushi/linux.git") + .arg("--depth") + .arg("1") + .arg("--branch") + .arg("master") + .arg("--single-branch") + .arg(&linux_dir) + .stdout(std::process::Stdio::inherit()) + .status() + .expect("failed to git clone linux"); + println!("did git clone linux...{:?}", output); + } +} + +fn main() { + checkout_linux_if_needed(); + let dir = bench_dir(); + let paths: (Vec<Utf32String>, Vec<String>) = walkdir::WalkDir::new(dir) + .into_iter() + .filter_map(|path| { + let dent = path.ok()?; + let path = dent.into_path().to_string_lossy().into_owned(); + Some((path.as_str().into(), path)) + }) + .unzip(); + let mut nucleo = nucleo::Matcher::new(nucleo::Config::DEFAULT.match_paths()); + let skim = fuzzy_matcher::skim::SkimMatcherV2::default(); + + // TODO: unicode? + let needles = ["never_matches", "copying", "/doc/kernel", "//.h"]; + // Announce that we've started. + ::std::eprint!("\x1b[1;38;5;199mStarting:\x1b[0m Running benchmark(s). Stand by!\n\n"); + let mut benches = Benches::default(); + // let mut scores = Vec::with_capacity(paths.0.len()); + for needle in needles { + println!("running {needle:?}..."); + benches.push(Bench::new(format!("nucleo {needle:?}")).run(|| { + // scores.clear(); + // scores.extend(paths.0.iter().filter_map(|haystack| { + for haystack in &paths.0 { + black_box( + nucleo.fuzzy_match(haystack.slice(..), Utf32Str::Ascii(needle.as_bytes())), + ); + } + // })); + // scores.sort_unstable(); + })); + benches.push(Bench::new(format!("skim {needle:?}")).run(|| { + for haystack in &paths.1 { + let res = skim.fuzzy_match(haystack, needle); + let _ = black_box(res); + } + })); + } + benches.finish(); +} diff --git a/crates/atuin-nucleo/matcher/Cargo.toml b/crates/atuin-nucleo/matcher/Cargo.toml new file mode 100644 index 00000000..4b90ddbb --- /dev/null +++ b/crates/atuin-nucleo/matcher/Cargo.toml @@ -0,0 +1,19 @@ +[package] +name = "nucleo-matcher" +description = "plug and play high performance fuzzy matcher" +authors = ["Pascal Kuthe <pascalkuthe@pm.me>"] +version = "0.3.1" +edition = "2021" +license = "MPL-2.0" +repository = "https://github.com/helix-editor/nucleo" +readme = "../README.md" + +[dependencies] +memchr = "2.5.0" +unicode-segmentation = { version = "1.10", optional = true } + +[features] +default = ["unicode-normalization", "unicode-casefold", "unicode-segmentation"] +unicode-normalization = [] +unicode-casefold = [] +unicode-segmentation = ["dep:unicode-segmentation"] diff --git a/crates/atuin-nucleo/matcher/LICENSE b/crates/atuin-nucleo/matcher/LICENSE new file mode 120000 index 00000000..ea5b6064 --- /dev/null +++ b/crates/atuin-nucleo/matcher/LICENSE @@ -0,0 +1 @@ +../LICENSE
\ No newline at end of file diff --git a/crates/atuin-nucleo/matcher/fuzz.sh b/crates/atuin-nucleo/matcher/fuzz.sh new file mode 100755 index 00000000..d3ffa2c9 --- /dev/null +++ b/crates/atuin-nucleo/matcher/fuzz.sh @@ -0,0 +1,3 @@ +#!/usr/bin/env bash + +cargo +nightly fuzz "${1}" fuzz_target_1 "${@:2:99}" diff --git a/crates/atuin-nucleo/matcher/fuzz/.gitignore b/crates/atuin-nucleo/matcher/fuzz/.gitignore new file mode 100644 index 00000000..1a45eee7 --- /dev/null +++ b/crates/atuin-nucleo/matcher/fuzz/.gitignore @@ -0,0 +1,4 @@ +target +corpus +artifacts +coverage diff --git a/crates/atuin-nucleo/matcher/fuzz/Cargo.toml b/crates/atuin-nucleo/matcher/fuzz/Cargo.toml new file mode 100644 index 00000000..1b9d8a7f --- /dev/null +++ b/crates/atuin-nucleo/matcher/fuzz/Cargo.toml @@ -0,0 +1,29 @@ +[package] +name = "fzf_oxide-fuzz" +version = "0.0.0" +publish = false +edition = "2021" + +[package.metadata] +cargo-fuzz = true + +[dependencies] +libfuzzer-sys = "0.4" +arbitrary = { version = "1", features = ["derive"] } + +[dependencies.fzf_oxide] +path = ".." + +# Prevent this from interfering with workspaces +[workspace] +members = ["."] + +[profile.release] +debug = 1 + +[[bin]] +name = "fuzz_target_1" +path = "fuzz_targets/fuzz_target_1.rs" +test = false +doc = false + diff --git a/crates/atuin-nucleo/matcher/fuzz/fuzz_targets/fuzz_target_1.rs b/crates/atuin-nucleo/matcher/fuzz/fuzz_targets/fuzz_target_1.rs new file mode 100644 index 00000000..d9df7d36 --- /dev/null +++ b/crates/atuin-nucleo/matcher/fuzz/fuzz_targets/fuzz_target_1.rs @@ -0,0 +1,78 @@ +#![no_main] + +use fzf_oxide::{chars, Matcher, MatcherConfig, Utf32Str}; +use libfuzzer_sys::arbitrary::Arbitrary; +use libfuzzer_sys::fuzz_target; + +#[derive(Arbitrary, Debug)] +pub struct Input<'a> { + haystack: &'a str, + needle: &'a str, + ignore_case: bool, + normalize: bool, +} + +fuzz_target!(|data: Input<'_>| { + let mut data = data; + let mut config = MatcherConfig::DEFAULT; + config.ignore_case = data.ignore_case; + config.normalize = data.normalize; + let mut matcher = Matcher::new(config); + let mut indices_optimal = Vec::new(); + let mut indices_greedy = Vec::new(); + let mut needle_buf = Vec::new(); + let mut haystack_buf = Vec::new(); + let normalize = |mut c: char| { + if config.normalize { + c = chars::normalize(c); + } + if config.ignore_case { + c = chars::to_lower_case(c); + } + c + }; + let needle: String = data.needle.chars().map(normalize).collect(); + let needle_chars: Vec<_> = needle.chars().collect(); + let needle = Utf32Str::new(&needle, &mut needle_buf); + let haystack = Utf32Str::new(data.haystack, &mut haystack_buf); + + let greedy_score = matcher.fuzzy_indices_greedy(haystack, needle, &mut indices_greedy); + if greedy_score.is_some() { + let match_chars: Vec<_> = indices_greedy + .iter() + .map(|&i| normalize(haystack.get(i))) + .collect(); + assert_eq!( + match_chars, needle_chars, + "failed match, found {indices_greedy:?} {match_chars:?} (greedy)" + ); + } + let optimal_score = matcher.fuzzy_indices(haystack, needle, &mut indices_optimal); + if optimal_score.is_some() { + let match_chars: Vec<_> = indices_optimal + .iter() + .map(|&i| normalize(haystack.get(i))) + .collect(); + assert_eq!( + match_chars, needle_chars, + "failed match, found {indices_optimal:?} {match_chars:?}" + ); + } + match (greedy_score, optimal_score) { + (None, Some(score)) => unreachable!("optimal matched {score} but greedy did not match"), + (Some(score), None) => unreachable!("greedy matched {score} but optimal did not match"), + (Some(greedy), Some(optimal)) => { + assert!( + greedy <= optimal, + "optimal score must be atleast the same as greedy score {greedy} {optimal}" + ); + if indices_greedy == indices_optimal { + assert_eq!( + greedy, optimal, + "if matching same char greedy and optimal score should be identical" + ) + } + } + (None, None) => (), + } +}); diff --git a/crates/atuin-nucleo/matcher/generate_case_fold_table.sh b/crates/atuin-nucleo/matcher/generate_case_fold_table.sh new file mode 100755 index 00000000..32a26697 --- /dev/null +++ b/crates/atuin-nucleo/matcher/generate_case_fold_table.sh @@ -0,0 +1,13 @@ +#!/usr/bin/env bash +set -e + +dir=$(pwd) +mkdir /tmp/ucd-15.0.0 +cd /tmp/ucd-15.0.0 +curl -LO https://www.unicode.org/Public/zipped/15.0.0/UCD.zip +unzip UCD.zip + +cd "${dir}" +cargo install ucd-generate +ucd-generate case-folding-simple /tmp/ucd-15.0.0 --chars > src/chars/case_fold.rs +rm -rf /tmp/ucd-15.0.0 diff --git a/crates/atuin-nucleo/matcher/src/chars.rs b/crates/atuin-nucleo/matcher/src/chars.rs new file mode 100644 index 00000000..d13a2466 --- /dev/null +++ b/crates/atuin-nucleo/matcher/src/chars.rs @@ -0,0 +1,207 @@ +//! Utilities for working with (unicode) characters/codepoints + +use std::fmt::{self, Debug, Display}; + +#[cfg(feature = "unicode-casefold")] +use crate::chars::case_fold::CASE_FOLDING_SIMPLE; +use crate::Config; + +//autogenerated by generate-ucd +#[allow(warnings)] +#[rustfmt::skip] +#[cfg(feature = "unicode-casefold")] +mod case_fold; +#[cfg(feature = "unicode-normalization")] +mod normalize; + +pub(crate) trait Char: Copy + Eq + Ord + fmt::Display { + const ASCII: bool; + fn char_class(self, config: &Config) -> CharClass; + fn char_class_and_normalize(self, config: &Config) -> (Self, CharClass); + fn normalize(self, config: &Config) -> Self; +} + +/// repr tansparent wrapper around u8 with better formatting and `PartialEq<char>` implementation +#[repr(transparent)] +#[derive(PartialEq, Eq, PartialOrd, Ord, Clone, Copy)] +pub(crate) struct AsciiChar(pub u8); + +impl AsciiChar { + pub fn cast(bytes: &[u8]) -> &[AsciiChar] { + unsafe { &*(bytes as *const [u8] as *const [AsciiChar]) } + } +} + +impl fmt::Display for AsciiChar { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + Display::fmt(&(self.0 as char), f) + } +} + +impl PartialEq<AsciiChar> for char { + fn eq(&self, other: &AsciiChar) -> bool { + other.0 as char == *self + } +} + +impl Char for AsciiChar { + const ASCII: bool = true; + #[inline] + fn char_class(self, config: &Config) -> CharClass { + let c = self.0; + // using manual if conditions instead optimizes better + if c >= b'a' && c <= b'z' { + CharClass::Lower + } else if c >= b'A' && c <= b'Z' { + CharClass::Upper + } else if c >= b'0' && c <= b'9' { + CharClass::Number + } else if c.is_ascii_whitespace() { + CharClass::Whitespace + } else if config.delimiter_chars.contains(&c) { + CharClass::Delimiter + } else { + CharClass::NonWord + } + } + + #[inline(always)] + fn char_class_and_normalize(mut self, config: &Config) -> (Self, CharClass) { + let char_class = self.char_class(config); + if config.ignore_case && char_class == CharClass::Upper { + self.0 += 32 + } + (self, char_class) + } + + #[inline(always)] + fn normalize(mut self, config: &Config) -> Self { + if config.ignore_case && self.0 >= b'A' && self.0 <= b'Z' { + self.0 += 32 + } + self + } +} +fn char_class_non_ascii(c: char) -> CharClass { + if c.is_lowercase() { + CharClass::Lower + } else if is_upper_case(c) { + CharClass::Upper + } else if c.is_numeric() { + CharClass::Number + } else if c.is_alphabetic() { + CharClass::Letter + } else if c.is_whitespace() { + CharClass::Whitespace + } else { + CharClass::NonWord + } +} +impl Char for char { + const ASCII: bool = false; + #[inline(always)] + fn char_class(self, config: &Config) -> CharClass { + if self.is_ascii() { + return AsciiChar(self as u8).char_class(config); + } + char_class_non_ascii(self) + } + + #[inline(always)] + fn char_class_and_normalize(mut self, config: &Config) -> (Self, CharClass) { + if self.is_ascii() { + let (c, class) = AsciiChar(self as u8).char_class_and_normalize(config); + return (c.0 as char, class); + } + let char_class = char_class_non_ascii(self); + #[cfg(feature = "unicode-casefold")] + let mut case_fold = char_class == CharClass::Upper; + #[cfg(feature = "unicode-normalization")] + if config.normalize { + self = normalize::normalize(self); + case_fold = true + } + #[cfg(feature = "unicode-casefold")] + if case_fold && config.ignore_case { + self = CASE_FOLDING_SIMPLE + .binary_search_by_key(&self, |(upper, _)| *upper) + .map_or(self, |idx| CASE_FOLDING_SIMPLE[idx].1) + } + (self, char_class) + } + + #[inline(always)] + fn normalize(mut self, config: &Config) -> Self { + #[cfg(feature = "unicode-normalization")] + if config.normalize { + self = normalize::normalize(self); + } + #[cfg(feature = "unicode-casefold")] + if config.ignore_case { + self = to_lower_case(self) + } + self + } +} + +#[cfg(feature = "unicode-normalization")] +pub use normalize::normalize; +#[cfg(feature = "unicode-segmentation")] +use unicode_segmentation::UnicodeSegmentation; + +/// Converts a character to lower case using simple unicode case folding +#[cfg(feature = "unicode-casefold")] +#[inline(always)] +pub fn to_lower_case(c: char) -> char { + CASE_FOLDING_SIMPLE + .binary_search_by_key(&c, |(upper, _)| *upper) + .map_or(c, |idx| CASE_FOLDING_SIMPLE[idx].1) +} + +/// Checks if a character is upper case according to simple unicode case folding. +/// if the `unicode-casefold` feature is disable the equivalent std function is used +#[inline(always)] +pub fn is_upper_case(c: char) -> bool { + #[cfg(feature = "unicode-casefold")] + let val = CASE_FOLDING_SIMPLE + .binary_search_by_key(&c, |(upper, _)| *upper) + .is_ok(); + #[cfg(not(feature = "unicode-casefold"))] + let val = c.is_uppercase(); + val +} + +#[derive(Debug, Eq, PartialEq, PartialOrd, Ord, Copy, Clone, Hash)] +pub(crate) enum CharClass { + Whitespace, + NonWord, + Delimiter, + Lower, + Upper, + Letter, + Number, +} + +/// Nucleo cannot match graphemes as single units. To work around +/// that we only use the first codepoint of each grapheme. This +/// iterator returns the first character of each unicode grapheme +/// in a string and is used for constructing `Utf32Str(ing)`. +pub fn graphemes(text: &str) -> impl Iterator<Item = char> + '_ { + #[cfg(feature = "unicode-segmentation")] + let res = text.graphemes(true).map(|grapheme| { + // we need to special-case this check since `\r\n` is a single grapheme and is + // therefore the exception to the rule that normalization of a grapheme should + // map to the first character. + if grapheme == "\r\n" { + '\n' + } else { + grapheme + .chars() + .next() + .expect("graphemes must be non-empty") + } + }); + #[cfg(not(feature = "unicode-segmentation"))] + let res = text.chars(); + res +} diff --git a/crates/atuin-nucleo/matcher/src/chars/case_fold.rs b/crates/atuin-nucleo/matcher/src/chars/case_fold.rs new file mode 100644 index 00000000..aacbe461 --- /dev/null +++ b/crates/atuin-nucleo/matcher/src/chars/case_fold.rs @@ -0,0 +1,347 @@ +// DO NOT EDIT THIS FILE. IT WAS AUTOMATICALLY GENERATED BY: +// +// ucd-generate case-folding-simple /tmp/ucd-15.0.0 --chars +// +// Unicode version: 15.0.0. +// +// ucd-generate 0.3.0 is available on crates.io. + +pub const CASE_FOLDING_SIMPLE: &'static [(char, char)] = &[ + ('A', 'a'), ('B', 'b'), ('C', 'c'), ('D', 'd'), ('E', 'e'), ('F', 'f'), + ('G', 'g'), ('H', 'h'), ('I', 'i'), ('J', 'j'), ('K', 'k'), ('L', 'l'), + ('M', 'm'), ('N', 'n'), ('O', 'o'), ('P', 'p'), ('Q', 'q'), ('R', 'r'), + ('S', 's'), ('T', 't'), ('U', 'u'), ('V', 'v'), ('W', 'w'), ('X', 'x'), + ('Y', 'y'), ('Z', 'z'), ('µ', 'μ'), ('À', 'à'), ('Á', 'á'), + ('Â', 'â'), ('Ã', 'ã'), ('Ä', 'ä'), ('Å', 'å'), ('Æ', 'æ'), + ('Ç', 'ç'), ('È', 'è'), ('É', 'é'), ('Ê', 'ê'), ('Ë', 'ë'), + ('Ì', 'ì'), ('Í', 'í'), ('Î', 'î'), ('Ï', 'ï'), ('Ð', 'ð'), + ('Ñ', 'ñ'), ('Ò', 'ò'), ('Ó', 'ó'), ('Ô', 'ô'), ('Õ', 'õ'), + ('Ö', 'ö'), ('Ø', 'ø'), ('Ù', 'ù'), ('Ú', 'ú'), ('Û', 'û'), + ('Ü', 'ü'), ('Ý', 'ý'), ('Þ', 'þ'), ('Ā', 'ā'), ('Ă', 'ă'), + ('Ą', 'ą'), ('Ć', 'ć'), ('Ĉ', 'ĉ'), ('Ċ', 'ċ'), ('Č', 'č'), + ('Ď', 'ď'), ('Đ', 'đ'), ('Ē', 'ē'), ('Ĕ', 'ĕ'), ('Ė', 'ė'), + ('Ę', 'ę'), ('Ě', 'ě'), ('Ĝ', 'ĝ'), ('Ğ', 'ğ'), ('Ġ', 'ġ'), + ('Ģ', 'ģ'), ('Ĥ', 'ĥ'), ('Ħ', 'ħ'), ('Ĩ', 'ĩ'), ('Ī', 'ī'), + ('Ĭ', 'ĭ'), ('Į', 'į'), ('IJ', 'ij'), ('Ĵ', 'ĵ'), ('Ķ', 'ķ'), + ('Ĺ', 'ĺ'), ('Ļ', 'ļ'), ('Ľ', 'ľ'), ('Ŀ', 'ŀ'), ('Ł', 'ł'), + ('Ń', 'ń'), ('Ņ', 'ņ'), ('Ň', 'ň'), ('Ŋ', 'ŋ'), ('Ō', 'ō'), + ('Ŏ', 'ŏ'), ('Ő', 'ő'), ('Œ', 'œ'), ('Ŕ', 'ŕ'), ('Ŗ', 'ŗ'), + ('Ř', 'ř'), ('Ś', 'ś'), ('Ŝ', 'ŝ'), ('Ş', 'ş'), ('Š', 'š'), + ('Ţ', 'ţ'), ('Ť', 'ť'), ('Ŧ', 'ŧ'), ('Ũ', 'ũ'), ('Ū', 'ū'), + ('Ŭ', 'ŭ'), ('Ů', 'ů'), ('Ű', 'ű'), ('Ų', 'ų'), ('Ŵ', 'ŵ'), + ('Ŷ', 'ŷ'), ('Ÿ', 'ÿ'), ('Ź', 'ź'), ('Ż', 'ż'), ('Ž', 'ž'), + ('ſ', 's'), ('Ɓ', 'ɓ'), ('Ƃ', 'ƃ'), ('Ƅ', 'ƅ'), ('Ɔ', 'ɔ'), + ('Ƈ', 'ƈ'), ('Ɖ', 'ɖ'), ('Ɗ', 'ɗ'), ('Ƌ', 'ƌ'), ('Ǝ', 'ǝ'), + ('Ə', 'ə'), ('Ɛ', 'ɛ'), ('Ƒ', 'ƒ'), ('Ɠ', 'ɠ'), ('Ɣ', 'ɣ'), + ('Ɩ', 'ɩ'), ('Ɨ', 'ɨ'), ('Ƙ', 'ƙ'), ('Ɯ', 'ɯ'), ('Ɲ', 'ɲ'), + ('Ɵ', 'ɵ'), ('Ơ', 'ơ'), ('Ƣ', 'ƣ'), ('Ƥ', 'ƥ'), ('Ʀ', 'ʀ'), + ('Ƨ', 'ƨ'), ('Ʃ', 'ʃ'), ('Ƭ', 'ƭ'), ('Ʈ', 'ʈ'), ('Ư', 'ư'), + ('Ʊ', 'ʊ'), ('Ʋ', 'ʋ'), ('Ƴ', 'ƴ'), ('Ƶ', 'ƶ'), ('Ʒ', 'ʒ'), + ('Ƹ', 'ƹ'), ('Ƽ', 'ƽ'), ('DŽ', 'dž'), ('Dž', 'dž'), ('LJ', 'lj'), + ('Lj', 'lj'), ('NJ', 'nj'), ('Nj', 'nj'), ('Ǎ', 'ǎ'), ('Ǐ', 'ǐ'), + ('Ǒ', 'ǒ'), ('Ǔ', 'ǔ'), ('Ǖ', 'ǖ'), ('Ǘ', 'ǘ'), ('Ǚ', 'ǚ'), + ('Ǜ', 'ǜ'), ('Ǟ', 'ǟ'), ('Ǡ', 'ǡ'), ('Ǣ', 'ǣ'), ('Ǥ', 'ǥ'), + ('Ǧ', 'ǧ'), ('Ǩ', 'ǩ'), ('Ǫ', 'ǫ'), ('Ǭ', 'ǭ'), ('Ǯ', 'ǯ'), + ('DZ', 'dz'), ('Dz', 'dz'), ('Ǵ', 'ǵ'), ('Ƕ', 'ƕ'), ('Ƿ', 'ƿ'), + ('Ǹ', 'ǹ'), ('Ǻ', 'ǻ'), ('Ǽ', 'ǽ'), ('Ǿ', 'ǿ'), ('Ȁ', 'ȁ'), + ('Ȃ', 'ȃ'), ('Ȅ', 'ȅ'), ('Ȇ', 'ȇ'), ('Ȉ', 'ȉ'), ('Ȋ', 'ȋ'), + ('Ȍ', 'ȍ'), ('Ȏ', 'ȏ'), ('Ȑ', 'ȑ'), ('Ȓ', 'ȓ'), ('Ȕ', 'ȕ'), + ('Ȗ', 'ȗ'), ('Ș', 'ș'), ('Ț', 'ț'), ('Ȝ', 'ȝ'), ('Ȟ', 'ȟ'), + ('Ƞ', 'ƞ'), ('Ȣ', 'ȣ'), ('Ȥ', 'ȥ'), ('Ȧ', 'ȧ'), ('Ȩ', 'ȩ'), + ('Ȫ', 'ȫ'), ('Ȭ', 'ȭ'), ('Ȯ', 'ȯ'), ('Ȱ', 'ȱ'), ('Ȳ', 'ȳ'), + ('Ⱥ', 'ⱥ'), ('Ȼ', 'ȼ'), ('Ƚ', 'ƚ'), ('Ⱦ', 'ⱦ'), ('Ɂ', 'ɂ'), + ('Ƀ', 'ƀ'), ('Ʉ', 'ʉ'), ('Ʌ', 'ʌ'), ('Ɇ', 'ɇ'), ('Ɉ', 'ɉ'), + ('Ɋ', 'ɋ'), ('Ɍ', 'ɍ'), ('Ɏ', 'ɏ'), ('\u{345}', 'ι'), ('Ͱ', 'ͱ'), + ('Ͳ', 'ͳ'), ('Ͷ', 'ͷ'), ('Ϳ', 'ϳ'), ('Ά', 'ά'), ('Έ', 'έ'), + ('Ή', 'ή'), ('Ί', 'ί'), ('Ό', 'ό'), ('Ύ', 'ύ'), ('Ώ', 'ώ'), + ('Α', 'α'), ('Β', 'β'), ('Γ', 'γ'), ('Δ', 'δ'), ('Ε', 'ε'), + ('Ζ', 'ζ'), ('Η', 'η'), ('Θ', 'θ'), ('Ι', 'ι'), ('Κ', 'κ'), + ('Λ', 'λ'), ('Μ', 'μ'), ('Ν', 'ν'), ('Ξ', 'ξ'), ('Ο', 'ο'), + ('Π', 'π'), ('Ρ', 'ρ'), ('Σ', 'σ'), ('Τ', 'τ'), ('Υ', 'υ'), + ('Φ', 'φ'), ('Χ', 'χ'), ('Ψ', 'ψ'), ('Ω', 'ω'), ('Ϊ', 'ϊ'), + ('Ϋ', 'ϋ'), ('ς', 'σ'), ('Ϗ', 'ϗ'), ('ϐ', 'β'), ('ϑ', 'θ'), + ('ϕ', 'φ'), ('ϖ', 'π'), ('Ϙ', 'ϙ'), ('Ϛ', 'ϛ'), ('Ϝ', 'ϝ'), + ('Ϟ', 'ϟ'), ('Ϡ', 'ϡ'), ('Ϣ', 'ϣ'), ('Ϥ', 'ϥ'), ('Ϧ', 'ϧ'), + ('Ϩ', 'ϩ'), ('Ϫ', 'ϫ'), ('Ϭ', 'ϭ'), ('Ϯ', 'ϯ'), ('ϰ', 'κ'), + ('ϱ', 'ρ'), ('ϴ', 'θ'), ('ϵ', 'ε'), ('Ϸ', 'ϸ'), ('Ϲ', 'ϲ'), + ('Ϻ', 'ϻ'), ('Ͻ', 'ͻ'), ('Ͼ', 'ͼ'), ('Ͽ', 'ͽ'), ('Ѐ', 'ѐ'), + ('Ё', 'ё'), ('Ђ', 'ђ'), ('Ѓ', 'ѓ'), ('Є', 'є'), ('Ѕ', 'ѕ'), + ('І', 'і'), ('Ї', 'ї'), ('Ј', 'ј'), ('Љ', 'љ'), ('Њ', 'њ'), + ('Ћ', 'ћ'), ('Ќ', 'ќ'), ('Ѝ', 'ѝ'), ('Ў', 'ў'), ('Џ', 'џ'), + ('А', 'а'), ('Б', 'б'), ('В', 'в'), ('Г', 'г'), ('Д', 'д'), + ('Е', 'е'), ('Ж', 'ж'), ('З', 'з'), ('И', 'и'), ('Й', 'й'), + ('К', 'к'), ('Л', 'л'), ('М', 'м'), ('Н', 'н'), ('О', 'о'), + ('П', 'п'), ('Р', 'р'), ('С', 'с'), ('Т', 'т'), ('У', 'у'), + ('Ф', 'ф'), ('Х', 'х'), ('Ц', 'ц'), ('Ч', 'ч'), ('Ш', 'ш'), + ('Щ', 'щ'), ('Ъ', 'ъ'), ('Ы', 'ы'), ('Ь', 'ь'), ('Э', 'э'), + ('Ю', 'ю'), ('Я', 'я'), ('Ѡ', 'ѡ'), ('Ѣ', 'ѣ'), ('Ѥ', 'ѥ'), + ('Ѧ', 'ѧ'), ('Ѩ', 'ѩ'), ('Ѫ', 'ѫ'), ('Ѭ', 'ѭ'), ('Ѯ', 'ѯ'), + ('Ѱ', 'ѱ'), ('Ѳ', 'ѳ'), ('Ѵ', 'ѵ'), ('Ѷ', 'ѷ'), ('Ѹ', 'ѹ'), + ('Ѻ', 'ѻ'), ('Ѽ', 'ѽ'), ('Ѿ', 'ѿ'), ('Ҁ', 'ҁ'), ('Ҋ', 'ҋ'), + ('Ҍ', 'ҍ'), ('Ҏ', 'ҏ'), ('Ґ', 'ґ'), ('Ғ', 'ғ'), ('Ҕ', 'ҕ'), + ('Җ', 'җ'), ('Ҙ', 'ҙ'), ('Қ', 'қ'), ('Ҝ', 'ҝ'), ('Ҟ', 'ҟ'), + ('Ҡ', 'ҡ'), ('Ң', 'ң'), ('Ҥ', 'ҥ'), ('Ҧ', 'ҧ'), ('Ҩ', 'ҩ'), + ('Ҫ', 'ҫ'), ('Ҭ', 'ҭ'), ('Ү', 'ү'), ('Ұ', 'ұ'), ('Ҳ', 'ҳ'), + ('Ҵ', 'ҵ'), ('Ҷ', 'ҷ'), ('Ҹ', 'ҹ'), ('Һ', 'һ'), ('Ҽ', 'ҽ'), + ('Ҿ', 'ҿ'), ('Ӏ', 'ӏ'), ('Ӂ', 'ӂ'), ('Ӄ', 'ӄ'), ('Ӆ', 'ӆ'), + ('Ӈ', 'ӈ'), ('Ӊ', 'ӊ'), ('Ӌ', 'ӌ'), ('Ӎ', 'ӎ'), ('Ӑ', 'ӑ'), + ('Ӓ', 'ӓ'), ('Ӕ', 'ӕ'), ('Ӗ', 'ӗ'), ('Ә', 'ә'), ('Ӛ', 'ӛ'), + ('Ӝ', 'ӝ'), ('Ӟ', 'ӟ'), ('Ӡ', 'ӡ'), ('Ӣ', 'ӣ'), ('Ӥ', 'ӥ'), + ('Ӧ', 'ӧ'), ('Ө', 'ө'), ('Ӫ', 'ӫ'), ('Ӭ', 'ӭ'), ('Ӯ', 'ӯ'), + ('Ӱ', 'ӱ'), ('Ӳ', 'ӳ'), ('Ӵ', 'ӵ'), ('Ӷ', 'ӷ'), ('Ӹ', 'ӹ'), + ('Ӻ', 'ӻ'), ('Ӽ', 'ӽ'), ('Ӿ', 'ӿ'), ('Ԁ', 'ԁ'), ('Ԃ', 'ԃ'), + ('Ԅ', 'ԅ'), ('Ԇ', 'ԇ'), ('Ԉ', 'ԉ'), ('Ԋ', 'ԋ'), ('Ԍ', 'ԍ'), + ('Ԏ', 'ԏ'), ('Ԑ', 'ԑ'), ('Ԓ', 'ԓ'), ('Ԕ', 'ԕ'), ('Ԗ', 'ԗ'), + ('Ԙ', 'ԙ'), ('Ԛ', 'ԛ'), ('Ԝ', 'ԝ'), ('Ԟ', 'ԟ'), ('Ԡ', 'ԡ'), + ('Ԣ', 'ԣ'), ('Ԥ', 'ԥ'), ('Ԧ', 'ԧ'), ('Ԩ', 'ԩ'), ('Ԫ', 'ԫ'), + ('Ԭ', 'ԭ'), ('Ԯ', 'ԯ'), ('Ա', 'ա'), ('Բ', 'բ'), ('Գ', 'գ'), + ('Դ', 'դ'), ('Ե', 'ե'), ('Զ', 'զ'), ('Է', 'է'), ('Ը', 'ը'), + ('Թ', 'թ'), ('Ժ', 'ժ'), ('Ի', 'ի'), ('Լ', 'լ'), ('Խ', 'խ'), + ('Ծ', 'ծ'), ('Կ', 'կ'), ('Հ', 'հ'), ('Ձ', 'ձ'), ('Ղ', 'ղ'), + ('Ճ', 'ճ'), ('Մ', 'մ'), ('Յ', 'յ'), ('Ն', 'ն'), ('Շ', 'շ'), + ('Ո', 'ո'), ('Չ', 'չ'), ('Պ', 'պ'), ('Ջ', 'ջ'), ('Ռ', 'ռ'), + ('Ս', 'ս'), ('Վ', 'վ'), ('Տ', 'տ'), ('Ր', 'ր'), ('Ց', 'ց'), + ('Ւ', 'ւ'), ('Փ', 'փ'), ('Ք', 'ք'), ('Օ', 'օ'), ('Ֆ', 'ֆ'), + ('Ⴀ', 'ⴀ'), ('Ⴁ', 'ⴁ'), ('Ⴂ', 'ⴂ'), ('Ⴃ', 'ⴃ'), + ('Ⴄ', 'ⴄ'), ('Ⴅ', 'ⴅ'), ('Ⴆ', 'ⴆ'), ('Ⴇ', 'ⴇ'), + ('Ⴈ', 'ⴈ'), ('Ⴉ', 'ⴉ'), ('Ⴊ', 'ⴊ'), ('Ⴋ', 'ⴋ'), + ('Ⴌ', 'ⴌ'), ('Ⴍ', 'ⴍ'), ('Ⴎ', 'ⴎ'), ('Ⴏ', 'ⴏ'), + ('Ⴐ', 'ⴐ'), ('Ⴑ', 'ⴑ'), ('Ⴒ', 'ⴒ'), ('Ⴓ', 'ⴓ'), + ('Ⴔ', 'ⴔ'), ('Ⴕ', 'ⴕ'), ('Ⴖ', 'ⴖ'), ('Ⴗ', 'ⴗ'), + ('Ⴘ', 'ⴘ'), ('Ⴙ', 'ⴙ'), ('Ⴚ', 'ⴚ'), ('Ⴛ', 'ⴛ'), + ('Ⴜ', 'ⴜ'), ('Ⴝ', 'ⴝ'), ('Ⴞ', 'ⴞ'), ('Ⴟ', 'ⴟ'), + ('Ⴠ', 'ⴠ'), ('Ⴡ', 'ⴡ'), ('Ⴢ', 'ⴢ'), ('Ⴣ', 'ⴣ'), + ('Ⴤ', 'ⴤ'), ('Ⴥ', 'ⴥ'), ('Ⴧ', 'ⴧ'), ('Ⴭ', 'ⴭ'), + ('ᏸ', 'Ᏸ'), ('ᏹ', 'Ᏹ'), ('ᏺ', 'Ᏺ'), ('ᏻ', 'Ᏻ'), + ('ᏼ', 'Ᏼ'), ('ᏽ', 'Ᏽ'), ('ᲀ', 'в'), ('ᲁ', 'д'), ('ᲂ', 'о'), + ('ᲃ', 'с'), ('ᲄ', 'т'), ('ᲅ', 'т'), ('ᲆ', 'ъ'), ('ᲇ', 'ѣ'), + ('ᲈ', 'ꙋ'), ('Ა', 'ა'), ('Ბ', 'ბ'), ('Გ', 'გ'), + ('Დ', 'დ'), ('Ე', 'ე'), ('Ვ', 'ვ'), ('Ზ', 'ზ'), + ('Თ', 'თ'), ('Ი', 'ი'), ('Კ', 'კ'), ('Ლ', 'ლ'), + ('Მ', 'მ'), ('Ნ', 'ნ'), ('Ო', 'ო'), ('Პ', 'პ'), + ('Ჟ', 'ჟ'), ('Რ', 'რ'), ('Ს', 'ს'), ('Ტ', 'ტ'), + ('Უ', 'უ'), ('Ფ', 'ფ'), ('Ქ', 'ქ'), ('Ღ', 'ღ'), + ('Ყ', 'ყ'), ('Შ', 'შ'), ('Ჩ', 'ჩ'), ('Ც', 'ც'), + ('Ძ', 'ძ'), ('Წ', 'წ'), ('Ჭ', 'ჭ'), ('Ხ', 'ხ'), + ('Ჯ', 'ჯ'), ('Ჰ', 'ჰ'), ('Ჱ', 'ჱ'), ('Ჲ', 'ჲ'), + ('Ჳ', 'ჳ'), ('Ჴ', 'ჴ'), ('Ჵ', 'ჵ'), ('Ჶ', 'ჶ'), + ('Ჷ', 'ჷ'), ('Ჸ', 'ჸ'), ('Ჹ', 'ჹ'), ('Ჺ', 'ჺ'), + ('Ჽ', 'ჽ'), ('Ჾ', 'ჾ'), ('Ჿ', 'ჿ'), ('Ḁ', 'ḁ'), + ('Ḃ', 'ḃ'), ('Ḅ', 'ḅ'), ('Ḇ', 'ḇ'), ('Ḉ', 'ḉ'), + ('Ḋ', 'ḋ'), ('Ḍ', 'ḍ'), ('Ḏ', 'ḏ'), ('Ḑ', 'ḑ'), + ('Ḓ', 'ḓ'), ('Ḕ', 'ḕ'), ('Ḗ', 'ḗ'), ('Ḙ', 'ḙ'), + ('Ḛ', 'ḛ'), ('Ḝ', 'ḝ'), ('Ḟ', 'ḟ'), ('Ḡ', 'ḡ'), + ('Ḣ', 'ḣ'), ('Ḥ', 'ḥ'), ('Ḧ', 'ḧ'), ('Ḩ', 'ḩ'), + ('Ḫ', 'ḫ'), ('Ḭ', 'ḭ'), ('Ḯ', 'ḯ'), ('Ḱ', 'ḱ'), + ('Ḳ', 'ḳ'), ('Ḵ', 'ḵ'), ('Ḷ', 'ḷ'), ('Ḹ', 'ḹ'), + ('Ḻ', 'ḻ'), ('Ḽ', 'ḽ'), ('Ḿ', 'ḿ'), ('Ṁ', 'ṁ'), + ('Ṃ', 'ṃ'), ('Ṅ', 'ṅ'), ('Ṇ', 'ṇ'), ('Ṉ', 'ṉ'), + ('Ṋ', 'ṋ'), ('Ṍ', 'ṍ'), ('Ṏ', 'ṏ'), ('Ṑ', 'ṑ'), + ('Ṓ', 'ṓ'), ('Ṕ', 'ṕ'), ('Ṗ', 'ṗ'), ('Ṙ', 'ṙ'), + ('Ṛ', 'ṛ'), ('Ṝ', 'ṝ'), ('Ṟ', 'ṟ'), ('Ṡ', 'ṡ'), + ('Ṣ', 'ṣ'), ('Ṥ', 'ṥ'), ('Ṧ', 'ṧ'), ('Ṩ', 'ṩ'), + ('Ṫ', 'ṫ'), ('Ṭ', 'ṭ'), ('Ṯ', 'ṯ'), ('Ṱ', 'ṱ'), + ('Ṳ', 'ṳ'), ('Ṵ', 'ṵ'), ('Ṷ', 'ṷ'), ('Ṹ', 'ṹ'), + ('Ṻ', 'ṻ'), ('Ṽ', 'ṽ'), ('Ṿ', 'ṿ'), ('Ẁ', 'ẁ'), + ('Ẃ', 'ẃ'), ('Ẅ', 'ẅ'), ('Ẇ', 'ẇ'), ('Ẉ', 'ẉ'), + ('Ẋ', 'ẋ'), ('Ẍ', 'ẍ'), ('Ẏ', 'ẏ'), ('Ẑ', 'ẑ'), + ('Ẓ', 'ẓ'), ('Ẕ', 'ẕ'), ('ẛ', 'ṡ'), ('ẞ', 'ß'), + ('Ạ', 'ạ'), ('Ả', 'ả'), ('Ấ', 'ấ'), ('Ầ', 'ầ'), + ('Ẩ', 'ẩ'), ('Ẫ', 'ẫ'), ('Ậ', 'ậ'), ('Ắ', 'ắ'), + ('Ằ', 'ằ'), ('Ẳ', 'ẳ'), ('Ẵ', 'ẵ'), ('Ặ', 'ặ'), + ('Ẹ', 'ẹ'), ('Ẻ', 'ẻ'), ('Ẽ', 'ẽ'), ('Ế', 'ế'), + ('Ề', 'ề'), ('Ể', 'ể'), ('Ễ', 'ễ'), ('Ệ', 'ệ'), + ('Ỉ', 'ỉ'), ('Ị', 'ị'), ('Ọ', 'ọ'), ('Ỏ', 'ỏ'), + ('Ố', 'ố'), ('Ồ', 'ồ'), ('Ổ', 'ổ'), ('Ỗ', 'ỗ'), + ('Ộ', 'ộ'), ('Ớ', 'ớ'), ('Ờ', 'ờ'), ('Ở', 'ở'), + ('Ỡ', 'ỡ'), ('Ợ', 'ợ'), ('Ụ', 'ụ'), ('Ủ', 'ủ'), + ('Ứ', 'ứ'), ('Ừ', 'ừ'), ('Ử', 'ử'), ('Ữ', 'ữ'), + ('Ự', 'ự'), ('Ỳ', 'ỳ'), ('Ỵ', 'ỵ'), ('Ỷ', 'ỷ'), + ('Ỹ', 'ỹ'), ('Ỻ', 'ỻ'), ('Ỽ', 'ỽ'), ('Ỿ', 'ỿ'), + ('Ἀ', 'ἀ'), ('Ἁ', 'ἁ'), ('Ἂ', 'ἂ'), ('Ἃ', 'ἃ'), + ('Ἄ', 'ἄ'), ('Ἅ', 'ἅ'), ('Ἆ', 'ἆ'), ('Ἇ', 'ἇ'), + ('Ἐ', 'ἐ'), ('Ἑ', 'ἑ'), ('Ἒ', 'ἒ'), ('Ἓ', 'ἓ'), + ('Ἔ', 'ἔ'), ('Ἕ', 'ἕ'), ('Ἠ', 'ἠ'), ('Ἡ', 'ἡ'), + ('Ἢ', 'ἢ'), ('Ἣ', 'ἣ'), ('Ἤ', 'ἤ'), ('Ἥ', 'ἥ'), + ('Ἦ', 'ἦ'), ('Ἧ', 'ἧ'), ('Ἰ', 'ἰ'), ('Ἱ', 'ἱ'), + ('Ἲ', 'ἲ'), ('Ἳ', 'ἳ'), ('Ἴ', 'ἴ'), ('Ἵ', 'ἵ'), + ('Ἶ', 'ἶ'), ('Ἷ', 'ἷ'), ('Ὀ', 'ὀ'), ('Ὁ', 'ὁ'), + ('Ὂ', 'ὂ'), ('Ὃ', 'ὃ'), ('Ὄ', 'ὄ'), ('Ὅ', 'ὅ'), + ('Ὑ', 'ὑ'), ('Ὓ', 'ὓ'), ('Ὕ', 'ὕ'), ('Ὗ', 'ὗ'), + ('Ὠ', 'ὠ'), ('Ὡ', 'ὡ'), ('Ὢ', 'ὢ'), ('Ὣ', 'ὣ'), + ('Ὤ', 'ὤ'), ('Ὥ', 'ὥ'), ('Ὦ', 'ὦ'), ('Ὧ', 'ὧ'), + ('ᾈ', 'ᾀ'), ('ᾉ', 'ᾁ'), ('ᾊ', 'ᾂ'), ('ᾋ', 'ᾃ'), + ('ᾌ', 'ᾄ'), ('ᾍ', 'ᾅ'), ('ᾎ', 'ᾆ'), ('ᾏ', 'ᾇ'), + ('ᾘ', 'ᾐ'), ('ᾙ', 'ᾑ'), ('ᾚ', 'ᾒ'), ('ᾛ', 'ᾓ'), + ('ᾜ', 'ᾔ'), ('ᾝ', 'ᾕ'), ('ᾞ', 'ᾖ'), ('ᾟ', 'ᾗ'), + ('ᾨ', 'ᾠ'), ('ᾩ', 'ᾡ'), ('ᾪ', 'ᾢ'), ('ᾫ', 'ᾣ'), + ('ᾬ', 'ᾤ'), ('ᾭ', 'ᾥ'), ('ᾮ', 'ᾦ'), ('ᾯ', 'ᾧ'), + ('Ᾰ', 'ᾰ'), ('Ᾱ', 'ᾱ'), ('Ὰ', 'ὰ'), ('Ά', 'ά'), + ('ᾼ', 'ᾳ'), ('ι', 'ι'), ('Ὲ', 'ὲ'), ('Έ', 'έ'), + ('Ὴ', 'ὴ'), ('Ή', 'ή'), ('ῌ', 'ῃ'), ('Ῐ', 'ῐ'), + ('Ῑ', 'ῑ'), ('Ὶ', 'ὶ'), ('Ί', 'ί'), ('Ῠ', 'ῠ'), + ('Ῡ', 'ῡ'), ('Ὺ', 'ὺ'), ('Ύ', 'ύ'), ('Ῥ', 'ῥ'), + ('Ὸ', 'ὸ'), ('Ό', 'ό'), ('Ὼ', 'ὼ'), ('Ώ', 'ώ'), + ('ῼ', 'ῳ'), ('Ω', 'ω'), ('K', 'k'), ('Å', 'å'), ('Ⅎ', 'ⅎ'), + ('Ⅰ', 'ⅰ'), ('Ⅱ', 'ⅱ'), ('Ⅲ', 'ⅲ'), ('Ⅳ', 'ⅳ'), + ('Ⅴ', 'ⅴ'), ('Ⅵ', 'ⅵ'), ('Ⅶ', 'ⅶ'), ('Ⅷ', 'ⅷ'), + ('Ⅸ', 'ⅸ'), ('Ⅹ', 'ⅹ'), ('Ⅺ', 'ⅺ'), ('Ⅻ', 'ⅻ'), + ('Ⅼ', 'ⅼ'), ('Ⅽ', 'ⅽ'), ('Ⅾ', 'ⅾ'), ('Ⅿ', 'ⅿ'), + ('Ↄ', 'ↄ'), ('Ⓐ', 'ⓐ'), ('Ⓑ', 'ⓑ'), ('Ⓒ', 'ⓒ'), + ('Ⓓ', 'ⓓ'), ('Ⓔ', 'ⓔ'), ('Ⓕ', 'ⓕ'), ('Ⓖ', 'ⓖ'), + ('Ⓗ', 'ⓗ'), ('Ⓘ', 'ⓘ'), ('Ⓙ', 'ⓙ'), ('Ⓚ', 'ⓚ'), + ('Ⓛ', 'ⓛ'), ('Ⓜ', 'ⓜ'), ('Ⓝ', 'ⓝ'), ('Ⓞ', 'ⓞ'), + ('Ⓟ', 'ⓟ'), ('Ⓠ', 'ⓠ'), ('Ⓡ', 'ⓡ'), ('Ⓢ', 'ⓢ'), + ('Ⓣ', 'ⓣ'), ('Ⓤ', 'ⓤ'), ('Ⓥ', 'ⓥ'), ('Ⓦ', 'ⓦ'), + ('Ⓧ', 'ⓧ'), ('Ⓨ', 'ⓨ'), ('Ⓩ', 'ⓩ'), ('Ⰰ', 'ⰰ'), + ('Ⰱ', 'ⰱ'), ('Ⰲ', 'ⰲ'), ('Ⰳ', 'ⰳ'), ('Ⰴ', 'ⰴ'), + ('Ⰵ', 'ⰵ'), ('Ⰶ', 'ⰶ'), ('Ⰷ', 'ⰷ'), ('Ⰸ', 'ⰸ'), + ('Ⰹ', 'ⰹ'), ('Ⰺ', 'ⰺ'), ('Ⰻ', 'ⰻ'), ('Ⰼ', 'ⰼ'), + ('Ⰽ', 'ⰽ'), ('Ⰾ', 'ⰾ'), ('Ⰿ', 'ⰿ'), ('Ⱀ', 'ⱀ'), + ('Ⱁ', 'ⱁ'), ('Ⱂ', 'ⱂ'), ('Ⱃ', 'ⱃ'), ('Ⱄ', 'ⱄ'), + ('Ⱅ', 'ⱅ'), ('Ⱆ', 'ⱆ'), ('Ⱇ', 'ⱇ'), ('Ⱈ', 'ⱈ'), + ('Ⱉ', 'ⱉ'), ('Ⱊ', 'ⱊ'), ('Ⱋ', 'ⱋ'), ('Ⱌ', 'ⱌ'), + ('Ⱍ', 'ⱍ'), ('Ⱎ', 'ⱎ'), ('Ⱏ', 'ⱏ'), ('Ⱐ', 'ⱐ'), + ('Ⱑ', 'ⱑ'), ('Ⱒ', 'ⱒ'), ('Ⱓ', 'ⱓ'), ('Ⱔ', 'ⱔ'), + ('Ⱕ', 'ⱕ'), ('Ⱖ', 'ⱖ'), ('Ⱗ', 'ⱗ'), ('Ⱘ', 'ⱘ'), + ('Ⱙ', 'ⱙ'), ('Ⱚ', 'ⱚ'), ('Ⱛ', 'ⱛ'), ('Ⱜ', 'ⱜ'), + ('Ⱝ', 'ⱝ'), ('Ⱞ', 'ⱞ'), ('Ⱟ', 'ⱟ'), ('Ⱡ', 'ⱡ'), + ('Ɫ', 'ɫ'), ('Ᵽ', 'ᵽ'), ('Ɽ', 'ɽ'), ('Ⱨ', 'ⱨ'), + ('Ⱪ', 'ⱪ'), ('Ⱬ', 'ⱬ'), ('Ɑ', 'ɑ'), ('Ɱ', 'ɱ'), ('Ɐ', 'ɐ'), + ('Ɒ', 'ɒ'), ('Ⱳ', 'ⱳ'), ('Ⱶ', 'ⱶ'), ('Ȿ', 'ȿ'), ('Ɀ', 'ɀ'), + ('Ⲁ', 'ⲁ'), ('Ⲃ', 'ⲃ'), ('Ⲅ', 'ⲅ'), ('Ⲇ', 'ⲇ'), + ('Ⲉ', 'ⲉ'), ('Ⲋ', 'ⲋ'), ('Ⲍ', 'ⲍ'), ('Ⲏ', 'ⲏ'), + ('Ⲑ', 'ⲑ'), ('Ⲓ', 'ⲓ'), ('Ⲕ', 'ⲕ'), ('Ⲗ', 'ⲗ'), + ('Ⲙ', 'ⲙ'), ('Ⲛ', 'ⲛ'), ('Ⲝ', 'ⲝ'), ('Ⲟ', 'ⲟ'), + ('Ⲡ', 'ⲡ'), ('Ⲣ', 'ⲣ'), ('Ⲥ', 'ⲥ'), ('Ⲧ', 'ⲧ'), + ('Ⲩ', 'ⲩ'), ('Ⲫ', 'ⲫ'), ('Ⲭ', 'ⲭ'), ('Ⲯ', 'ⲯ'), + ('Ⲱ', 'ⲱ'), ('Ⲳ', 'ⲳ'), ('Ⲵ', 'ⲵ'), ('Ⲷ', 'ⲷ'), + ('Ⲹ', 'ⲹ'), ('Ⲻ', 'ⲻ'), ('Ⲽ', 'ⲽ'), ('Ⲿ', 'ⲿ'), + ('Ⳁ', 'ⳁ'), ('Ⳃ', 'ⳃ'), ('Ⳅ', 'ⳅ'), ('Ⳇ', 'ⳇ'), + ('Ⳉ', 'ⳉ'), ('Ⳋ', 'ⳋ'), ('Ⳍ', 'ⳍ'), ('Ⳏ', 'ⳏ'), + ('Ⳑ', 'ⳑ'), ('Ⳓ', 'ⳓ'), ('Ⳕ', 'ⳕ'), ('Ⳗ', 'ⳗ'), + ('Ⳙ', 'ⳙ'), ('Ⳛ', 'ⳛ'), ('Ⳝ', 'ⳝ'), ('Ⳟ', 'ⳟ'), + ('Ⳡ', 'ⳡ'), ('Ⳣ', 'ⳣ'), ('Ⳬ', 'ⳬ'), ('Ⳮ', 'ⳮ'), + ('Ⳳ', 'ⳳ'), ('Ꙁ', 'ꙁ'), ('Ꙃ', 'ꙃ'), ('Ꙅ', 'ꙅ'), + ('Ꙇ', 'ꙇ'), ('Ꙉ', 'ꙉ'), ('Ꙋ', 'ꙋ'), ('Ꙍ', 'ꙍ'), + ('Ꙏ', 'ꙏ'), ('Ꙑ', 'ꙑ'), ('Ꙓ', 'ꙓ'), ('Ꙕ', 'ꙕ'), + ('Ꙗ', 'ꙗ'), ('Ꙙ', 'ꙙ'), ('Ꙛ', 'ꙛ'), ('Ꙝ', 'ꙝ'), + ('Ꙟ', 'ꙟ'), ('Ꙡ', 'ꙡ'), ('Ꙣ', 'ꙣ'), ('Ꙥ', 'ꙥ'), + ('Ꙧ', 'ꙧ'), ('Ꙩ', 'ꙩ'), ('Ꙫ', 'ꙫ'), ('Ꙭ', 'ꙭ'), + ('Ꚁ', 'ꚁ'), ('Ꚃ', 'ꚃ'), ('Ꚅ', 'ꚅ'), ('Ꚇ', 'ꚇ'), + ('Ꚉ', 'ꚉ'), ('Ꚋ', 'ꚋ'), ('Ꚍ', 'ꚍ'), ('Ꚏ', 'ꚏ'), + ('Ꚑ', 'ꚑ'), ('Ꚓ', 'ꚓ'), ('Ꚕ', 'ꚕ'), ('Ꚗ', 'ꚗ'), + ('Ꚙ', 'ꚙ'), ('Ꚛ', 'ꚛ'), ('Ꜣ', 'ꜣ'), ('Ꜥ', 'ꜥ'), + ('Ꜧ', 'ꜧ'), ('Ꜩ', 'ꜩ'), ('Ꜫ', 'ꜫ'), ('Ꜭ', 'ꜭ'), + ('Ꜯ', 'ꜯ'), ('Ꜳ', 'ꜳ'), ('Ꜵ', 'ꜵ'), ('Ꜷ', 'ꜷ'), + ('Ꜹ', 'ꜹ'), ('Ꜻ', 'ꜻ'), ('Ꜽ', 'ꜽ'), ('Ꜿ', 'ꜿ'), + ('Ꝁ', 'ꝁ'), ('Ꝃ', 'ꝃ'), ('Ꝅ', 'ꝅ'), ('Ꝇ', 'ꝇ'), + ('Ꝉ', 'ꝉ'), ('Ꝋ', 'ꝋ'), ('Ꝍ', 'ꝍ'), ('Ꝏ', 'ꝏ'), + ('Ꝑ', 'ꝑ'), ('Ꝓ', 'ꝓ'), ('Ꝕ', 'ꝕ'), ('Ꝗ', 'ꝗ'), + ('Ꝙ', 'ꝙ'), ('Ꝛ', 'ꝛ'), ('Ꝝ', 'ꝝ'), ('Ꝟ', 'ꝟ'), + ('Ꝡ', 'ꝡ'), ('Ꝣ', 'ꝣ'), ('Ꝥ', 'ꝥ'), ('Ꝧ', 'ꝧ'), + ('Ꝩ', 'ꝩ'), ('Ꝫ', 'ꝫ'), ('Ꝭ', 'ꝭ'), ('Ꝯ', 'ꝯ'), + ('Ꝺ', 'ꝺ'), ('Ꝼ', 'ꝼ'), ('Ᵹ', 'ᵹ'), ('Ꝿ', 'ꝿ'), + ('Ꞁ', 'ꞁ'), ('Ꞃ', 'ꞃ'), ('Ꞅ', 'ꞅ'), ('Ꞇ', 'ꞇ'), + ('Ꞌ', 'ꞌ'), ('Ɥ', 'ɥ'), ('Ꞑ', 'ꞑ'), ('Ꞓ', 'ꞓ'), + ('Ꞗ', 'ꞗ'), ('Ꞙ', 'ꞙ'), ('Ꞛ', 'ꞛ'), ('Ꞝ', 'ꞝ'), + ('Ꞟ', 'ꞟ'), ('Ꞡ', 'ꞡ'), ('Ꞣ', 'ꞣ'), ('Ꞥ', 'ꞥ'), + ('Ꞧ', 'ꞧ'), ('Ꞩ', 'ꞩ'), ('Ɦ', 'ɦ'), ('Ɜ', 'ɜ'), ('Ɡ', 'ɡ'), + ('Ɬ', 'ɬ'), ('Ɪ', 'ɪ'), ('Ʞ', 'ʞ'), ('Ʇ', 'ʇ'), ('Ʝ', 'ʝ'), + ('Ꭓ', 'ꭓ'), ('Ꞵ', 'ꞵ'), ('Ꞷ', 'ꞷ'), ('Ꞹ', 'ꞹ'), + ('Ꞻ', 'ꞻ'), ('Ꞽ', 'ꞽ'), ('Ꞿ', 'ꞿ'), ('Ꟁ', 'ꟁ'), + ('Ꟃ', 'ꟃ'), ('Ꞔ', 'ꞔ'), ('Ʂ', 'ʂ'), ('Ᶎ', 'ᶎ'), + ('Ꟈ', 'ꟈ'), ('Ꟊ', 'ꟊ'), ('Ꟑ', 'ꟑ'), ('Ꟗ', 'ꟗ'), + ('Ꟙ', 'ꟙ'), ('Ꟶ', 'ꟶ'), ('ꭰ', 'Ꭰ'), ('ꭱ', 'Ꭱ'), + ('ꭲ', 'Ꭲ'), ('ꭳ', 'Ꭳ'), ('ꭴ', 'Ꭴ'), ('ꭵ', 'Ꭵ'), + ('ꭶ', 'Ꭶ'), ('ꭷ', 'Ꭷ'), ('ꭸ', 'Ꭸ'), ('ꭹ', 'Ꭹ'), + ('ꭺ', 'Ꭺ'), ('ꭻ', 'Ꭻ'), ('ꭼ', 'Ꭼ'), ('ꭽ', 'Ꭽ'), + ('ꭾ', 'Ꭾ'), ('ꭿ', 'Ꭿ'), ('ꮀ', 'Ꮀ'), ('ꮁ', 'Ꮁ'), + ('ꮂ', 'Ꮂ'), ('ꮃ', 'Ꮃ'), ('ꮄ', 'Ꮄ'), ('ꮅ', 'Ꮅ'), + ('ꮆ', 'Ꮆ'), ('ꮇ', 'Ꮇ'), ('ꮈ', 'Ꮈ'), ('ꮉ', 'Ꮉ'), + ('ꮊ', 'Ꮊ'), ('ꮋ', 'Ꮋ'), ('ꮌ', 'Ꮌ'), ('ꮍ', 'Ꮍ'), + ('ꮎ', 'Ꮎ'), ('ꮏ', 'Ꮏ'), ('ꮐ', 'Ꮐ'), ('ꮑ', 'Ꮑ'), + ('ꮒ', 'Ꮒ'), ('ꮓ', 'Ꮓ'), ('ꮔ', 'Ꮔ'), ('ꮕ', 'Ꮕ'), + ('ꮖ', 'Ꮖ'), ('ꮗ', 'Ꮗ'), ('ꮘ', 'Ꮘ'), ('ꮙ', 'Ꮙ'), + ('ꮚ', 'Ꮚ'), ('ꮛ', 'Ꮛ'), ('ꮜ', 'Ꮜ'), ('ꮝ', 'Ꮝ'), + ('ꮞ', 'Ꮞ'), ('ꮟ', 'Ꮟ'), ('ꮠ', 'Ꮠ'), ('ꮡ', 'Ꮡ'), + ('ꮢ', 'Ꮢ'), ('ꮣ', 'Ꮣ'), ('ꮤ', 'Ꮤ'), ('ꮥ', 'Ꮥ'), + ('ꮦ', 'Ꮦ'), ('ꮧ', 'Ꮧ'), ('ꮨ', 'Ꮨ'), ('ꮩ', 'Ꮩ'), + ('ꮪ', 'Ꮪ'), ('ꮫ', 'Ꮫ'), ('ꮬ', 'Ꮬ'), ('ꮭ', 'Ꮭ'), + ('ꮮ', 'Ꮮ'), ('ꮯ', 'Ꮯ'), ('ꮰ', 'Ꮰ'), ('ꮱ', 'Ꮱ'), + ('ꮲ', 'Ꮲ'), ('ꮳ', 'Ꮳ'), ('ꮴ', 'Ꮴ'), ('ꮵ', 'Ꮵ'), + ('ꮶ', 'Ꮶ'), ('ꮷ', 'Ꮷ'), ('ꮸ', 'Ꮸ'), ('ꮹ', 'Ꮹ'), + ('ꮺ', 'Ꮺ'), ('ꮻ', 'Ꮻ'), ('ꮼ', 'Ꮼ'), ('ꮽ', 'Ꮽ'), + ('ꮾ', 'Ꮾ'), ('ꮿ', 'Ꮿ'), ('Ａ', 'ａ'), ('Ｂ', 'ｂ'), + ('Ｃ', 'ｃ'), ('Ｄ', 'ｄ'), ('Ｅ', 'ｅ'), ('Ｆ', 'ｆ'), + ('Ｇ', 'ｇ'), ('Ｈ', 'ｈ'), ('Ｉ', 'ｉ'), ('Ｊ', 'ｊ'), + ('Ｋ', 'ｋ'), ('Ｌ', 'ｌ'), ('Ｍ', 'ｍ'), ('Ｎ', 'ｎ'), + ('Ｏ', 'ｏ'), ('Ｐ', 'ｐ'), ('Ｑ', 'ｑ'), ('Ｒ', 'ｒ'), + ('Ｓ', 'ｓ'), ('Ｔ', 'ｔ'), ('Ｕ', 'ｕ'), ('Ｖ', 'ｖ'), + ('Ｗ', 'ｗ'), ('Ｘ', 'ｘ'), ('Ｙ', 'ｙ'), ('Ｚ', 'ｚ'), + ('𐐀', '𐐨'), ('𐐁', '𐐩'), ('𐐂', '𐐪'), ('𐐃', '𐐫'), + ('𐐄', '𐐬'), ('𐐅', '𐐭'), ('𐐆', '𐐮'), ('𐐇', '𐐯'), + ('𐐈', '𐐰'), ('𐐉', '𐐱'), ('𐐊', '𐐲'), ('𐐋', '𐐳'), + ('𐐌', '𐐴'), ('𐐍', '𐐵'), ('𐐎', '𐐶'), ('𐐏', '𐐷'), + ('𐐐', '𐐸'), ('𐐑', '𐐹'), ('𐐒', '𐐺'), ('𐐓', '𐐻'), + ('𐐔', '𐐼'), ('𐐕', '𐐽'), ('𐐖', '𐐾'), ('𐐗', '𐐿'), + ('𐐘', '𐑀'), ('𐐙', '𐑁'), ('𐐚', '𐑂'), ('𐐛', '𐑃'), + ('𐐜', '𐑄'), ('𐐝', '𐑅'), ('𐐞', '𐑆'), ('𐐟', '𐑇'), + ('𐐠', '𐑈'), ('𐐡', '𐑉'), ('𐐢', '𐑊'), ('𐐣', '𐑋'), + ('𐐤', '𐑌'), ('𐐥', '𐑍'), ('𐐦', '𐑎'), ('𐐧', '𐑏'), + ('𐒰', '𐓘'), ('𐒱', '𐓙'), ('𐒲', '𐓚'), ('𐒳', '𐓛'), + ('𐒴', '𐓜'), ('𐒵', '𐓝'), ('𐒶', '𐓞'), ('𐒷', '𐓟'), + ('𐒸', '𐓠'), ('𐒹', '𐓡'), ('𐒺', '𐓢'), ('𐒻', '𐓣'), + ('𐒼', '𐓤'), ('𐒽', '𐓥'), ('𐒾', '𐓦'), ('𐒿', '𐓧'), + ('𐓀', '𐓨'), ('𐓁', '𐓩'), ('𐓂', '𐓪'), ('𐓃', '𐓫'), + ('𐓄', '𐓬'), ('𐓅', '𐓭'), ('𐓆', '𐓮'), ('𐓇', '𐓯'), + ('𐓈', '𐓰'), ('𐓉', '𐓱'), ('𐓊', '𐓲'), ('𐓋', '𐓳'), + ('𐓌', '𐓴'), ('𐓍', '𐓵'), ('𐓎', '𐓶'), ('𐓏', '𐓷'), + ('𐓐', '𐓸'), ('𐓑', '𐓹'), ('𐓒', '𐓺'), ('𐓓', '𐓻'), + ('𐕰', '𐖗'), ('𐕱', '𐖘'), ('𐕲', '𐖙'), ('𐕳', '𐖚'), + ('𐕴', '𐖛'), ('𐕵', '𐖜'), ('𐕶', '𐖝'), ('𐕷', '𐖞'), + ('𐕸', '𐖟'), ('𐕹', '𐖠'), ('𐕺', '𐖡'), ('𐕼', '𐖣'), + ('𐕽', '𐖤'), ('𐕾', '𐖥'), ('𐕿', '𐖦'), ('𐖀', '𐖧'), + ('𐖁', '𐖨'), ('𐖂', '𐖩'), ('𐖃', '𐖪'), ('𐖄', '𐖫'), + ('𐖅', '𐖬'), ('𐖆', '𐖭'), ('𐖇', '𐖮'), ('𐖈', '𐖯'), + ('𐖉', '𐖰'), ('𐖊', '𐖱'), ('𐖌', '𐖳'), ('𐖍', '𐖴'), + ('𐖎', '𐖵'), ('𐖏', '𐖶'), ('𐖐', '𐖷'), ('𐖑', '𐖸'), + ('𐖒', '𐖹'), ('𐖔', '𐖻'), ('𐖕', '𐖼'), ('𐲀', '𐳀'), + ('𐲁', '𐳁'), ('𐲂', '𐳂'), ('𐲃', '𐳃'), ('𐲄', '𐳄'), + ('𐲅', '𐳅'), ('𐲆', '𐳆'), ('𐲇', '𐳇'), ('𐲈', '𐳈'), + ('𐲉', '𐳉'), ('𐲊', '𐳊'), ('𐲋', '𐳋'), ('𐲌', '𐳌'), + ('𐲍', '𐳍'), ('𐲎', '𐳎'), ('𐲏', '𐳏'), ('𐲐', '𐳐'), + ('𐲑', '𐳑'), ('𐲒', '𐳒'), ('𐲓', '𐳓'), ('𐲔', '𐳔'), + ('𐲕', '𐳕'), ('𐲖', '𐳖'), ('𐲗', '𐳗'), ('𐲘', '𐳘'), + ('𐲙', '𐳙'), ('𐲚', '𐳚'), ('𐲛', '𐳛'), ('𐲜', '𐳜'), + ('𐲝', '𐳝'), ('𐲞', '𐳞'), ('𐲟', '𐳟'), ('𐲠', '𐳠'), + ('𐲡', '𐳡'), ('𐲢', '𐳢'), ('𐲣', '𐳣'), ('𐲤', '𐳤'), + ('𐲥', '𐳥'), ('𐲦', '𐳦'), ('𐲧', '𐳧'), ('𐲨', '𐳨'), + ('𐲩', '𐳩'), ('𐲪', '𐳪'), ('𐲫', '𐳫'), ('𐲬', '𐳬'), + ('𐲭', '𐳭'), ('𐲮', '𐳮'), ('𐲯', '𐳯'), ('𐲰', '𐳰'), + ('𐲱', '𐳱'), ('𐲲', '𐳲'), ('𑢠', '𑣀'), ('𑢡', '𑣁'), + ('𑢢', '𑣂'), ('𑢣', '𑣃'), ('𑢤', '𑣄'), ('𑢥', '𑣅'), + ('𑢦', '𑣆'), ('𑢧', '𑣇'), ('𑢨', '𑣈'), ('𑢩', '𑣉'), + ('𑢪', '𑣊'), ('𑢫', '𑣋'), ('𑢬', '𑣌'), ('𑢭', '𑣍'), + ('𑢮', '𑣎'), ('𑢯', '𑣏'), ('𑢰', '𑣐'), ('𑢱', '𑣑'), + ('𑢲', '𑣒'), ('𑢳', '𑣓'), ('𑢴', '𑣔'), ('𑢵', '𑣕'), + ('𑢶', '𑣖'), ('𑢷', '𑣗'), ('𑢸', '𑣘'), ('𑢹', '𑣙'), + ('𑢺', '𑣚'), ('𑢻', '𑣛'), ('𑢼', '𑣜'), ('𑢽', '𑣝'), + ('𑢾', '𑣞'), ('𑢿', '𑣟'), ('𖹀', '𖹠'), ('𖹁', '𖹡'), + ('𖹂', '𖹢'), ('𖹃', '𖹣'), ('𖹄', '𖹤'), ('𖹅', '𖹥'), + ('𖹆', '𖹦'), ('𖹇', '𖹧'), ('𖹈', '𖹨'), ('𖹉', '𖹩'), + ('𖹊', '𖹪'), ('𖹋', '𖹫'), ('𖹌', '𖹬'), ('𖹍', '𖹭'), + ('𖹎', '𖹮'), ('𖹏', '𖹯'), ('𖹐', '𖹰'), ('𖹑', '𖹱'), + ('𖹒', '𖹲'), ('𖹓', '𖹳'), ('𖹔', '𖹴'), ('𖹕', '𖹵'), + ('𖹖', '𖹶'), ('𖹗', '𖹷'), ('𖹘', '𖹸'), ('𖹙', '𖹹'), + ('𖹚', '𖹺'), ('𖹛', '𖹻'), ('𖹜', '𖹼'), ('𖹝', '𖹽'), + ('𖹞', '𖹾'), ('𖹟', '𖹿'), ('𞤀', '𞤢'), ('𞤁', '𞤣'), + ('𞤂', '𞤤'), ('𞤃', '𞤥'), ('𞤄', '𞤦'), ('𞤅', '𞤧'), + ('𞤆', '𞤨'), ('𞤇', '𞤩'), ('𞤈', '𞤪'), ('𞤉', '𞤫'), + ('𞤊', '𞤬'), ('𞤋', '𞤭'), ('𞤌', '𞤮'), ('𞤍', '𞤯'), + ('𞤎', '𞤰'), ('𞤏', '𞤱'), ('𞤐', '𞤲'), ('𞤑', '𞤳'), + ('𞤒', '𞤴'), ('𞤓', '𞤵'), ('𞤔', '𞤶'), ('𞤕', '𞤷'), + ('𞤖', '𞤸'), ('𞤗', '𞤹'), ('𞤘', '𞤺'), ('𞤙', '𞤻'), + ('𞤚', '𞤼'), ('𞤛', '𞤽'), ('𞤜', '𞤾'), ('𞤝', '𞤿'), + ('𞤞', '𞥀'), ('𞤟', '𞥁'), ('𞤠', '𞥂'), ('𞤡', '𞥃'), +]; diff --git a/crates/atuin-nucleo/matcher/src/chars/normalize.rs b/crates/atuin-nucleo/matcher/src/chars/normalize.rs new file mode 100644 index 00000000..3de501aa --- /dev/null +++ b/crates/atuin-nucleo/matcher/src/chars/normalize.rs @@ -0,0 +1,972 @@ +/// Normalize a Unicode character by converting Latin characters which are variants +/// of ASCII characters to their latin equivalent. +/// +/// Note that this method acts on single `char`s: if you want to perform full normalization, you +/// should first split on graphemes, and then normalize each grapheme by normalizing the first +/// `char` in the grapheme. +/// +/// If a character does not normalize to a single ASCII character, no normalization is performed. +/// +/// This performs normalization within the following Unicode blocks: +/// +/// - [Latin-1 Supplement](https://en.wikipedia.org/wiki/Latin-1_Supplement) +/// - [Latin Extended-A](https://en.wikipedia.org/wiki/Latin_Extended-A) +/// - [Latin Extended-B](https://en.wikipedia.org/wiki/Latin_Extended-B) +/// - [Latin Extended Additional](https://en.wikipedia.org/wiki/Latin_Extended_Additional) +/// - [Superscripts and Subscripts](https://en.wikipedia.org/wiki/Superscripts_and_Subscripts) +/// +/// If the character does not fall in this block, it is not normalized. +/// +/// # Example +/// ``` +/// # use nucleo_matcher::chars::normalize; +/// assert_eq!(normalize('ä'), 'a'); +/// assert_eq!(normalize('Æ'), 'Æ'); +/// assert_eq!(normalize('ữ'), 'u'); +/// ``` +pub fn normalize(c: char) -> char { + // outside checked blocks + if c < '\u{a0}' || c >= '\u{20A0}' { + return c; + } + // Latin-1 Supplement, Extended-A, Extended-B + if c <= '\u{29f}' { + return LATIN_1AB[c as usize - '\u{a0}' as usize]; + } + // between blocks + if c < '\u{1e00}' { + return c; + } + // Latin Extended Additional + if c <= '\u{1eff}' { + return LATIN_EXTENDED_ADDITIONAL[c as usize - '\u{1e00}' as usize]; + } + // between blocks + if c < '\u{2070}' { + return c; + } + // Superscripts and subscripts + SUPERSCRIPTS_AND_SUBSCRIPTS[c as usize - '\u{2070}' as usize] +} + +/// A char array corresponding to the following contiguous Unicode blocks: +/// +/// - [Latin-1 Supplement](https://en.wikipedia.org/wiki/Latin-1_Supplement) +/// - [Latin Extended-A](https://en.wikipedia.org/wiki/Latin_Extended-A) +/// - [Latin Extended-B](https://en.wikipedia.org/wiki/Latin_Extended-B) +/// +/// This covers the range `'\u{a0}'..='\u{29f}'`. +static LATIN_1AB: [char; 512] = [ + '\u{a0}', // invisible NON BREAKING SPACE + '!', // '¡'; '\u{a1}' + '¢', // '¢'; '\u{a2}' + '£', // '£'; '\u{a3}' + '¤', // '¤'; '\u{a4}' + '¥', // '¥'; '\u{a5}' + '¦', // '¦'; '\u{a6}' + '§', // '§'; '\u{a7}' + '¨', // '¨'; '\u{a8}' + '©', // '©'; '\u{a9}' + 'a', // 'ª'; '\u{aa}' + '«', // '«'; '\u{ab}' + '¬', // '¬'; '\u{ac}' + '\u{ad}', // invisible SOFT HYPHEN + '®', // '®'; '\u{ae}' + '¯', // '¯'; '\u{af}' + '°', // '°'; '\u{b0}' + '±', // '±'; '\u{b1}' + '2', // '²'; '\u{b2}' + '3', // '³'; '\u{b3}' + '´', // '´'; '\u{b4}' + 'µ', // 'µ'; '\u{b5}' + '¶', // '¶'; '\u{b6}' + '·', // '·'; '\u{b7}' + '¸', // '¸'; '\u{b8}' + '1', // '¹'; '\u{b9}' + '0', // 'º'; '\u{ba}' + '»', // '»'; '\u{bb}' + '¼', // '¼'; '\u{bc}' + '½', // '½'; '\u{bd}' + '¾', // '¾'; '\u{be}' + '?', // '¿'; '\u{bf}' + 'A', // 'À'; '\u{c0}' + 'A', // 'Á'; '\u{c1}' + 'A', // 'Â'; '\u{c2}' + 'A', // 'Ã'; '\u{c3}' + 'A', // 'Ä'; '\u{c4}' + 'A', // 'Å'; '\u{c5}' + 'Æ', // 'Æ'; '\u{c6}' + 'C', // 'Ç'; '\u{c7}' + 'E', // 'È'; '\u{c8}' + 'E', // 'É'; '\u{c9}' + 'E', // 'Ê'; '\u{ca}' + 'E', // 'Ë'; '\u{cb}' + 'I', // 'Ì'; '\u{cc}' + 'I', // 'Í'; '\u{cd}' + 'I', // 'Î'; '\u{ce}' + 'I', // 'Ï'; '\u{cf}' + 'D', // 'Ð'; '\u{d0}' + 'N', // 'Ñ'; '\u{d1}' + 'O', // 'Ò'; '\u{d2}' + 'O', // 'Ó'; '\u{d3}' + 'O', // 'Ô'; '\u{d4}' + 'O', // 'Õ'; '\u{d5}' + 'O', // 'Ö'; '\u{d6}' + '×', // '×'; '\u{d7}' + 'O', // 'Ø'; '\u{d8}' + 'U', // 'Ù'; '\u{d9}' + 'U', // 'Ú'; '\u{da}' + 'U', // 'Û'; '\u{db}' + 'U', // 'Ü'; '\u{dc}' + 'Y', // 'Ý'; '\u{dd}' + 'Þ', // 'Þ'; '\u{de}' + 's', // 'ß'; '\u{df}' + 'a', // 'à'; '\u{e0}' + 'a', // 'á'; '\u{e1}' + 'a', // 'â'; '\u{e2}' + 'a', // 'ã'; '\u{e3}' + 'a', // 'ä'; '\u{e4}' + 'a', // 'å'; '\u{e5}' + 'æ', // 'æ'; '\u{e6}' + 'c', // 'ç'; '\u{e7}' + 'e', // 'è'; '\u{e8}' + 'e', // 'é'; '\u{e9}' + 'e', // 'ê'; '\u{ea}' + 'e', // 'ë'; '\u{eb}' + 'i', // 'ì'; '\u{ec}' + 'i', // 'í'; '\u{ed}' + 'i', // 'î'; '\u{ee}' + 'i', // 'ï'; '\u{ef}' + 'd', // 'ð'; '\u{f0}' + 'n', // 'ñ'; '\u{f1}' + 'o', // 'ò'; '\u{f2}' + 'o', // 'ó'; '\u{f3}' + 'o', // 'ô'; '\u{f4}' + 'o', // 'õ'; '\u{f5}' + 'o', // 'ö'; '\u{f6}' + '÷', // '÷'; '\u{f7}' + 'o', // 'ø'; '\u{f8}' + 'u', // 'ù'; '\u{f9}' + 'u', // 'ú'; '\u{fa}' + 'u', // 'û'; '\u{fb}' + 'u', // 'ü'; '\u{fc}' + 'y', // 'ý'; '\u{fd}' + 'þ', // 'þ'; '\u{fe}' + 'y', // 'ÿ'; '\u{ff}' + 'A', // 'Ā'; '\u{100}' + 'a', // 'ā'; '\u{101}' + 'A', // 'Ă'; '\u{102}' + 'a', // 'ă'; '\u{103}' + 'A', // 'Ą'; '\u{104}' + 'a', // 'ą'; '\u{105}' + 'C', // 'Ć'; '\u{106}' + 'c', // 'ć'; '\u{107}' + 'C', // 'Ĉ'; '\u{108}' + 'c', // 'ĉ'; '\u{109}' + 'C', // 'Ċ'; '\u{10a}' + 'c', // 'ċ'; '\u{10b}' + 'C', // 'Č'; '\u{10c}' + 'c', // 'č'; '\u{10d}' + 'D', // 'Ď'; '\u{10e}' + 'd', // 'ď'; '\u{10f}' + 'D', // 'Đ'; '\u{110}' + 'd', // 'đ'; '\u{111}' + 'E', // 'Ē'; '\u{112}' + 'e', // 'ē'; '\u{113}' + 'E', // 'Ĕ'; '\u{114}' + 'e', // 'ĕ'; '\u{115}' + 'E', // 'Ė'; '\u{116}' + 'e', // 'ė'; '\u{117}' + 'E', // 'Ę'; '\u{118}' + 'e', // 'ę'; '\u{119}' + 'E', // 'Ě'; '\u{11a}' + 'e', // 'ě'; '\u{11b}' + 'G', // 'Ĝ'; '\u{11c}' + 'g', // 'ĝ'; '\u{11d}' + 'G', // 'Ğ'; '\u{11e}' + 'g', // 'ğ'; '\u{11f}' + 'G', // 'Ġ'; '\u{120}' + 'g', // 'ġ'; '\u{121}' + 'G', // 'Ģ'; '\u{122}' + 'g', // 'ģ'; '\u{123}' + 'H', // 'Ĥ'; '\u{124}' + 'h', // 'ĥ'; '\u{125}' + 'H', // 'Ħ'; '\u{126}' + 'h', // 'ħ'; '\u{127}' + 'I', // 'Ĩ'; '\u{128}' + 'i', // 'ĩ'; '\u{129}' + 'I', // 'Ī'; '\u{12a}' + 'i', // 'ī'; '\u{12b}' + 'I', // 'Ĭ'; '\u{12c}' + 'i', // 'ĭ'; '\u{12d}' + 'I', // 'Į'; '\u{12e}' + 'i', // 'į'; '\u{12f}' + 'I', // 'İ'; '\u{130}' + 'i', // 'ı'; '\u{131}' + 'IJ', // 'IJ'; '\u{132}' + 'ij', // 'ij'; '\u{133}' + 'J', // 'Ĵ'; '\u{134}' + 'j', // 'ĵ'; '\u{135}' + 'K', // 'Ķ'; '\u{136}' + 'k', // 'ķ'; '\u{137}' + 'ĸ', // 'ĸ'; '\u{138}' + 'L', // 'Ĺ'; '\u{139}' + 'l', // 'ĺ'; '\u{13a}' + 'L', // 'Ļ'; '\u{13b}' + 'l', // 'ļ'; '\u{13c}' + 'L', // 'Ľ'; '\u{13d}' + 'l', // 'ľ'; '\u{13e}' + 'L', // 'Ŀ'; '\u{13f}' + 'l', // 'ŀ'; '\u{140}' + 'L', // 'Ł'; '\u{141}' + 'l', // 'ł'; '\u{142}' + 'N', // 'Ń'; '\u{143}' + 'n', // 'ń'; '\u{144}' + 'N', // 'Ņ'; '\u{145}' + 'n', // 'ņ'; '\u{146}' + 'N', // 'Ň'; '\u{147}' + 'n', // 'ň'; '\u{148}' + 'n', // 'ʼn'; '\u{149}' + 'N', // 'Ŋ'; '\u{14a}' + 'n', // 'ŋ'; '\u{14b}' + 'O', // 'Ō'; '\u{14c}' + 'o', // 'ō'; '\u{14d}' + 'O', // 'Ŏ'; '\u{14e}' + 'o', // 'ŏ'; '\u{14f}' + 'O', // 'Ő'; '\u{150}' + 'o', // 'ő'; '\u{151}' + 'Œ', // 'Œ'; '\u{152}' + 'œ', // 'œ'; '\u{153}' + 'R', // 'Ŕ'; '\u{154}' + 'r', // 'ŕ'; '\u{155}' + 'R', // 'Ŗ'; '\u{156}' + 'r', // 'ŗ'; '\u{157}' + 'R', // 'Ř'; '\u{158}' + 'r', // 'ř'; '\u{159}' + 'S', // 'Ś'; '\u{15a}' + 's', // 'ś'; '\u{15b}' + 'S', // 'Ŝ'; '\u{15c}' + 's', // 'ŝ'; '\u{15d}' + 'S', // 'Ş'; '\u{15e}' + 's', // 'ş'; '\u{15f}' + 'S', // 'Š'; '\u{160}' + 's', // 'š'; '\u{161}' + 'T', // 'Ţ'; '\u{162}' + 't', // 'ţ'; '\u{163}' + 'T', // 'Ť'; '\u{164}' + 't', // 'ť'; '\u{165}' + 'T', // 'Ŧ'; '\u{166}' + 't', // 'ŧ'; '\u{167}' + 'U', // 'Ũ'; '\u{168}' + 'u', // 'ũ'; '\u{169}' + 'U', // 'Ū'; '\u{16a}' + 'u', // 'ū'; '\u{16b}' + 'U', // 'Ŭ'; '\u{16c}' + 'u', // 'ŭ'; '\u{16d}' + 'U', // 'Ů'; '\u{16e}' + 'u', // 'ů'; '\u{16f}' + 'U', // 'Ű'; '\u{170}' + 'u', // 'ű'; '\u{171}' + 'U', // 'Ų'; '\u{172}' + 'u', // 'ų'; '\u{173}' + 'W', // 'Ŵ'; '\u{174}' + 'w', // 'ŵ'; '\u{175}' + 'Y', // 'Ŷ'; '\u{176}' + 'y', // 'ŷ'; '\u{177}' + 'Y', // 'Ÿ'; '\u{178}' + 'Z', // 'Ź'; '\u{179}' + 'z', // 'ź'; '\u{17a}' + 'Z', // 'Ż'; '\u{17b}' + 'z', // 'ż'; '\u{17c}' + 'Z', // 'Ž'; '\u{17d}' + 'z', // 'ž'; '\u{17e}' + 's', // 'ſ'; '\u{17f}' + 'b', // 'ƀ'; '\u{180}' + 'B', // 'Ɓ'; '\u{181}' + 'b', // 'Ƃ'; '\u{182}' + 'b', // 'ƃ'; '\u{183}' + 'b', // 'Ƅ'; '\u{184}' + 'ƅ', // 'ƅ'; '\u{185}' + 'O', // 'Ɔ'; '\u{186}' + 'C', // 'Ƈ'; '\u{187}' + 'c', // 'ƈ'; '\u{188}' + 'D', // 'Ɖ'; '\u{189}' + 'D', // 'Ɗ'; '\u{18a}' + 'd', // 'Ƌ'; '\u{18b}' + 'd', // 'ƌ'; '\u{18c}' + 'ƍ', // 'ƍ'; '\u{18d}' + 'E', // 'Ǝ'; '\u{18e}' + 'e', // 'Ə'; '\u{18f}' + 'E', // 'Ɛ'; '\u{190}' + 'F', // 'Ƒ'; '\u{191}' + 'f', // 'ƒ'; '\u{192}' + 'G', // 'Ɠ'; '\u{193}' + 'Ɣ', // 'Ɣ'; '\u{194}' + 'h', // 'ƕ'; '\u{195}' + 'I', // 'Ɩ'; '\u{196}' + 'I', // 'Ɨ'; '\u{197}' + 'Ƙ', // 'Ƙ'; '\u{198}' + 'k', // 'ƙ'; '\u{199}' + 'l', // 'ƚ'; '\u{19a}' + 'ƛ', // 'ƛ'; '\u{19b}' + 'M', // 'Ɯ'; '\u{19c}' + 'N', // 'Ɲ'; '\u{19d}' + 'n', // 'ƞ'; '\u{19e}' + 'O', // 'Ɵ'; '\u{19f}' + 'O', // 'Ơ'; '\u{1a0}' + 'o', // 'ơ'; '\u{1a1}' + 'Ƣ', // 'Ƣ'; '\u{1a2}' + 'ƣ', // 'ƣ'; '\u{1a3}' + 'P', // 'Ƥ'; '\u{1a4}' + 'p', // 'ƥ'; '\u{1a5}' + 'R', // 'Ʀ'; '\u{1a6}' + 'S', // 'Ƨ'; '\u{1a7}' + 's', // 'ƨ'; '\u{1a8}' + 'Ʃ', // 'Ʃ'; '\u{1a9}' + 'l', // 'ƪ'; '\u{1aa}' + 't', // 'ƫ'; '\u{1ab}' + 'T', // 'Ƭ'; '\u{1ac}' + 't', // 'ƭ'; '\u{1ad}' + 'T', // 'Ʈ'; '\u{1ae}' + 'U', // 'Ư'; '\u{1af}' + 'u', // 'ư'; '\u{1b0}' + 'Ʊ', // 'Ʊ'; '\u{1b1}' + 'V', // 'Ʋ'; '\u{1b2}' + 'Y', // 'Ƴ'; '\u{1b3}' + 'y', // 'ƴ'; '\u{1b4}' + 'Z', // 'Ƶ'; '\u{1b5}' + 'z', // 'ƶ'; '\u{1b6}' + 'Ʒ', // 'Ʒ'; '\u{1b7}' + 'Ƹ', // 'Ƹ'; '\u{1b8}' + 'ƹ', // 'ƹ'; '\u{1b9}' + 'ƺ', // 'ƺ'; '\u{1ba}' + 'ƻ', // 'ƻ'; '\u{1bb}' + 'Ƽ', // 'Ƽ'; '\u{1bc}' + 'ƽ', // 'ƽ'; '\u{1bd}' + 'ƾ', // 'ƾ'; '\u{1be}' + 'ƿ', // 'ƿ'; '\u{1bf}' + 'ǀ', // 'ǀ'; '\u{1c0}' + 'ǁ', // 'ǁ'; '\u{1c1}' + 'ǂ', // 'ǂ'; '\u{1c2}' + '!', // 'ǃ'; '\u{1c3}' + 'DŽ', // 'DŽ'; '\u{1c4}' + 'Dž', // 'Dž'; '\u{1c5}' + 'dž', // 'dž'; '\u{1c6}' + 'LJ', // 'LJ'; '\u{1c7}' + 'Lj', // 'Lj'; '\u{1c8}' + 'lj', // 'lj'; '\u{1c9}' + 'NJ', // 'NJ'; '\u{1ca}' + 'Nj', // 'Nj'; '\u{1cb}' + 'nj', // 'nj'; '\u{1cc}' + 'A', // 'Ǎ'; '\u{1cd}' + 'a', // 'ǎ'; '\u{1ce}' + 'I', // 'Ǐ'; '\u{1cf}' + 'i', // 'ǐ'; '\u{1d0}' + 'O', // 'Ǒ'; '\u{1d1}' + 'o', // 'ǒ'; '\u{1d2}' + 'U', // 'Ǔ'; '\u{1d3}' + 'u', // 'ǔ'; '\u{1d4}' + 'U', // 'Ǖ'; '\u{1d5}' + 'u', // 'ǖ'; '\u{1d6}' + 'U', // 'Ǘ'; '\u{1d7}' + 'u', // 'ǘ'; '\u{1d8}' + 'U', // 'Ǚ'; '\u{1d9}' + 'u', // 'ǚ'; '\u{1da}' + 'U', // 'Ǜ'; '\u{1db}' + 'u', // 'ǜ'; '\u{1dc}' + 'e', // 'ǝ'; '\u{1dd}' + 'A', // 'Ǟ'; '\u{1de}' + 'a', // 'ǟ'; '\u{1df}' + 'A', // 'Ǡ'; '\u{1e0}' + 'a', // 'ǡ'; '\u{1e1}' + 'Æ', // 'Ǣ'; '\u{1e2}' + 'æ', // 'ǣ'; '\u{1e3}' + 'G', // 'Ǥ'; '\u{1e4}' + 'g', // 'ǥ'; '\u{1e5}' + 'G', // 'Ǧ'; '\u{1e6}' + 'g', // 'ǧ'; '\u{1e7}' + 'K', // 'Ǩ'; '\u{1e8}' + 'k', // 'ǩ'; '\u{1e9}' + 'O', // 'Ǫ'; '\u{1ea}' + 'o', // 'ǫ'; '\u{1eb}' + 'O', // 'Ǭ'; '\u{1ec}' + 'o', // 'ǭ'; '\u{1ed}' + 'Ǯ', // 'Ǯ'; '\u{1ee}' + 'ǯ', // 'ǯ'; '\u{1ef}' + 'j', // 'ǰ'; '\u{1f0}' + 'DZ', // 'DZ'; '\u{1f1}' + 'Dz', // 'Dz'; '\u{1f2}' + 'dz', // 'dz'; '\u{1f3}' + 'G', // 'Ǵ'; '\u{1f4}' + 'g', // 'ǵ'; '\u{1f5}' + 'Ƕ', // 'Ƕ'; '\u{1f6}' + 'Ƿ', // 'Ƿ'; '\u{1f7}' + 'N', // 'Ǹ'; '\u{1f8}' + 'n', // 'ǹ'; '\u{1f9}' + 'A', // 'Ǻ'; '\u{1fa}' + 'a', // 'ǻ'; '\u{1fb}' + 'Æ', // 'Ǽ'; '\u{1fc}' + 'æ', // 'ǽ'; '\u{1fd}' + 'O', // 'Ǿ'; '\u{1fe}' + 'o', // 'ǿ'; '\u{1ff}' + 'A', // 'Ȁ'; '\u{200}' + 'a', // 'ȁ'; '\u{201}' + 'A', // 'Ȃ'; '\u{202}' + 'a', // 'ȃ'; '\u{203}' + 'E', // 'Ȅ'; '\u{204}' + 'e', // 'ȅ'; '\u{205}' + 'E', // 'Ȇ'; '\u{206}' + 'e', // 'ȇ'; '\u{207}' + 'I', // 'Ȉ'; '\u{208}' + 'i', // 'ȉ'; '\u{209}' + 'I', // 'Ȋ'; '\u{20a}' + 'i', // 'ȋ'; '\u{20b}' + 'O', // 'Ȍ'; '\u{20c}' + 'o', // 'ȍ'; '\u{20d}' + 'O', // 'Ȏ'; '\u{20e}' + 'o', // 'ȏ'; '\u{20f}' + 'R', // 'Ȑ'; '\u{210}' + 'r', // 'ȑ'; '\u{211}' + 'R', // 'Ȓ'; '\u{212}' + 'r', // 'ȓ'; '\u{213}' + 'U', // 'Ȕ'; '\u{214}' + 'u', // 'ȕ'; '\u{215}' + 'U', // 'Ȗ'; '\u{216}' + 'u', // 'ȗ'; '\u{217}' + 'S', // 'Ș'; '\u{218}' + 's', // 'ș'; '\u{219}' + 'T', // 'Ț'; '\u{21a}' + 't', // 'ț'; '\u{21b}' + 'Ȝ', // 'Ȝ'; '\u{21c}' + 'ȝ', // 'ȝ'; '\u{21d}' + 'H', // 'Ȟ'; '\u{21e}' + 'h', // 'ȟ'; '\u{21f}' + 'N', // 'Ƞ'; '\u{220}' + 'd', // 'ȡ'; '\u{221}' + 'Ȣ', // 'Ȣ'; '\u{222}' + 'ȣ', // 'ȣ'; '\u{223}' + 'Z', // 'Ȥ'; '\u{224}' + 'z', // 'ȥ'; '\u{225}' + 'A', // 'Ȧ'; '\u{226}' + 'a', // 'ȧ'; '\u{227}' + 'E', // 'Ȩ'; '\u{228}' + 'e', // 'ȩ'; '\u{229}' + 'O', // 'Ȫ'; '\u{22a}' + 'o', // 'ȫ'; '\u{22b}' + 'O', // 'Ȭ'; '\u{22c}' + 'o', // 'ȭ'; '\u{22d}' + 'O', // 'Ȯ'; '\u{22e}' + 'o', // 'ȯ'; '\u{22f}' + 'O', // 'Ȱ'; '\u{230}' + 'o', // 'ȱ'; '\u{231}' + 'Y', // 'Ȳ'; '\u{232}' + 'y', // 'ȳ'; '\u{233}' + 'l', // 'ȴ'; '\u{234}' + 'n', // 'ȵ'; '\u{235}' + 't', // 'ȶ'; '\u{236}' + 'j', // 'ȷ'; '\u{237}' + 'ȸ', // 'ȸ'; '\u{238}' + 'ȹ', // 'ȹ'; '\u{239}' + 'A', // 'Ⱥ'; '\u{23a}' + 'C', // 'Ȼ'; '\u{23b}' + 'c', // 'ȼ'; '\u{23c}' + 'L', // 'Ƚ'; '\u{23d}' + 'T', // 'Ⱦ'; '\u{23e}' + 's', // 'ȿ'; '\u{23f}' + 'z', // 'ɀ'; '\u{240}' + 'Ɂ', // 'Ɂ'; '\u{241}' + 'ɂ', // 'ɂ'; '\u{242}' + 'B', // 'Ƀ'; '\u{243}' + 'U', // 'Ʉ'; '\u{244}' + 'V', // 'Ʌ'; '\u{245}' + 'E', // 'Ɇ'; '\u{246}' + 'e', // 'ɇ'; '\u{247}' + 'J', // 'Ɉ'; '\u{248}' + 'j', // 'ɉ'; '\u{249}' + 'Q', // 'Ɋ'; '\u{24a}' + 'q', // 'ɋ'; '\u{24b}' + 'R', // 'Ɍ'; '\u{24c}' + 'r', // 'ɍ'; '\u{24d}' + 'Y', // 'Ɏ'; '\u{24e}' + 'y', // 'ɏ'; '\u{24f}' + 'a', // 'ɐ'; '\u{250}' + 'a', // 'ɑ'; '\u{251}' + 'a', // 'ɒ'; '\u{252}' + 'b', // 'ɓ'; '\u{253}' + 'c', // 'ɔ'; '\u{254}' + 'c', // 'ɕ'; '\u{255}' + 'd', // 'ɖ'; '\u{256}' + 'd', // 'ɗ'; '\u{257}' + 'e', // 'ɘ'; '\u{258}' + 'e', // 'ə'; '\u{259}' + 'e', // 'ɚ'; '\u{25a}' + 'e', // 'ɛ'; '\u{25b}' + 'e', // 'ɜ'; '\u{25c}' + 'e', // 'ɝ'; '\u{25d}' + 'e', // 'ɞ'; '\u{25e}' + 'j', // 'ɟ'; '\u{25f}' + 'g', // 'ɠ'; '\u{260}' + 'g', // 'ɡ'; '\u{261}' + 'G', // 'ɢ'; '\u{262}' + 'g', // 'ɣ'; '\u{263}' + 'u', // 'ɤ'; '\u{264}' + 'h', // 'ɥ'; '\u{265}' + 'h', // 'ɦ'; '\u{266}' + 'h', // 'ɧ'; '\u{267}' + 'i', // 'ɨ'; '\u{268}' + 'i', // 'ɩ'; '\u{269}' + 'I', // 'ɪ'; '\u{26a}' + 'l', // 'ɫ'; '\u{26b}' + 'l', // 'ɬ'; '\u{26c}' + 'l', // 'ɭ'; '\u{26d}' + 'ɮ', // 'ɮ'; '\u{26e}' + 'm', // 'ɯ'; '\u{26f}' + 'm', // 'ɰ'; '\u{270}' + 'm', // 'ɱ'; '\u{271}' + 'n', // 'ɲ'; '\u{272}' + 'n', // 'ɳ'; '\u{273}' + 'N', // 'ɴ'; '\u{274}' + 'o', // 'ɵ'; '\u{275}' + 'ɶ', // 'ɶ'; '\u{276}' + 'ɷ', // 'ɷ'; '\u{277}' + 'ɸ', // 'ɸ'; '\u{278}' + 'r', // 'ɹ'; '\u{279}' + 'r', // 'ɺ'; '\u{27a}' + 'r', // 'ɻ'; '\u{27b}' + 'r', // 'ɼ'; '\u{27c}' + 'r', // 'ɽ'; '\u{27d}' + 'r', // 'ɾ'; '\u{27e}' + 'r', // 'ɿ'; '\u{27f}' + 'R', // 'ʀ'; '\u{280}' + 'R', // 'ʁ'; '\u{281}' + 's', // 'ʂ'; '\u{282}' + 'ʃ', // 'ʃ'; '\u{283}' + 'ʄ', // 'ʄ'; '\u{284}' + 'ʅ', // 'ʅ'; '\u{285}' + 'ʆ', // 'ʆ'; '\u{286}' + 't', // 'ʇ'; '\u{287}' + 't', // 'ʈ'; '\u{288}' + 'u', // 'ʉ'; '\u{289}' + 'ʊ', // 'ʊ'; '\u{28a}' + 'v', // 'ʋ'; '\u{28b}' + 'v', // 'ʌ'; '\u{28c}' + 'w', // 'ʍ'; '\u{28d}' + 'y', // 'ʎ'; '\u{28e}' + 'Y', // 'ʏ'; '\u{28f}' + 'z', // 'ʐ'; '\u{290}' + 'z', // 'ʑ'; '\u{291}' + 'ʒ', // 'ʒ'; '\u{292}' + 'ʓ', // 'ʓ'; '\u{293}' + 'ʔ', // 'ʔ'; '\u{294}' + 'ʕ', // 'ʕ'; '\u{295}' + 'ʖ', // 'ʖ'; '\u{296}' + 'c', // 'ʗ'; '\u{297}' + 'ʘ', // 'ʘ'; '\u{298}' + 'B', // 'ʙ'; '\u{299}' + 'e', // 'ʚ'; '\u{29a}' + 'G', // 'ʛ'; '\u{29b}' + 'H', // 'ʜ'; '\u{29c}' + 'j', // 'ʝ'; '\u{29d}' + 'k', // 'ʞ'; '\u{29e}' + 'L', // 'ʟ'; '\u{29f}' +]; + +/// A char array corresponding to the following Unicode block: +/// +/// - [Latin Extended Additional](https://en.wikipedia.org/wiki/Latin_Extended_Additional) +/// +/// This covers the range `'\u{1e00}'..='\u{1eff}'`. +static LATIN_EXTENDED_ADDITIONAL: [char; 256] = [ + 'A', // 'Ḁ'; '\u{1e00}' + 'a', // 'ḁ'; '\u{1e01}' + 'B', // 'Ḃ'; '\u{1e02}' + 'b', // 'ḃ'; '\u{1e03}' + 'B', // 'Ḅ'; '\u{1e04}' + 'b', // 'ḅ'; '\u{1e05}' + 'B', // 'Ḇ'; '\u{1e06}' + 'b', // 'ḇ'; '\u{1e07}' + 'C', // 'Ḉ'; '\u{1e08}' + 'c', // 'ḉ'; '\u{1e09}' + 'D', // 'Ḋ'; '\u{1e0a}' + 'e', // 'ḋ'; '\u{1e0b}' + 'D', // 'Ḍ'; '\u{1e0c}' + 'd', // 'ḍ'; '\u{1e0d}' + 'D', // 'Ḏ'; '\u{1e0e}' + 'd', // 'ḏ'; '\u{1e0f}' + 'D', // 'Ḑ'; '\u{1e10}' + 'd', // 'ḑ'; '\u{1e11}' + 'D', // 'Ḓ'; '\u{1e12}' + 'd', // 'ḓ'; '\u{1e13}' + 'E', // 'Ḕ'; '\u{1e14}' + 'e', // 'ḕ'; '\u{1e15}' + 'E', // 'Ḗ'; '\u{1e16}' + 'e', // 'ḗ'; '\u{1e17}' + 'E', // 'Ḙ'; '\u{1e18}' + 'e', // 'ḙ'; '\u{1e19}' + 'E', // 'Ḛ'; '\u{1e1a}' + 'e', // 'ḛ'; '\u{1e1b}' + 'E', // 'Ḝ'; '\u{1e1c}' + 'e', // 'ḝ'; '\u{1e1d}' + 'F', // 'Ḟ'; '\u{1e1e}' + 'f', // 'ḟ'; '\u{1e1f}' + 'G', // 'Ḡ'; '\u{1e20}' + 'g', // 'ḡ'; '\u{1e21}' + 'H', // 'Ḣ'; '\u{1e22}' + 'g', // 'ḣ'; '\u{1e23}' + 'H', // 'Ḥ'; '\u{1e24}' + 'g', // 'ḥ'; '\u{1e25}' + 'H', // 'Ḧ'; '\u{1e26}' + 'g', // 'ḧ'; '\u{1e27}' + 'H', // 'Ḩ'; '\u{1e28}' + 'g', // 'ḩ'; '\u{1e29}' + 'H', // 'Ḫ'; '\u{1e2a}' + 'h', // 'ḫ'; '\u{1e2b}' + 'I', // 'Ḭ'; '\u{1e2c}' + 'i', // 'ḭ'; '\u{1e2d}' + 'I', // 'Ḯ'; '\u{1e2e}' + 'i', // 'ḯ'; '\u{1e2f}' + 'K', // 'Ḱ'; '\u{1e30}' + 'k', // 'ḱ'; '\u{1e31}' + 'K', // 'Ḳ'; '\u{1e32}' + 'k', // 'ḳ'; '\u{1e33}' + 'K', // 'Ḵ'; '\u{1e34}' + 'k', // 'ḵ'; '\u{1e35}' + 'L', // 'Ḷ'; '\u{1e36}' + 'l', // 'ḷ'; '\u{1e37}' + 'L', // 'Ḹ'; '\u{1e38}' + 'l', // 'ḹ'; '\u{1e39}' + 'L', // 'Ḻ'; '\u{1e3a}' + 'l', // 'ḻ'; '\u{1e3b}' + 'L', // 'Ḽ'; '\u{1e3c}' + 'l', // 'ḽ'; '\u{1e3d}' + 'M', // 'Ḿ'; '\u{1e3e}' + 'm', // 'ḿ'; '\u{1e3f}' + 'M', // 'Ṁ'; '\u{1e40}' + 'm', // 'ṁ'; '\u{1e41}' + 'M', // 'Ṃ'; '\u{1e42}' + 'm', // 'ṃ'; '\u{1e43}' + 'N', // 'Ṅ'; '\u{1e44}' + 'n', // 'ṅ'; '\u{1e45}' + 'N', // 'Ṇ'; '\u{1e46}' + 'n', // 'ṇ'; '\u{1e47}' + 'N', // 'Ṉ'; '\u{1e48}' + 'n', // 'ṉ'; '\u{1e49}' + 'N', // 'Ṋ'; '\u{1e4a}' + 'n', // 'ṋ'; '\u{1e4b}' + 'O', // 'Ṍ'; '\u{1e4c}' + 'o', // 'ṍ'; '\u{1e4d}' + 'O', // 'Ṏ'; '\u{1e4e}' + 'o', // 'ṏ'; '\u{1e4f}' + 'O', // 'Ṑ'; '\u{1e50}' + 'o', // 'ṑ'; '\u{1e51}' + 'O', // 'Ṓ'; '\u{1e52}' + 'o', // 'ṓ'; '\u{1e53}' + 'P', // 'Ṕ'; '\u{1e54}' + 'p', // 'ṕ'; '\u{1e55}' + 'P', // 'Ṗ'; '\u{1e56}' + 'p', // 'ṗ'; '\u{1e57}' + 'R', // 'Ṙ'; '\u{1e58}' + 'r', // 'ṙ'; '\u{1e59}' + 'R', // 'Ṛ'; '\u{1e5a}' + 'r', // 'ṛ'; '\u{1e5b}' + 'R', // 'Ṝ'; '\u{1e5c}' + 'r', // 'ṝ'; '\u{1e5d}' + 'R', // 'Ṟ'; '\u{1e5e}' + 'r', // 'ṟ'; '\u{1e5f}' + 'S', // 'Ṡ'; '\u{1e60}' + 's', // 'ṡ'; '\u{1e61}' + 'S', // 'Ṣ'; '\u{1e62}' + 's', // 'ṣ'; '\u{1e63}' + 'S', // 'Ṥ'; '\u{1e64}' + 's', // 'ṥ'; '\u{1e65}' + 'S', // 'Ṧ'; '\u{1e66}' + 's', // 'ṧ'; '\u{1e67}' + 'S', // 'Ṩ'; '\u{1e68}' + 's', // 'ṩ'; '\u{1e69}' + 'T', // 'Ṫ'; '\u{1e6a}' + 't', // 'ṫ'; '\u{1e6b}' + 'T', // 'Ṭ'; '\u{1e6c}' + 't', // 'ṭ'; '\u{1e6d}' + 'T', // 'Ṯ'; '\u{1e6e}' + 't', // 'ṯ'; '\u{1e6f}' + 'T', // 'Ṱ'; '\u{1e70}' + 't', // 'ṱ'; '\u{1e71}' + 'U', // 'Ṳ'; '\u{1e72}' + 'u', // 'ṳ'; '\u{1e73}' + 'U', // 'Ṵ'; '\u{1e74}' + 'u', // 'ṵ'; '\u{1e75}' + 'U', // 'Ṷ'; '\u{1e76}' + 'u', // 'ṷ'; '\u{1e77}' + 'U', // 'Ṹ'; '\u{1e78}' + 'u', // 'ṹ'; '\u{1e79}' + 'U', // 'Ṻ'; '\u{1e7a}' + 'u', // 'ṻ'; '\u{1e7b}' + 'V', // 'Ṽ'; '\u{1e7c}' + 'v', // 'ṽ'; '\u{1e7d}' + 'V', // 'Ṿ'; '\u{1e7e}' + 'v', // 'ṿ'; '\u{1e7f}' + 'W', // 'Ẁ'; '\u{1e80}' + 'w', // 'ẁ'; '\u{1e81}' + 'W', // 'Ẃ'; '\u{1e82}' + 'w', // 'ẃ'; '\u{1e83}' + 'W', // 'Ẅ'; '\u{1e84}' + 'w', // 'ẅ'; '\u{1e85}' + 'W', // 'Ẇ'; '\u{1e86}' + 'w', // 'ẇ'; '\u{1e87}' + 'W', // 'Ẉ'; '\u{1e88}' + 'j', // 'ẉ'; '\u{1e89}' + 'X', // 'Ẋ'; '\u{1e8a}' + 'x', // 'ẋ'; '\u{1e8b}' + 'X', // 'Ẍ'; '\u{1e8c}' + 'x', // 'ẍ'; '\u{1e8d}' + 'Y', // 'Ẏ'; '\u{1e8e}' + 'y', // 'ẏ'; '\u{1e8f}' + 'Z', // 'Ẑ'; '\u{1e90}' + 'z', // 'ẑ'; '\u{1e91}' + 'Z', // 'Ẓ'; '\u{1e92}' + 'z', // 'ẓ'; '\u{1e93}' + 'Z', // 'Ẕ'; '\u{1e94}' + 'z', // 'ẕ'; '\u{1e95}' + 'h', // 'ẖ'; '\u{1e96}' + 't', // 'ẗ'; '\u{1e97}' + 'w', // 'ẘ'; '\u{1e98}' + 'y', // 'ẙ'; '\u{1e99}' + 'a', // 'ẚ'; '\u{1e9a}' + 'i', // 'ẛ'; '\u{1e9b}' + 'f', // 'ẜ'; '\u{1e9c}' + 'f', // 'ẝ'; '\u{1e9d}' + 'ẞ', // 'ẞ'; '\u{1e9e}' + 'ẟ', // 'ẟ'; '\u{1e9f}' + 'A', // 'Ạ'; '\u{1ea0}' + 'a', // 'ạ'; '\u{1ea1}' + 'A', // 'Ả'; '\u{1ea2}' + 'a', // 'ả'; '\u{1ea3}' + 'A', // 'Ấ'; '\u{1ea4}' + 'a', // 'ấ'; '\u{1ea5}' + 'A', // 'Ầ'; '\u{1ea6}' + 'a', // 'ầ'; '\u{1ea7}' + 'A', // 'Ẩ'; '\u{1ea8}' + 'a', // 'ẩ'; '\u{1ea9}' + 'A', // 'Ẫ'; '\u{1eaa}' + 'a', // 'ẫ'; '\u{1eab}' + 'A', // 'Ậ'; '\u{1eac}' + 'a', // 'ậ'; '\u{1ead}' + 'A', // 'Ắ'; '\u{1eae}' + 'a', // 'ắ'; '\u{1eaf}' + 'A', // 'Ằ'; '\u{1eb0}' + 'a', // 'ằ'; '\u{1eb1}' + 'A', // 'Ẳ'; '\u{1eb2}' + 'a', // 'ẳ'; '\u{1eb3}' + 'A', // 'Ẵ'; '\u{1eb4}' + 'a', // 'ẵ'; '\u{1eb5}' + 'A', // 'Ặ'; '\u{1eb6}' + 'a', // 'ặ'; '\u{1eb7}' + 'E', // 'Ẹ'; '\u{1eb8}' + 'e', // 'ẹ'; '\u{1eb9}' + 'E', // 'Ẻ'; '\u{1eba}' + 'e', // 'ẻ'; '\u{1ebb}' + 'E', // 'Ẽ'; '\u{1ebc}' + 'e', // 'ẽ'; '\u{1ebd}' + 'E', // 'Ế'; '\u{1ebe}' + 'e', // 'ế'; '\u{1ebf}' + 'E', // 'Ề'; '\u{1ec0}' + 'e', // 'ề'; '\u{1ec1}' + 'E', // 'Ể'; '\u{1ec2}' + 'e', // 'ể'; '\u{1ec3}' + 'E', // 'Ễ'; '\u{1ec4}' + 'e', // 'ễ'; '\u{1ec5}' + 'E', // 'Ệ'; '\u{1ec6}' + 'e', // 'ệ'; '\u{1ec7}' + 'I', // 'Ỉ'; '\u{1ec8}' + 'i', // 'ỉ'; '\u{1ec9}' + 'I', // 'Ị'; '\u{1eca}' + 'i', // 'ị'; '\u{1ecb}' + 'O', // 'Ọ'; '\u{1ecc}' + 'o', // 'ọ'; '\u{1ecd}' + 'O', // 'Ỏ'; '\u{1ece}' + 'o', // 'ỏ'; '\u{1ecf}' + 'O', // 'Ố'; '\u{1ed0}' + 'o', // 'ố'; '\u{1ed1}' + 'O', // 'Ồ'; '\u{1ed2}' + 'o', // 'ồ'; '\u{1ed3}' + 'O', // 'Ổ'; '\u{1ed4}' + 'o', // 'ổ'; '\u{1ed5}' + 'O', // 'Ỗ'; '\u{1ed6}' + 'o', // 'ỗ'; '\u{1ed7}' + 'O', // 'Ộ'; '\u{1ed8}' + 'o', // 'ộ'; '\u{1ed9}' + 'O', // 'Ớ'; '\u{1eda}' + 'o', // 'ớ'; '\u{1edb}' + 'O', // 'Ờ'; '\u{1edc}' + 'o', // 'ờ'; '\u{1edd}' + 'O', // 'Ở'; '\u{1ede}' + 'o', // 'ở'; '\u{1edf}' + 'O', // 'Ỡ'; '\u{1ee0}' + 'o', // 'ỡ'; '\u{1ee1}' + 'O', // 'Ợ'; '\u{1ee2}' + 'o', // 'ợ'; '\u{1ee3}' + 'U', // 'Ụ'; '\u{1ee4}' + 'u', // 'ụ'; '\u{1ee5}' + 'U', // 'Ủ'; '\u{1ee6}' + 'u', // 'ủ'; '\u{1ee7}' + 'U', // 'Ứ'; '\u{1ee8}' + 'u', // 'ứ'; '\u{1ee9}' + 'U', // 'Ừ'; '\u{1eea}' + 'u', // 'ừ'; '\u{1eeb}' + 'U', // 'Ử'; '\u{1eec}' + 'u', // 'ử'; '\u{1eed}' + 'U', // 'Ữ'; '\u{1eee}' + 'u', // 'ữ'; '\u{1eef}' + 'U', // 'Ự'; '\u{1ef0}' + 'u', // 'ự'; '\u{1ef1}' + 'Y', // 'Ỳ'; '\u{1ef2}' + 'y', // 'ỳ'; '\u{1ef3}' + 'Y', // 'Ỵ'; '\u{1ef4}' + 'y', // 'ỵ'; '\u{1ef5}' + 'Y', // 'Ỷ'; '\u{1ef6}' + 'y', // 'ỷ'; '\u{1ef7}' + 'Y', // 'Ỹ'; '\u{1ef8}' + 'y', // 'ỹ'; '\u{1ef9}' + 'Ỻ', // 'Ỻ'; '\u{1efa}' + 'ỻ', // 'ỻ'; '\u{1efb}' + 'Ỽ', // 'Ỽ'; '\u{1efc}' + 'ỽ', // 'ỽ'; '\u{1efd}' + 'Ỿ', // 'Ỿ'; '\u{1efe}' + 'ỿ', // 'ỿ'; '\u{1eff}' +]; + +/// A char array corresponding to the following Unicode block: +/// +/// - [Superscripts and Subscripts](https://en.wikipedia.org/wiki/Superscripts_and_Subscripts) +/// +/// This covers the range `'\u{2070}'..='\u{209f}'`. +static SUPERSCRIPTS_AND_SUBSCRIPTS: [char; 48] = [ + '0', // '⁰'; '\u{2070}' + 'i', // 'ⁱ'; '\u{2071}' + '', // ''; '\u{2072}' + '', // ''; '\u{2073}' + '4', // '⁴'; '\u{2074}' + '5', // '⁵'; '\u{2075}' + '6', // '⁶'; '\u{2076}' + '7', // '⁷'; '\u{2077}' + '8', // '⁸'; '\u{2078}' + '0', // '⁹'; '\u{2079}' + '+', // '⁺'; '\u{207a}' + '-', // '⁻'; '\u{207b}' + '=', // '⁼'; '\u{207c}' + '(', // '⁽'; '\u{207d}' + ')', // '⁾'; '\u{207e}' + 'n', // 'ⁿ'; '\u{207f}' + '0', // '₀'; '\u{2080}' + '1', // '₁'; '\u{2081}' + '2', // '₂'; '\u{2082}' + '3', // '₃'; '\u{2083}' + '4', // '₄'; '\u{2084}' + '5', // '₅'; '\u{2085}' + '6', // '₆'; '\u{2086}' + '7', // '₇'; '\u{2087}' + '8', // '₈'; '\u{2088}' + '9', // '₉'; '\u{2089}' + '+', // '₊'; '\u{208a}' + '-', // '₋'; '\u{208b}' + '=', // '₌'; '\u{208c}' + '(', // '₍'; '\u{208d}' + ')', // '₎'; '\u{208e}' + '', // ''; '\u{208f}' + 'a', // 'ₐ'; '\u{2090}' + 'e', // 'ₑ'; '\u{2091}' + 'o', // 'ₒ'; '\u{2092}' + 'x', // 'ₓ'; '\u{2093}' + 'e', // 'ₔ'; '\u{2094}' + 'h', // 'ₕ'; '\u{2095}' + 'k', // 'ₖ'; '\u{2096}' + 'l', // 'ₗ'; '\u{2097}' + 'm', // 'ₘ'; '\u{2098}' + 'n', // 'ₙ'; '\u{2099}' + 'p', // 'ₚ'; '\u{209a}' + 's', // 'ₛ'; '\u{209b}' + 't', // 'ₜ'; '\u{209c}' + '', // ''; '\u{209d}' + '', // ''; '\u{209e}' + '', // ''; '\u{209f}' +]; + +#[cfg(test)] +mod tests { + use super::*; + + /// Helper function for test assertions. + fn check_conversions(pairs: &[(char, char)]) { + for (original, normalized) in pairs { + assert_eq!(normalize(*original), *normalized); + } + } + + /// General conversion checks + #[test] + fn general() { + check_conversions(&[ + ('ą', 'a'), + ('À', 'A'), + ('ć', 'c'), + ('ę', 'e'), + ('ł', 'l'), + ('ń', 'n'), + ('ó', 'o'), + ('ś', 's'), + ('ź', 'z'), + ('ż', 'z'), + ('Ą', 'A'), + ('Ć', 'C'), + ('Ę', 'E'), + ('ł', 'l'), + ('Ł', 'L'), + ('Ń', 'N'), + ('Ó', 'O'), + ('Ś', 'S'), + ('Ź', 'Z'), + ('Ż', 'Z'), + ('¡', '!'), + ]); + } + + /// Some checks for characters which are not visible. + #[test] + fn invisible_chars() { + check_conversions(&[('\u{a0}', '\u{a0}'), ('\u{ad}', '\u{ad}')]); + } + + /// Check boundary cases in case ranges are modified. + #[test] + fn boundary_cases() { + check_conversions(&[ + ('\u{9f}', '\u{9f}'), + ('\u{a0}', '\u{a0}'), + ('¡', '!'), + ('ʟ', 'L'), + ('\u{2a0}', '\u{2a0}'), + ('\u{1dff}', '\u{1dff}'), + ('Ḁ', 'A'), + ('ỹ', 'y'), + ('\u{1eff}', '\u{1eff}'), + ('\u{1f00}', '\u{1f00}'), + ('⁰', '0'), + ('\u{209c}', 't'), + ('\u{209f}', '\u{209f}'), + ('\u{20a0}', '\u{20a0}'), + ]); + } + + /// Check that conversions outside the blocks are unchanged. + #[test] + fn unchanged_outside_blocks() { + check_conversions(&[ + ('a', 'a'), + ('⟁', '⟁'), + ('┍', '┍'), + ('ω', 'ω'), + ('⁕', '⁕'), + ('ה', 'ה'), + ]); + } +} diff --git a/crates/atuin-nucleo/matcher/src/config.rs b/crates/atuin-nucleo/matcher/src/config.rs new file mode 100644 index 00000000..eca7ae38 --- /dev/null +++ b/crates/atuin-nucleo/matcher/src/config.rs @@ -0,0 +1,70 @@ +use crate::chars::CharClass; +use crate::score::BONUS_BOUNDARY; + +/// Configuration data that controls how a matcher behaves +#[non_exhaustive] +#[derive(PartialEq, Eq, Debug, Clone)] +pub struct Config { + /// Characters that act as delimiters and provide bonus + /// for matching the following char + pub(crate) delimiter_chars: &'static [u8], + /// Extra bonus for word boundary after whitespace character or beginning of the string + pub(crate) bonus_boundary_white: u16, + /// Extra bonus for word boundary after slash, colon, semi-colon, and comma + pub(crate) bonus_boundary_delimiter: u16, + pub(crate) initial_char_class: CharClass, + + /// Whether to normalize latin script characters to ASCII (enabled by default) + pub normalize: bool, + /// whether to ignore casing + pub ignore_case: bool, + /// Whether to provide a bonus to matches by their distance from the start + /// of the haystack. The bonus is fairly small compared to the normal gap + /// penalty to avoid messing with the normal score heuristic. This setting + /// is not turned on by default and only recommended for autocompletion + /// usecases where the expectation is that the user is typing the entire + /// match. For a full fzf-like fuzzy matcher/picker word segmentation and + /// explicit prefix literals should be used instead. + pub prefer_prefix: bool, +} + +impl Config { + /// The default config for nucleo, implemented as a constant since + /// Default::default can not be called in a const context + pub const DEFAULT: Self = { + Config { + delimiter_chars: b"/,:;|", + bonus_boundary_white: BONUS_BOUNDARY + 2, + bonus_boundary_delimiter: BONUS_BOUNDARY + 1, + initial_char_class: CharClass::Whitespace, + normalize: true, + ignore_case: true, + prefer_prefix: false, + } + }; +} + +impl Config { + /// Configures the matcher with bonuses appropriate for matching file paths. + pub fn set_match_paths(&mut self) { + if cfg!(windows) { + self.delimiter_chars = b"/:\\"; + } else { + self.delimiter_chars = b"/:"; + } + self.bonus_boundary_white = BONUS_BOUNDARY; + self.initial_char_class = CharClass::Delimiter; + } + + /// Configures the matcher with bonuses appropriate for matching file paths. + pub const fn match_paths(mut self) -> Self { + if cfg!(windows) { + self.delimiter_chars = b"/\\"; + } else { + self.delimiter_chars = b"/"; + } + self.bonus_boundary_white = BONUS_BOUNDARY; + self.initial_char_class = CharClass::Delimiter; + self + } +} diff --git a/crates/atuin-nucleo/matcher/src/debug.rs b/crates/atuin-nucleo/matcher/src/debug.rs new file mode 100644 index 00000000..b8369f32 --- /dev/null +++ b/crates/atuin-nucleo/matcher/src/debug.rs @@ -0,0 +1,14 @@ +use crate::matrix::{MatrixCell, ScoreCell}; +use std::fmt::{Debug, Formatter, Result}; + +impl Debug for ScoreCell { + fn fmt(&self, f: &mut Formatter<'_>) -> Result { + write!(f, "({}, {})", self.score, self.matched) + } +} + +impl Debug for MatrixCell { + fn fmt(&self, f: &mut Formatter<'_>) -> Result { + write!(f, "({}, {})", (self.0 & 1) != 0, (self.0 & 2) != 0) + } +} diff --git a/crates/atuin-nucleo/matcher/src/exact.rs b/crates/atuin-nucleo/matcher/src/exact.rs new file mode 100644 index 00000000..3cb3ceb2 --- /dev/null +++ b/crates/atuin-nucleo/matcher/src/exact.rs @@ -0,0 +1,275 @@ +use memchr::memmem; +use memchr::{Memchr, Memchr2}; + +use crate::chars::{AsciiChar, Char}; +use crate::score::{BONUS_FIRST_CHAR_MULTIPLIER, SCORE_MATCH}; +use crate::Matcher; + +impl Matcher { + pub(crate) fn substring_match_1_ascii<const INDICES: bool>( + &mut self, + haystack: &[u8], + c: u8, + indices: &mut Vec<u32>, + ) -> Option<u16> { + let mut max_score = 0; + let mut max_pos = 0; + if self.config.ignore_case && c >= b'a' && c <= b'z' { + for i in Memchr2::new(c, c - 32, haystack) { + let prev_char_class = i + .checked_sub(1) + .map(|i| AsciiChar(haystack[i]).char_class(&self.config)) + .unwrap_or(self.config.initial_char_class); + let char_class = AsciiChar(haystack[i]).char_class(&self.config); + let bonus = self.config.bonus_for(prev_char_class, char_class); + let score = bonus * BONUS_FIRST_CHAR_MULTIPLIER + SCORE_MATCH; + if score > max_score { + max_pos = i as u32; + max_score = score; + // can't get better than this + if bonus >= self.config.bonus_boundary_white { + break; + } + } + } + } else { + let char_class = AsciiChar(c).char_class(&self.config); + for i in Memchr::new(c, haystack) { + let prev_char_class = i + .checked_sub(1) + .map(|i| AsciiChar(haystack[i]).char_class(&self.config)) + .unwrap_or(self.config.initial_char_class); + let bonus = self.config.bonus_for(prev_char_class, char_class); + let score = bonus * BONUS_FIRST_CHAR_MULTIPLIER + SCORE_MATCH; + if score > max_score { + max_pos = i as u32; + max_score = score; + // can't get better than this + if bonus >= self.config.bonus_boundary_white { + break; + } + } + } + } + if max_score == 0 { + return None; + } + + if INDICES { + indices.push(max_pos); + } + Some(max_score) + } + + pub(crate) fn substring_match_ascii_with_prefilter( + &mut self, + haystack: &[u8], + needle: &[u8], + prefilter_len: usize, + prefilter: impl Iterator<Item = usize>, + ) -> (u16, usize) { + let needle_without_prefilter = &needle[prefilter_len..]; + let mut max_score = 0; + let mut max_pos = 0; + for i in prefilter { + let prev_char_class = i + .checked_sub(1) + .map(|i| AsciiChar(haystack[i]).char_class(&self.config)) + .unwrap_or(self.config.initial_char_class); + let char_class = AsciiChar(haystack[i]).char_class(&self.config); + let bonus = self.config.bonus_for(prev_char_class, char_class); + let score = bonus * BONUS_FIRST_CHAR_MULTIPLIER + SCORE_MATCH; + if score > max_score + && haystack[i + prefilter_len..(i + needle.len()).min(haystack.len())] + .iter() + .map(|&c| AsciiChar(c).normalize(&self.config).0) + .eq(needle_without_prefilter.iter().copied()) + { + max_pos = i; + max_score = score; + // can't get better than this + if bonus >= self.config.bonus_boundary_white { + break; + } + } + } + (max_score, max_pos) + } + + pub(crate) fn substring_match_ascii<const INDICES: bool>( + &mut self, + haystack: &[u8], + needle: &[u8], + indices: &mut Vec<u32>, + ) -> Option<u16> { + let mut max_score = 0; + let mut max_pos = 0; + if self.config.ignore_case { + match needle.iter().position(|&c| c >= b'a' && c <= b'z') { + // start with char do case insensitive search + Some(0) => { + (max_score, max_pos) = self.substring_match_ascii_with_prefilter( + haystack, + needle, + 1, + Memchr2::new( + needle[0], + needle[0] - 32, + &haystack[..haystack.len() - needle.len() + 1], + ), + ); + if max_score == 0 { + return None; + } + } + Some(1) => { + (max_score, max_pos) = self.substring_match_ascii_with_prefilter( + haystack, + needle, + 1, + Memchr::new(needle[0], &haystack[..haystack.len() - needle.len() + 1]), + ); + if max_score == 0 { + return None; + } + } + Some(len) => { + (max_score, max_pos) = self.substring_match_ascii_with_prefilter( + haystack, + needle, + 1, + memmem::find_iter(&haystack[..haystack.len() - needle.len() + len], needle), + ); + if max_score == 0 { + return None; + } + } + // in case we don't have any letter in the needle + // we can treat the search as case sensitive and use memmem directly which is way faster + None => (), + } + } + + if max_score == 0 { + let char_class = AsciiChar(needle[0]).char_class(&self.config); + for i in memmem::find_iter(haystack, needle) { + let prev_char_class = i + .checked_sub(1) + .map(|i| AsciiChar(haystack[i]).char_class(&self.config)) + .unwrap_or(self.config.initial_char_class); + let bonus = self.config.bonus_for(prev_char_class, char_class); + let score = bonus * BONUS_FIRST_CHAR_MULTIPLIER + SCORE_MATCH; + if score > max_score { + max_pos = i; + max_score = score; + // can't get better than this + if bonus >= self.config.bonus_boundary_white { + break; + } + } + } + if max_score == 0 { + return None; + } + } + let score = self.calculate_score::<INDICES, _, _>( + AsciiChar::cast(haystack), + AsciiChar::cast(needle), + max_pos, + max_pos + needle.len(), + indices, + ); + Some(score) + } + + pub(crate) fn substring_match_1_non_ascii<const INDICES: bool>( + &mut self, + haystack: &[char], + needle: char, + start: usize, + indices: &mut Vec<u32>, + ) -> u16 { + let mut max_score = 0; + let mut max_pos = 0; + let mut prev_class = start + .checked_sub(1) + .map(|i| haystack[i].char_class(&self.config)) + .unwrap_or(self.config.initial_char_class); + for (i, &c) in haystack[start..].iter().enumerate() { + let (c, char_class) = c.char_class_and_normalize(&self.config); + if c != needle { + continue; + } + let bonus = self.config.bonus_for(prev_class, char_class); + prev_class = char_class; + let score = bonus * BONUS_FIRST_CHAR_MULTIPLIER + SCORE_MATCH; + if score > max_score { + max_pos = i as u32; + max_score = score; + // can't get better than this + if bonus >= self.config.bonus_boundary_white { + break; + } + } + } + + if INDICES { + indices.push(max_pos + start as u32); + } + max_score + } + + pub(crate) fn substring_match_non_ascii<const INDICES: bool, N>( + &mut self, + haystack: &[char], + needle: &[N], + start: usize, + indices: &mut Vec<u32>, + ) -> Option<u16> + where + N: Char, + char: PartialEq<N>, + { + let mut max_score = 0; + let mut max_pos = 0; + let mut prev_class = start + .checked_sub(1) + .map(|i| haystack[i].char_class(&self.config)) + .unwrap_or(self.config.initial_char_class); + let end = haystack.len() - needle.len(); + for (i, &c) in haystack[start..end].iter().enumerate() { + let (c, char_class) = c.char_class_and_normalize(&self.config); + if c != needle[0] { + continue; + } + let bonus = self.config.bonus_for(prev_class, char_class); + prev_class = char_class; + let score = bonus * BONUS_FIRST_CHAR_MULTIPLIER + SCORE_MATCH; + if score > max_score + && haystack[start + i + 1..start + i + needle.len()] + .iter() + .map(|c| c.normalize(&self.config)) + .eq(needle[1..].iter().copied()) + { + max_pos = i; + max_score = score; + // can't get better than this + if bonus >= self.config.bonus_boundary_white { + break; + } + } + } + if max_score == 0 { + return None; + } + + let score = self.calculate_score::<INDICES, _, _>( + haystack, + needle, + start + max_pos, + start + max_pos + needle.len(), + indices, + ); + Some(score) + } +} diff --git a/crates/atuin-nucleo/matcher/src/fuzzy_greedy.rs b/crates/atuin-nucleo/matcher/src/fuzzy_greedy.rs new file mode 100644 index 00000000..8215bf31 --- /dev/null +++ b/crates/atuin-nucleo/matcher/src/fuzzy_greedy.rs @@ -0,0 +1,51 @@ +use crate::chars::Char; +use crate::Matcher; + +impl Matcher { + /// greedy fallback algorithm, much faster (linear time) but reported scores/indicies + /// might not be the best match + pub(crate) fn fuzzy_match_greedy_<const INDICES: bool, H: Char + PartialEq<N>, N: Char>( + &mut self, + haystack: &[H], + needle: &[N], + mut start: usize, + mut end: usize, + indices: &mut Vec<u32>, + ) -> Option<u16> { + let first_char_end = if H::ASCII && N::ASCII { start + 1 } else { end }; + 'nonascii: { + if !H::ASCII || !N::ASCII { + let mut needle_iter = needle[1..].iter().copied(); + if let Some(mut needle_char) = needle_iter.next() { + for (i, &c) in haystack[first_char_end..].iter().enumerate() { + if c.normalize(&self.config) == needle_char { + let Some(next_needle_char) = needle_iter.next() else { + // we found a match so we are now in the same state + // as the prefilter would produce + end = first_char_end + i + 1; + break 'nonascii; + }; + needle_char = next_needle_char; + } + } + // some needle chars were not matched bail out + return None; + } + } + } // minimize the greedly match by greedy matching in reverse + + let mut needle_iter = needle.iter().rev().copied(); + let mut needle_char = needle_iter.next().unwrap(); + for (i, &c) in haystack[start..end].iter().enumerate().rev() { + let c = c.normalize(&self.config); + if c == needle_char { + let Some(next_needle_char) = needle_iter.next() else { + start += i; + break; + }; + needle_char = next_needle_char; + } + } + Some(self.calculate_score::<INDICES, H, N>(haystack, needle, start, end, indices)) + } +} diff --git a/crates/atuin-nucleo/matcher/src/fuzzy_optimal.rs b/crates/atuin-nucleo/matcher/src/fuzzy_optimal.rs new file mode 100644 index 00000000..5d53ecfb --- /dev/null +++ b/crates/atuin-nucleo/matcher/src/fuzzy_optimal.rs @@ -0,0 +1,348 @@ +use std::cmp::max; + +use crate::chars::{Char, CharClass}; +use crate::matrix::{MatcherDataView, MatrixCell, ScoreCell}; +use crate::score::{ + BONUS_BOUNDARY, BONUS_CONSECUTIVE, BONUS_FIRST_CHAR_MULTIPLIER, MAX_PREFIX_BONUS, + PENALTY_GAP_EXTENSION, PENALTY_GAP_START, PREFIX_BONUS_SCALE, SCORE_MATCH, +}; +use crate::{Config, Matcher}; + +impl Matcher { + pub(crate) fn fuzzy_match_optimal<const INDICES: bool, H: Char + PartialEq<N>, N: Char>( + &mut self, + haystack: &[H], + needle: &[N], + start: usize, + greedy_end: usize, + end: usize, + indices: &mut Vec<u32>, + ) -> Option<u16> { + // construct a matrix (and copy the haystack), the matrix and haystack size are bounded + // to avoid the slow O(mn) time complexity for large inputs. Furthermore, it allows + // us to treat needle indices as u16 + let Some(mut matrix) = self.slab.alloc(&haystack[start..end], needle.len()) else { + return self.fuzzy_match_greedy_::<INDICES, H, N>( + haystack, needle, start, greedy_end, indices, + ); + }; + + let prev_class = start + .checked_sub(1) + .map(|i| haystack[i].char_class(&self.config)) + .unwrap_or(self.config.initial_char_class); + let matched = matrix.setup::<INDICES, _>(needle, prev_class, &self.config, start as u32); + // this only happened with unicode haystacks, for ASCII the prefilter handles all rejects + if !matched { + assert!( + !N::ASCII || !H::ASCII, + "Non-match should have been caught by prefilter. Maybe `needle` is not normalized?" + ); + return None; + } + + // populate the matrix and find the best score + let matrix_len = matrix.populate_matrix::<INDICES, _>(needle); + let last_row_off = matrix.row_offs[needle.len() - 1]; + let relative_last_row_off = last_row_off as usize + 1 - needle.len(); + let (match_end, match_score_cell) = matrix.current_row[relative_last_row_off..] + .iter() + .enumerate() + .max_by_key(|(_, cell)| cell.score) + .expect("there must be atleast one match"); + if INDICES { + matrix.reconstruct_optimal_path(match_end as u16, indices, matrix_len, start as u32); + } + Some(match_score_cell.score) + } +} + +const UNMATCHED: ScoreCell = ScoreCell { + score: 0, + // if matched is true then the consecutive bonus + // is always atleast BONUS_CONSECUTIVE so + // this constant can never occur naturally + consecutive_bonus: 0, + matched: true, +}; + +fn next_m_cell(p_score: u16, bonus: u16, m_cell: ScoreCell) -> ScoreCell { + if m_cell == UNMATCHED { + return ScoreCell { + score: p_score + bonus + SCORE_MATCH, + matched: false, + consecutive_bonus: bonus as u8, + }; + } + + let mut consecutive_bonus = max(m_cell.consecutive_bonus as u16, BONUS_CONSECUTIVE); + if bonus >= BONUS_BOUNDARY && bonus > consecutive_bonus { + consecutive_bonus = bonus + } + + let score_match = m_cell.score + max(consecutive_bonus, bonus); + let score_skip = p_score + bonus; + if score_match > score_skip { + ScoreCell { + score: score_match + SCORE_MATCH, + matched: true, + consecutive_bonus: consecutive_bonus as u8, + } + } else { + ScoreCell { + score: score_skip + SCORE_MATCH, + matched: false, + consecutive_bonus: bonus as u8, + } + } +} + +fn p_score(prev_p_score: u16, prev_m_score: u16) -> (u16, bool) { + let score_match = prev_m_score.saturating_sub(PENALTY_GAP_START); + let score_skip = prev_p_score.saturating_sub(PENALTY_GAP_EXTENSION); + if score_match > score_skip { + (score_match, true) + } else { + (score_skip, false) + } +} + +impl<H: Char> MatcherDataView<'_, H> { + fn setup<const INDICES: bool, N: Char>( + &mut self, + needle: &[N], + mut prev_class: CharClass, + config: &Config, + start: u32, + ) -> bool + where + H: PartialEq<N>, + { + let mut row_iter = needle.iter().copied().zip(self.row_offs.iter_mut()); + let (mut needle_char, mut row_start) = row_iter.next().unwrap(); + + let col_iter = self + .haystack + .iter_mut() + .zip(self.bonus.iter_mut()) + .enumerate(); + + let mut matched = false; + for (i, (c_, bonus_)) in col_iter { + let (c, class) = c_.char_class_and_normalize(config); + *c_ = c; + + let bonus = config.bonus_for(prev_class, class); + // save bonus for later so we don't have to recompute it each time + *bonus_ = bonus as u8; + prev_class = class; + + let i = i as u16; + if c == needle_char { + // save the first idx of each char + if let Some(next) = row_iter.next() { + *row_start = i; + (needle_char, row_start) = next; + } else if !matched { + *row_start = i; + // we have atleast one match + matched = true; + } + } + } + if !matched { + return false; + } + debug_assert_eq!(self.row_offs[0], 0); + Self::score_row::<true, INDICES, _>( + self.current_row, + self.matrix_cells, + self.haystack, + self.bonus, + 0, + self.row_offs[1], + 0, + needle[0], + needle[1], + if config.prefer_prefix { + if start == 0 { + MAX_PREFIX_BONUS * PREFIX_BONUS_SCALE + } else { + (MAX_PREFIX_BONUS * PREFIX_BONUS_SCALE - PENALTY_GAP_START).saturating_sub( + (start - 1).min(u16::MAX as u32) as u16 * PENALTY_GAP_EXTENSION, + ) + } + } else { + 0 + }, + ); + true + } + + #[allow(clippy::too_many_arguments)] + fn score_row<const FIRST_ROW: bool, const INDICES: bool, N: Char>( + current_row: &mut [ScoreCell], + matrix_cells: &mut [MatrixCell], + haystack: &[H], + bonus: &[u8], + row_off: u16, + mut next_row_off: u16, + needle_idx: u16, + needle_char: N, + next_needle_char: N, + mut prefix_bonus: u16, + ) where + H: PartialEq<N>, + { + next_row_off -= 1; + let relative_row_off = row_off - needle_idx; + let next_relative_row_off = next_row_off - needle_idx; + let skipped_col_iter = haystack[row_off as usize..next_row_off as usize] + .iter() + .zip(bonus[row_off as usize..next_row_off as usize].iter()) + .zip(current_row[relative_row_off as usize..next_relative_row_off as usize].iter_mut()) + .zip(matrix_cells.iter_mut()); + let mut prev_p_score = 0; + let mut prev_m_score = 0; + for (((&c, bonus), score_cell), matrix_cell) in skipped_col_iter { + let (p_score, p_matched) = p_score(prev_p_score, prev_m_score); + let m_cell = if FIRST_ROW { + let cell = if c == needle_char { + ScoreCell { + score: *bonus as u16 * BONUS_FIRST_CHAR_MULTIPLIER + + SCORE_MATCH + + prefix_bonus / PREFIX_BONUS_SCALE, + matched: false, + consecutive_bonus: *bonus, + } + } else { + UNMATCHED + }; + prefix_bonus = prefix_bonus.saturating_sub(PENALTY_GAP_EXTENSION); + cell + } else { + *score_cell + }; + if INDICES { + matrix_cell.set(p_matched, m_cell.matched); + } + prev_p_score = p_score; + prev_m_score = m_cell.score; + } + let col_iter = haystack[next_row_off as usize..] + .windows(2) + .zip(bonus[next_row_off as usize..].windows(2)) + .zip(current_row[next_relative_row_off as usize..].iter_mut()) + .zip(matrix_cells[(next_relative_row_off - relative_row_off) as usize..].iter_mut()); + for (((c, bonus), score_cell), matrix_cell) in col_iter { + let (p_score, p_matched) = p_score(prev_p_score, prev_m_score); + let m_cell = if FIRST_ROW { + let cell = if c[0] == needle_char { + ScoreCell { + score: bonus[0] as u16 * BONUS_FIRST_CHAR_MULTIPLIER + + SCORE_MATCH + + prefix_bonus / PREFIX_BONUS_SCALE, + matched: false, + consecutive_bonus: bonus[0], + } + } else { + UNMATCHED + }; + prefix_bonus = prefix_bonus.saturating_sub(PENALTY_GAP_EXTENSION); + cell + } else { + *score_cell + }; + *score_cell = if c[1] == next_needle_char { + next_m_cell(p_score, bonus[1] as u16, m_cell) + } else { + UNMATCHED + }; + if INDICES { + matrix_cell.set(p_matched, m_cell.matched); + } + prev_p_score = p_score; + prev_m_score = m_cell.score; + } + } + + fn populate_matrix<const INDICES: bool, N: Char>(&mut self, needle: &[N]) -> usize + where + H: PartialEq<N>, + { + let mut matrix_cells = &mut self.matrix_cells[self.current_row.len()..]; + let mut row_iter = needle[1..] + .iter() + .copied() + .zip(self.row_offs[1..].iter().copied()) + .enumerate(); + let (mut needle_idx, (mut needle_char, mut row_off)) = row_iter.next().unwrap(); + for (next_needle_idx, (next_needle_char, next_row_off)) in row_iter { + Self::score_row::<false, INDICES, _>( + self.current_row, + matrix_cells, + self.haystack, + self.bonus, + row_off, + next_row_off, + needle_idx as u16 + 1, + needle_char, + next_needle_char, + 0, + ); + let len = self.current_row.len() + needle_idx + 1 - row_off as usize; + matrix_cells = &mut matrix_cells[len..]; + (needle_idx, needle_char, row_off) = (next_needle_idx, next_needle_char, next_row_off); + } + matrix_cells.as_ptr() as usize - self.matrix_cells.as_ptr() as usize + } + + fn reconstruct_optimal_path( + &self, + max_score_end: u16, + indices: &mut Vec<u32>, + matrix_len: usize, + start: u32, + ) { + let indices_start = indices.len(); + indices.resize(indices_start + self.row_offs.len(), 0); + let indices = &mut indices[indices_start..]; + let last_row_off = *self.row_offs.last().unwrap(); + indices[self.row_offs.len() - 1] = start + max_score_end as u32 + last_row_off as u32; + + let mut matrix_cells = &self.matrix_cells[..matrix_len]; + let width = self.current_row.len(); + let mut row_iter = self.row_offs[..self.row_offs.len() - 1] + .iter() + .copied() + .enumerate() + .rev() + .map(|(i, off)| { + let relative_off = off as usize - i; + let row; + (matrix_cells, row) = + matrix_cells.split_at(matrix_cells.len() - (width - relative_off)); + (i, off, row) + }); + let (mut row_idx, mut row_off, mut row) = row_iter.next().unwrap(); + let mut col = max_score_end; + let relative_last_row_off = last_row_off as usize + 1 - self.row_offs.len(); + let mut matched = self.current_row[col as usize + relative_last_row_off].matched; + col += last_row_off - row_off - 1; + loop { + if matched { + indices[row_idx] = start + col as u32 + row_off as u32; + } + let next_matched = row[col as usize].get(matched); + if matched { + let Some((next_row_idx, next_row_off, next_row)) = row_iter.next() else { + break; + }; + col += row_off - next_row_off; + (row_idx, row_off, row) = (next_row_idx, next_row_off, next_row) + } + col -= 1; + matched = next_matched; + } + } +} diff --git a/crates/atuin-nucleo/matcher/src/lib.rs b/crates/atuin-nucleo/matcher/src/lib.rs new file mode 100644 index 00000000..3e8874c5 --- /dev/null +++ b/crates/atuin-nucleo/matcher/src/lib.rs @@ -0,0 +1,780 @@ +/*! +`nucleo_matcher` is a low level crate that contains the matcher implementation +used by the high level `nucleo` crate. + +**NOTE**: If you are building an fzf-like interactive fuzzy finder that is +meant to match a reasonably large number of items (> 100) using the high level +`nucleo` crate is highly recommended. Using `nucleo-matcher` directly in you ui +loop will be very slow. Implementing this logic yourself is very complex. + +The matcher is hightly optimized and can significantly outperform `fzf` and +`skim` (the `fuzzy-matcher` crate). However some of these optimizations require +a slightly less convenient API. Be sure to carefully read the documentation of +the [`Matcher`] to avoid unexpected behaviour. +# Examples + +For almost all usecases the [`pattern`] API should be used instead of calling +the matcher methods directly. [`Pattern::parse`](pattern::Pattern::parse) will +construct a single Atom (a single match operation) for each word. The pattern +can contain special characters to control what kind of match is performed (see +[`AtomKind`](crate::pattern::AtomKind)). + +``` +# use nucleo_matcher::{Matcher, Config}; +# use nucleo_matcher::pattern::{Pattern, Normalization, CaseMatching}; +let paths = ["foo/bar", "bar/foo", "foobar"]; +let mut matcher = Matcher::new(Config::DEFAULT.match_paths()); +let matches = Pattern::parse("foo bar", CaseMatching::Ignore, Normalization::Smart).match_list(paths, &mut matcher); +assert_eq!(matches, vec![("foo/bar", 168), ("bar/foo", 168), ("foobar", 140)]); +let matches = Pattern::parse("^foo bar", CaseMatching::Ignore, Normalization::Smart).match_list(paths, &mut matcher); +assert_eq!(matches, vec![("foo/bar", 168), ("foobar", 140)]); +``` + +If the pattern should be matched literally (without this special parsing) +[`Pattern::new`](pattern::Pattern::new) can be used instead. + +``` +# use nucleo_matcher::{Matcher, Config}; +# use nucleo_matcher::pattern::{Pattern, CaseMatching, AtomKind, Normalization}; +let paths = ["foo/bar", "bar/foo", "foobar"]; +let mut matcher = Matcher::new(Config::DEFAULT.match_paths()); +let matches = Pattern::new("foo bar", CaseMatching::Ignore, Normalization::Smart, AtomKind::Fuzzy).match_list(paths, &mut matcher); +assert_eq!(matches, vec![("foo/bar", 168), ("bar/foo", 168), ("foobar", 140)]); +let paths = ["^foo/bar", "bar/^foo", "foobar"]; +let matches = Pattern::new("^foo bar", CaseMatching::Ignore, Normalization::Smart, AtomKind::Fuzzy).match_list(paths, &mut matcher); +assert_eq!(matches, vec![("^foo/bar", 188), ("bar/^foo", 188)]); +``` + +Word segmentation is performed automatically on any unescaped character for which [`is_whitespace`](char::is_whitespace) returns true. +This is relevant, for instance, with non-english keyboard input. + +``` +# use nucleo_matcher::pattern::{Atom, Pattern, Normalization, CaseMatching}; +assert_eq!( + // double-width 'Ideographic Space', i.e. `'\u{3000}'` + Pattern::parse("ほげ ふが", CaseMatching::Smart, Normalization::Smart).atoms, + vec![ + Atom::parse("ほげ", CaseMatching::Smart, Normalization::Smart), + Atom::parse("ふが", CaseMatching::Smart, Normalization::Smart), + ], +); +``` + +If word segmentation is also not desired, a single `Atom` can be constructed directly. + +``` +# use nucleo_matcher::{Matcher, Config}; +# use nucleo_matcher::pattern::{Pattern, Atom, CaseMatching, Normalization, AtomKind}; +let paths = ["foobar", "foo bar"]; +let mut matcher = Matcher::new(Config::DEFAULT); +let matches = Atom::new("foo bar", CaseMatching::Ignore, Normalization::Smart, AtomKind::Fuzzy, false).match_list(paths, &mut matcher); +assert_eq!(matches, vec![("foo bar", 192)]); +``` + + +# Status + +Nucleo is used in the helix-editor and therefore has a large user base with lots or real world testing. The core matcher implementation is considered complete and is unlikely to see major changes. The `nucleo-matcher` crate is finished and ready for widespread use, breaking changes should be very rare (a 1.0 release should not be far away). + +*/ + +// sadly ranges don't optmimzie well +#![allow(clippy::manual_range_contains)] +#![warn(missing_docs)] + +pub mod chars; +mod config; +#[cfg(test)] +mod debug; +mod exact; +mod fuzzy_greedy; +mod fuzzy_optimal; +mod matrix; +pub mod pattern; +mod prefilter; +mod score; +mod utf32_str; + +#[cfg(test)] +mod tests; + +pub use crate::config::Config; +pub use crate::utf32_str::{Utf32Str, Utf32String}; + +use crate::chars::{AsciiChar, Char}; +use crate::matrix::MatrixSlab; + +/// A matcher engine that can execute (fuzzy) matches. +/// +/// A matches contains **heap allocated** scratch memory that is reused during +/// matching. This scratch memory allows the matcher to guarantee that it will +/// **never allocate** during matching (with the exception of pushing to the +/// `indices` vector if there isn't enough capacity). However this scratch +/// memory is fairly large (around 135KB) so creating a matcher is expensive. +/// +/// All `.._match` functions will not compute the indices of the matched +/// characters. These should be used to prefilter to filter and rank all +/// matches. All `.._indices` functions will also compute the indices of the +/// matched characters but are slower compared to the `..match` variant. These +/// should be used when rendering the best N matches. Note that the `indices` +/// argument is **never cleared**. This allows running multiple different +/// matches on the same haystack and merging the indices by sorting and +/// deduplicating the vector. +/// +/// The `needle` argument for each function must always be normalized by the +/// caller (unicode normalization and case folding). Otherwise, the matcher +/// may fail to produce a match. The [`pattern`] modules provides utilities +/// to preprocess needles and **should usually be preferred over invoking the +/// matcher directly**. Additionally it's recommend to perform separate matches +/// for each word in the needle. Consider the folloling example: +/// +/// If `foo bar` is used as the needle it matches both `foo test baaar` and +/// `foo hello-world bar`. However, `foo test baaar` will receive a higher +/// score than `foo hello-world bar`. `baaar` contains a 2 character gap which +/// will receive a penalty and therefore the user will likely expect it to rank +/// lower. However, if `foo bar` is matched as a single query `hello-world` and +/// `test` are both considered gaps too. As `hello-world` is a much longer gap +/// then `test` the extra penalty for `baaar` is canceled out. If both words +/// are matched individually the interspersed words do not receive a penalty and +/// `foo hello-world bar` ranks higher. +/// +/// In general nucleo is a **substring matching tool** (except for the prefix/ +/// postfix matching modes) with no penalty assigned to matches that start +/// later within the same pattern (which enables matching words individually +/// as shown above). If patterns show a large variety in length and the syntax +/// described above is not used it may be preferable to give preference to +/// matches closer to the start of a haystack. To accommodate that usecase the +/// [`prefer_prefix`](Config::prefer_prefix) option can be set to true. +/// +/// Matching is limited to 2^32-1 codepoints, if the haystack is longer than +/// that the matcher **will panic**. The caller must decide whether it wants to +/// filter out long haystacks or truncate them. +pub struct Matcher { + #[allow(missing_docs)] + pub config: Config, + slab: MatrixSlab, +} + +// this is just here for convenience not sure if we should implement this +impl Clone for Matcher { + fn clone(&self) -> Self { + Matcher { + config: self.config.clone(), + slab: MatrixSlab::new(), + } + } +} + +impl std::fmt::Debug for Matcher { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("Matcher") + .field("config", &self.config) + .finish_non_exhaustive() + } +} + +impl Default for Matcher { + fn default() -> Self { + Matcher { + config: Config::DEFAULT, + slab: MatrixSlab::new(), + } + } +} + +impl Matcher { + /// Creates a new matcher instance, note that this will eagerly allocate a + /// fairly large chunk of heap memory (around 135KB currently but subject to + /// change) so matchers should be reused if called often (like in a loop). + pub fn new(config: Config) -> Self { + Self { + config, + slab: MatrixSlab::new(), + } + } + + /// Find the fuzzy match with the highest score in the `haystack`. + /// + /// This functions has `O(mn)` time complexity for short inputs. + /// To avoid slowdowns it automatically falls back to + /// [greedy matching](crate::Matcher::fuzzy_match_greedy) for large + /// needles and haystacks. + /// + /// See the [matcher documentation](crate::Matcher) for more details. + pub fn fuzzy_match(&mut self, haystack: Utf32Str<'_>, needle: Utf32Str<'_>) -> Option<u16> { + assert!(haystack.len() <= u32::MAX as usize); + self.fuzzy_matcher_impl::<false>(haystack, needle, &mut Vec::new()) + } + + /// Find the fuzzy match with the highest score in the `haystack` and + /// compute its indices. + /// + /// This functions has `O(mn)` time complexity for short inputs. To + /// avoid slowdowns it automatically falls back to + /// [greedy matching](crate::Matcher::fuzzy_match_greedy) for large needles + /// and haystacks + /// + /// See the [matcher documentation](crate::Matcher) for more details. + pub fn fuzzy_indices( + &mut self, + haystack: Utf32Str<'_>, + needle: Utf32Str<'_>, + indices: &mut Vec<u32>, + ) -> Option<u16> { + assert!(haystack.len() <= u32::MAX as usize); + self.fuzzy_matcher_impl::<true>(haystack, needle, indices) + } + + fn fuzzy_matcher_impl<const INDICES: bool>( + &mut self, + haystack_: Utf32Str<'_>, + needle_: Utf32Str<'_>, + indices: &mut Vec<u32>, + ) -> Option<u16> { + if needle_.len() > haystack_.len() { + return None; + } + if needle_.is_empty() { + return Some(0); + } + if needle_.len() == haystack_.len() { + return self.exact_match_impl::<INDICES>( + haystack_, + needle_, + 0, + haystack_.len(), + indices, + ); + } + assert!( + haystack_.len() <= u32::MAX as usize, + "fuzzy matching is only support for up to 2^32-1 codepoints" + ); + match (haystack_, needle_) { + (Utf32Str::Ascii(haystack), Utf32Str::Ascii(needle)) => { + if let &[needle] = needle { + return self.substring_match_1_ascii::<INDICES>(haystack, needle, indices); + } + let (start, greedy_end, end) = self.prefilter_ascii(haystack, needle, false)?; + if needle_.len() == end - start { + return Some(self.calculate_score::<INDICES, _, _>( + AsciiChar::cast(haystack), + AsciiChar::cast(needle), + start, + greedy_end, + indices, + )); + } + self.fuzzy_match_optimal::<INDICES, AsciiChar, AsciiChar>( + AsciiChar::cast(haystack), + AsciiChar::cast(needle), + start, + greedy_end, + end, + indices, + ) + } + (Utf32Str::Ascii(_), Utf32Str::Unicode(_)) => { + // a purely ascii haystack can never be transformed to match + // a needle that contains non-ascii chars since we don't allow gaps + None + } + (Utf32Str::Unicode(haystack), Utf32Str::Ascii(needle)) => { + if let &[needle] = needle { + let (start, _) = self.prefilter_non_ascii(haystack, needle_, true)?; + let res = self.substring_match_1_non_ascii::<INDICES>( + haystack, + needle as char, + start, + indices, + ); + return Some(res); + } + let (start, end) = self.prefilter_non_ascii(haystack, needle_, false)?; + if needle_.len() == end - start { + return self + .exact_match_impl::<INDICES>(haystack_, needle_, start, end, indices); + } + self.fuzzy_match_optimal::<INDICES, char, AsciiChar>( + haystack, + AsciiChar::cast(needle), + start, + start + 1, + end, + indices, + ) + } + (Utf32Str::Unicode(haystack), Utf32Str::Unicode(needle)) => { + if let &[needle] = needle { + let (start, _) = self.prefilter_non_ascii(haystack, needle_, true)?; + let res = self + .substring_match_1_non_ascii::<INDICES>(haystack, needle, start, indices); + return Some(res); + } + let (start, end) = self.prefilter_non_ascii(haystack, needle_, false)?; + if needle_.len() == end - start { + return self + .exact_match_impl::<INDICES>(haystack_, needle_, start, end, indices); + } + self.fuzzy_match_optimal::<INDICES, char, char>( + haystack, + needle, + start, + start + 1, + end, + indices, + ) + } + } + } + + /// Greedly find a fuzzy match in the `haystack`. + /// + /// This functions has `O(n)` time complexity but may provide unintutive (non-optimal) + /// indices and scores. Usually [fuzzy_match](crate::Matcher::fuzzy_match) should + /// be preferred. + /// + /// See the [matcher documentation](crate::Matcher) for more details. + pub fn fuzzy_match_greedy( + &mut self, + haystack: Utf32Str<'_>, + needle: Utf32Str<'_>, + ) -> Option<u16> { + assert!(haystack.len() <= u32::MAX as usize); + self.fuzzy_match_greedy_impl::<false>(haystack, needle, &mut Vec::new()) + } + + /// Greedly find a fuzzy match in the `haystack` and compute its indices. + /// + /// This functions has `O(n)` time complexity but may provide unintuitive (non-optimal) + /// indices and scores. Usually [fuzzy_indices](crate::Matcher::fuzzy_indices) should + /// be preferred. + /// + /// See the [matcher documentation](crate::Matcher) for more details. + pub fn fuzzy_indices_greedy( + &mut self, + haystack: Utf32Str<'_>, + needle: Utf32Str<'_>, + indices: &mut Vec<u32>, + ) -> Option<u16> { + assert!(haystack.len() <= u32::MAX as usize); + self.fuzzy_match_greedy_impl::<true>(haystack, needle, indices) + } + + fn fuzzy_match_greedy_impl<const INDICES: bool>( + &mut self, + haystack: Utf32Str<'_>, + needle_: Utf32Str<'_>, + indices: &mut Vec<u32>, + ) -> Option<u16> { + if needle_.len() > haystack.len() { + return None; + } + if needle_.is_empty() { + return Some(0); + } + if needle_.len() == haystack.len() { + return self.exact_match_impl::<INDICES>(haystack, needle_, 0, haystack.len(), indices); + } + assert!( + haystack.len() <= u32::MAX as usize, + "matching is only support for up to 2^32-1 codepoints" + ); + match (haystack, needle_) { + (Utf32Str::Ascii(haystack), Utf32Str::Ascii(needle)) => { + let (start, greedy_end, _) = self.prefilter_ascii(haystack, needle, true)?; + if needle_.len() == greedy_end - start { + return Some(self.calculate_score::<INDICES, _, _>( + AsciiChar::cast(haystack), + AsciiChar::cast(needle), + start, + greedy_end, + indices, + )); + } + self.fuzzy_match_greedy_::<INDICES, AsciiChar, AsciiChar>( + AsciiChar::cast(haystack), + AsciiChar::cast(needle), + start, + greedy_end, + indices, + ) + } + (Utf32Str::Ascii(_), Utf32Str::Unicode(_)) => { + // a purely ascii haystack can never be transformed to match + // a needle that contains non-ascii chars since we don't allow gaps + None + } + (Utf32Str::Unicode(haystack), Utf32Str::Ascii(needle)) => { + let (start, _) = self.prefilter_non_ascii(haystack, needle_, true)?; + self.fuzzy_match_greedy_::<INDICES, char, AsciiChar>( + haystack, + AsciiChar::cast(needle), + start, + start + 1, + indices, + ) + } + (Utf32Str::Unicode(haystack), Utf32Str::Unicode(needle)) => { + let (start, _) = self.prefilter_non_ascii(haystack, needle_, true)?; + self.fuzzy_match_greedy_::<INDICES, char, char>( + haystack, + needle, + start, + start + 1, + indices, + ) + } + } + } + + /// Finds the substring match with the highest score in the `haystack`. + /// + /// This functions has `O(nm)` time complexity. However many cases can + /// be significantly accelerated using prefilters so it's usually very fast + /// in practice. + /// + /// See the [matcher documentation](crate::Matcher) for more details. + pub fn substring_match( + &mut self, + haystack: Utf32Str<'_>, + needle_: Utf32Str<'_>, + ) -> Option<u16> { + self.substring_match_impl::<false>(haystack, needle_, &mut Vec::new()) + } + + /// Finds the substring match with the highest score in the `haystack` and + /// compute its indices. + /// + /// This functions has `O(nm)` time complexity. However many cases can + /// be significantly accelerated using prefilters so it's usually fast + /// in practice. + /// + /// See the [matcher documentation](crate::Matcher) for more details. + pub fn substring_indices( + &mut self, + haystack: Utf32Str<'_>, + needle_: Utf32Str<'_>, + indices: &mut Vec<u32>, + ) -> Option<u16> { + self.substring_match_impl::<true>(haystack, needle_, indices) + } + + fn substring_match_impl<const INDICES: bool>( + &mut self, + haystack: Utf32Str<'_>, + needle_: Utf32Str<'_>, + indices: &mut Vec<u32>, + ) -> Option<u16> { + if needle_.len() > haystack.len() { + return None; + } + if needle_.is_empty() { + return Some(0); + } + if needle_.len() == haystack.len() { + return self.exact_match_impl::<INDICES>(haystack, needle_, 0, haystack.len(), indices); + } + assert!( + haystack.len() <= u32::MAX as usize, + "matching is only support for up to 2^32-1 codepoints" + ); + match (haystack, needle_) { + (Utf32Str::Ascii(haystack), Utf32Str::Ascii(needle)) => { + if let &[needle] = needle { + return self.substring_match_1_ascii::<INDICES>(haystack, needle, indices); + } + self.substring_match_ascii::<INDICES>(haystack, needle, indices) + } + (Utf32Str::Ascii(_), Utf32Str::Unicode(_)) => { + // a purely ascii haystack can never be transformed to match + // a needle that contains non-ascii chars since we don't allow gaps + None + } + (Utf32Str::Unicode(haystack), Utf32Str::Ascii(needle)) => { + if let &[needle] = needle { + let (start, _) = self.prefilter_non_ascii(haystack, needle_, true)?; + let res = self.substring_match_1_non_ascii::<INDICES>( + haystack, + needle as char, + start, + indices, + ); + return Some(res); + } + let (start, _) = self.prefilter_non_ascii(haystack, needle_, false)?; + self.substring_match_non_ascii::<INDICES, _>( + haystack, + AsciiChar::cast(needle), + start, + indices, + ) + } + (Utf32Str::Unicode(haystack), Utf32Str::Unicode(needle)) => { + if let &[needle] = needle { + let (start, _) = self.prefilter_non_ascii(haystack, needle_, true)?; + let res = self + .substring_match_1_non_ascii::<INDICES>(haystack, needle, start, indices); + return Some(res); + } + let (start, _) = self.prefilter_non_ascii(haystack, needle_, false)?; + self.substring_match_non_ascii::<INDICES, _>(haystack, needle, start, indices) + } + } + } + + /// Checks whether needle and haystack match exactly. + /// + /// This functions has `O(n)` time complexity. + /// + /// See the [matcher documentation](crate::Matcher) for more details. + pub fn exact_match(&mut self, haystack: Utf32Str<'_>, needle: Utf32Str<'_>) -> Option<u16> { + if needle.is_empty() { + return Some(0); + } + let mut leading_space = 0; + let mut trailing_space = 0; + if !needle.first().is_whitespace() { + leading_space = haystack.leading_white_space() + } + if !needle.last().is_whitespace() { + trailing_space = haystack.trailing_white_space() + } + // avoid wraparound in size check + if trailing_space == haystack.len() { + return None; + } + self.exact_match_impl::<false>( + haystack, + needle, + leading_space, + haystack.len() - trailing_space, + &mut Vec::new(), + ) + } + + /// Checks whether needle and haystack match exactly and compute the matches indices. + /// + /// This functions has `O(n)` time complexity. + /// + /// See the [matcher documentation](crate::Matcher) for more details. + pub fn exact_indices( + &mut self, + haystack: Utf32Str<'_>, + needle: Utf32Str<'_>, + indices: &mut Vec<u32>, + ) -> Option<u16> { + if needle.is_empty() { + return Some(0); + } + let mut leading_space = 0; + let mut trailing_space = 0; + if !needle.first().is_whitespace() { + leading_space = haystack.leading_white_space() + } + if !needle.last().is_whitespace() { + trailing_space = haystack.trailing_white_space() + } + // avoid wraparound in size check + if trailing_space == haystack.len() { + return None; + } + self.exact_match_impl::<true>( + haystack, + needle, + leading_space, + haystack.len() - trailing_space, + indices, + ) + } + + /// Checks whether needle is a prefix of the haystack. + /// + /// This functions has `O(n)` time complexity. + /// + /// See the [matcher documentation](crate::Matcher) for more details. + pub fn prefix_match(&mut self, haystack: Utf32Str<'_>, needle: Utf32Str<'_>) -> Option<u16> { + if needle.is_empty() { + return Some(0); + } + let mut leading_space = 0; + if !needle.first().is_whitespace() { + leading_space = haystack.leading_white_space() + } + if haystack.len() - leading_space < needle.len() { + None + } else { + self.exact_match_impl::<false>( + haystack, + needle, + leading_space, + needle.len() + leading_space, + &mut Vec::new(), + ) + } + } + + /// Checks whether needle is a prefix of the haystack and compute the matches indices. + /// + /// This functions has `O(n)` time complexity. + /// + /// See the [matcher documentation](crate::Matcher) for more details. + pub fn prefix_indices( + &mut self, + haystack: Utf32Str<'_>, + needle: Utf32Str<'_>, + indices: &mut Vec<u32>, + ) -> Option<u16> { + if needle.is_empty() { + return Some(0); + } + let mut leading_space = 0; + if !needle.first().is_whitespace() { + leading_space = haystack.leading_white_space() + } + if haystack.len() - leading_space < needle.len() { + None + } else { + self.exact_match_impl::<true>( + haystack, + needle, + leading_space, + needle.len() + leading_space, + indices, + ) + } + } + + /// Checks whether needle is a postfix of the haystack. + /// + /// This functions has `O(n)` time complexity. + /// + /// See the [matcher documentation](crate::Matcher) for more details. + pub fn postfix_match(&mut self, haystack: Utf32Str<'_>, needle: Utf32Str<'_>) -> Option<u16> { + if needle.is_empty() { + return Some(0); + } + let mut trailing_spaces = 0; + if !needle.last().is_whitespace() { + trailing_spaces = haystack.trailing_white_space() + } + if haystack.len() - trailing_spaces < needle.len() { + None + } else { + self.exact_match_impl::<false>( + haystack, + needle, + haystack.len() - needle.len() - trailing_spaces, + haystack.len() - trailing_spaces, + &mut Vec::new(), + ) + } + } + + /// Checks whether needle is a postfix of the haystack and compute the matches indices. + /// + /// This functions has `O(n)` time complexity. + /// + /// See the [matcher documentation](crate::Matcher) for more details. + pub fn postfix_indices( + &mut self, + haystack: Utf32Str<'_>, + needle: Utf32Str<'_>, + indices: &mut Vec<u32>, + ) -> Option<u16> { + if needle.is_empty() { + return Some(0); + } + let mut trailing_spaces = 0; + if !needle.last().is_whitespace() { + trailing_spaces = haystack.trailing_white_space() + } + if haystack.len() - trailing_spaces < needle.len() { + None + } else { + self.exact_match_impl::<true>( + haystack, + needle, + haystack.len() - needle.len() - trailing_spaces, + haystack.len() - trailing_spaces, + indices, + ) + } + } + + fn exact_match_impl<const INDICES: bool>( + &mut self, + haystack: Utf32Str<'_>, + needle_: Utf32Str<'_>, + start: usize, + end: usize, + indices: &mut Vec<u32>, + ) -> Option<u16> { + if needle_.len() != end - start { + return None; + } + assert!( + haystack.len() <= u32::MAX as usize, + "matching is only support for up to 2^32-1 codepoints" + ); + let score = match (haystack, needle_) { + (Utf32Str::Ascii(haystack), Utf32Str::Ascii(needle)) => { + let matched = if self.config.ignore_case { + AsciiChar::cast(haystack)[start..end] + .iter() + .map(|c| c.normalize(&self.config)) + .eq(AsciiChar::cast(needle) + .iter() + .map(|c| c.normalize(&self.config))) + } else { + &haystack[start..end] == needle + }; + if !matched { + return None; + } + self.calculate_score::<INDICES, _, _>( + AsciiChar::cast(haystack), + AsciiChar::cast(needle), + start, + end, + indices, + ) + } + (Utf32Str::Ascii(_), Utf32Str::Unicode(_)) => { + // a purely ascii haystack can never be transformed to match + // a needle that contains non-ascii chars since we don't allow gaps + return None; + } + (Utf32Str::Unicode(haystack), Utf32Str::Ascii(needle)) => { + let matched = haystack[start..end] + .iter() + .map(|c| c.normalize(&self.config)) + .eq(AsciiChar::cast(needle) + .iter() + .map(|c| c.normalize(&self.config))); + if !matched { + return None; + } + + self.calculate_score::<INDICES, _, _>( + haystack, + AsciiChar::cast(needle), + start, + end, + indices, + ) + } + (Utf32Str::Unicode(haystack), Utf32Str::Unicode(needle)) => { + let matched = haystack[start..end] + .iter() + .map(|c| c.normalize(&self.config)) + .eq(needle.iter().map(|c| c.normalize(&self.config))); + if !matched { + return None; + } + self.calculate_score::<INDICES, _, _>(haystack, needle, start, end, indices) + } + }; + Some(score) + } +} diff --git a/crates/atuin-nucleo/matcher/src/matrix.rs b/crates/atuin-nucleo/matcher/src/matrix.rs new file mode 100644 index 00000000..a91ed95f --- /dev/null +++ b/crates/atuin-nucleo/matcher/src/matrix.rs @@ -0,0 +1,198 @@ +use std::alloc::{alloc_zeroed, dealloc, handle_alloc_error, Layout}; +use std::marker::PhantomData; +use std::mem::size_of; +use std::panic::{RefUnwindSafe, UnwindSafe}; +use std::ptr::{slice_from_raw_parts_mut, NonNull}; + +use crate::chars::Char; + +const MAX_MATRIX_SIZE: usize = 100 * 1024; // 100*1024 = 100KB + +// these two aren't hard maxima, instead we simply allow whatever will fit into memory +const MAX_HAYSTACK_LEN: usize = 2048; // 64KB +const MAX_NEEDLE_LEN: usize = 2048; // 64KB + +struct MatrixLayout<C: Char> { + haystack_len: usize, + needle_len: usize, + layout: Layout, + haystack_off: usize, + bonus_off: usize, + rows_off: usize, + score_off: usize, + matrix_off: usize, + _phantom: PhantomData<C>, +} +impl<C: Char> MatrixLayout<C> { + fn new(haystack_len: usize, needle_len: usize) -> MatrixLayout<C> { + assert!(haystack_len >= needle_len); + assert!(haystack_len <= u32::MAX as usize); + let mut layout = Layout::from_size_align(0, 1).unwrap(); + let haystack_layout = Layout::array::<C>(haystack_len).unwrap(); + let bonus_layout = Layout::array::<u8>(haystack_len).unwrap(); + let rows_layout = Layout::array::<u16>(needle_len).unwrap(); + let score_layout = Layout::array::<ScoreCell>(haystack_len + 1 - needle_len).unwrap(); + let matrix_layout = + Layout::array::<MatrixCell>((haystack_len + 1 - needle_len) * needle_len).unwrap(); + + let haystack_off; + (layout, haystack_off) = layout.extend(haystack_layout).unwrap(); + let bonus_off; + (layout, bonus_off) = layout.extend(bonus_layout).unwrap(); + let rows_off; + (layout, rows_off) = layout.extend(rows_layout).unwrap(); + let score_off; + (layout, score_off) = layout.extend(score_layout).unwrap(); + let matrix_off; + (layout, matrix_off) = layout.extend(matrix_layout).unwrap(); + MatrixLayout { + haystack_len, + needle_len, + layout, + haystack_off, + bonus_off, + rows_off, + score_off, + matrix_off, + _phantom: PhantomData, + } + } + /// # Safety + /// + /// `ptr` must point at an allocated with MARTIX_ALLOC_LAYOUT + #[allow(clippy::type_complexity)] + unsafe fn fieds_from_ptr( + &self, + ptr: NonNull<u8>, + ) -> ( + *mut [C], + *mut [u8], + *mut [u16], + *mut [ScoreCell], + *mut [MatrixCell], + ) { + let base = ptr.as_ptr(); + let haystack = base.add(self.haystack_off) as *mut C; + let haystack = slice_from_raw_parts_mut(haystack, self.haystack_len); + let bonus = base.add(self.bonus_off); + let bonus = slice_from_raw_parts_mut(bonus, self.haystack_len); + let rows = base.add(self.rows_off) as *mut u16; + let rows = slice_from_raw_parts_mut(rows, self.needle_len); + let cells = base.add(self.score_off) as *mut ScoreCell; + let cells = slice_from_raw_parts_mut(cells, self.haystack_len + 1 - self.needle_len); + let matrix = base.add(self.matrix_off) as *mut MatrixCell; + let matrix = slice_from_raw_parts_mut( + matrix, + (self.haystack_len + 1 - self.needle_len) * self.haystack_len, + ); + (haystack, bonus, rows, cells, matrix) + } +} + +const _SIZE_CHECK: () = { + if size_of::<ScoreCell>() != 8 { + panic!() + } +}; + +// make this act like a u64 +#[repr(align(8))] +#[derive(Clone, Copy, PartialEq, Eq)] +pub(crate) struct ScoreCell { + pub score: u16, + pub consecutive_bonus: u8, + pub matched: bool, +} + +pub(crate) struct MatcherDataView<'a, C: Char> { + pub haystack: &'a mut [C], + // stored as a separate array instead of struct + // to avoid padding since char is too large and u8 too small :/ + pub bonus: &'a mut [u8], + pub current_row: &'a mut [ScoreCell], + pub row_offs: &'a mut [u16], + pub matrix_cells: &'a mut [MatrixCell], +} +#[repr(transparent)] +pub struct MatrixCell(pub(crate) u8); + +impl MatrixCell { + pub fn set(&mut self, p_match: bool, m_match: bool) { + self.0 = p_match as u8 | ((m_match as u8) << 1); + } + + pub fn get(&self, m_matrix: bool) -> bool { + let mask = m_matrix as u8 + 1; + (self.0 & mask) != 0 + } +} + +// we only use this to construct the layout for the slab allocation +#[allow(unused)] +struct MatcherData { + haystack: [char; MAX_HAYSTACK_LEN], + bonus: [u8; MAX_HAYSTACK_LEN], + row_offs: [u16; MAX_NEEDLE_LEN], + scratch_space: [ScoreCell; MAX_HAYSTACK_LEN], + matrix: [u8; MAX_MATRIX_SIZE], +} + +pub(crate) struct MatrixSlab(NonNull<u8>); +unsafe impl Sync for MatrixSlab {} +unsafe impl Send for MatrixSlab {} +impl UnwindSafe for MatrixSlab {} +impl RefUnwindSafe for MatrixSlab {} + +impl MatrixSlab { + pub fn new() -> Self { + let layout = Layout::new::<MatcherData>(); + // safety: the matrix is never zero sized (hardcoded constants) + let ptr = unsafe { alloc_zeroed(layout) }; + let Some(ptr) = NonNull::new(ptr) else { + handle_alloc_error(layout) + }; + MatrixSlab(ptr.cast()) + } + + pub(crate) fn alloc<C: Char>( + &mut self, + haystack_: &[C], + needle_len: usize, + ) -> Option<MatcherDataView<'_, C>> { + let cells = haystack_.len() * needle_len; + if cells > MAX_MATRIX_SIZE + || haystack_.len() > u16::MAX as usize + // ensures that scores never overflow + || needle_len > MAX_NEEDLE_LEN + { + return None; + } + let matrix_layout = MatrixLayout::<C>::new(haystack_.len(), needle_len); + if matrix_layout.layout.size() > size_of::<MatcherData>() { + return None; + } + unsafe { + // safely: this allocation is valid for MATRIX_ALLOC_LAYOUT + let (haystack, bonus, rows, current_row, matrix_cells) = + matrix_layout.fieds_from_ptr(self.0); + // copy haystack before creating references to ensure we don't create + // references to invalid chars (which may or may not be UB) + haystack_ + .as_ptr() + .copy_to_nonoverlapping(haystack as *mut _, haystack_.len()); + Some(MatcherDataView { + haystack: &mut *haystack, + row_offs: &mut *rows, + bonus: &mut *bonus, + current_row: &mut *current_row, + matrix_cells: &mut *matrix_cells, + }) + } + } +} + +impl Drop for MatrixSlab { + fn drop(&mut self) { + unsafe { dealloc(self.0.as_ptr(), Layout::new::<MatcherData>()) }; + } +} diff --git a/crates/atuin-nucleo/matcher/src/pattern.rs b/crates/atuin-nucleo/matcher/src/pattern.rs new file mode 100644 index 00000000..495feede --- /dev/null +++ b/crates/atuin-nucleo/matcher/src/pattern.rs @@ -0,0 +1,566 @@ +//! This module provides a slightly higher level API for matching strings. + +use std::cmp::Reverse; + +use crate::{chars, Matcher, Utf32Str}; + +#[cfg(test)] +mod tests; + +use crate::Utf32String; + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)] +#[non_exhaustive] +/// How to treat a case mismatch between two characters. +pub enum CaseMatching { + /// Characters never match their case folded version (`a != A`). + #[cfg_attr(not(feature = "unicode-casefold"), default)] + Respect, + /// Characters always match their case folded version (`a == A`). + #[cfg(feature = "unicode-casefold")] + Ignore, + /// Acts like [`Ignore`](CaseMatching::Ignore) if all characters in a pattern atom are + /// lowercase and like [`Respect`](CaseMatching::Respect) otherwise. + #[default] + #[cfg(feature = "unicode-casefold")] + Smart, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)] +#[non_exhaustive] +/// How to handle unicode normalization, +pub enum Normalization { + /// Characters never match their normalized version (`a != ä`). + #[cfg_attr(not(feature = "unicode-normalization"), default)] + Never, + /// Acts like [`Never`](Normalization::Never) if any character in a pattern atom + /// would need to be normalized. Otherwise normalization occurs (`a == ä` but `ä != a`). + #[default] + #[cfg(feature = "unicode-normalization")] + Smart, +} + +#[derive(Debug, PartialEq, Eq, Clone, Copy)] +#[non_exhaustive] +/// The kind of matching algorithm to run for an atom. +pub enum AtomKind { + /// Fuzzy matching where the needle must match any haystack characters + /// (match can contain gaps). This atom kind is used by default if no + /// special syntax is used. There is no negated fuzzy matching (too + /// many false positives). + /// + /// See also [`Matcher::fuzzy_match`](crate::Matcher::fuzzy_match). + Fuzzy, + /// The needle must match a contiguous sequence of haystack characters + /// without gaps. This atom kind is parsed from the following syntax: + /// `'foo` and `!foo` (negated). + /// + /// See also [`Matcher::substring_match`](crate::Matcher::substring_match). + Substring, + /// The needle must match all leading haystack characters without gaps or + /// prefix. This atom kind is parsed from the following syntax: `^foo` and + /// `!^foo` (negated). + /// + /// See also [`Matcher::prefix_match`](crate::Matcher::prefix_match). + Prefix, + /// The needle must match all trailing haystack characters without gaps or + /// postfix. This atom kind is parsed from the following syntax: `foo$` and + /// `!foo$` (negated). + /// + /// See also [`Matcher::postfix_match`](crate::Matcher::postfix_match). + Postfix, + /// The needle must match all haystack characters without gaps or prefix. + /// This atom kind is parsed from the following syntax: `^foo$` and `!^foo$` + /// (negated). + /// + /// See also [`Matcher::exact_match`](crate::Matcher::exact_match). + Exact, +} + +/// A single pattern component that is matched with a single [`Matcher`] function +#[derive(Debug, PartialEq, Eq, Clone)] +pub struct Atom { + /// Whether this pattern atom is a negative match. + /// A negative pattern atom will prevent haystacks matching it from + /// being matchend. It does not contribute to scoring/indices + pub negative: bool, + /// The kind of match that this pattern performs + pub kind: AtomKind, + needle: Utf32String, + ignore_case: bool, + normalize: bool, +} + +impl Atom { + /// Creates a single [`Atom`] from a string by performing unicode + /// normalization and case folding (if necessary). Optionally `\ ` can + /// be escaped to ` `. + pub fn new( + needle: &str, + case: CaseMatching, + normalize: Normalization, + kind: AtomKind, + escape_whitespace: bool, + ) -> Atom { + Atom::new_inner(needle, case, normalize, kind, escape_whitespace, false) + } + + fn new_inner( + needle: &str, + case: CaseMatching, + normalization: Normalization, + kind: AtomKind, + escape_whitespace: bool, + append_dollar: bool, + ) -> Atom { + let mut ignore_case; + let mut normalize; + #[cfg(feature = "unicode-normalization")] + { + normalize = matches!(normalization, Normalization::Smart); + } + #[cfg(not(feature = "unicode-normalization"))] + { + normalize = false; + } + let needle = if needle.is_ascii() { + let mut needle = if escape_whitespace { + if let Some((start, rem)) = needle.split_once("\\ ") { + let mut needle = start.to_owned(); + for rem in rem.split("\\ ") { + needle.push(' '); + needle.push_str(rem); + } + needle + } else { + needle.to_owned() + } + } else { + needle.to_owned() + }; + + match case { + #[cfg(feature = "unicode-casefold")] + CaseMatching::Ignore => { + ignore_case = true; + needle.make_ascii_lowercase() + } + #[cfg(feature = "unicode-casefold")] + CaseMatching::Smart => { + ignore_case = !needle.bytes().any(|b| b.is_ascii_uppercase()) + } + CaseMatching::Respect => ignore_case = false, + } + if append_dollar { + needle.push('$'); + } + Utf32String::Ascii(needle.into_boxed_str()) + } else { + let mut needle_ = Vec::with_capacity(needle.len()); + #[cfg(feature = "unicode-casefold")] + { + ignore_case = matches!(case, CaseMatching::Ignore | CaseMatching::Smart); + } + #[cfg(not(feature = "unicode-casefold"))] + { + ignore_case = false; + } + #[cfg(feature = "unicode-normalization")] + { + normalize = matches!(normalization, Normalization::Smart); + } + if escape_whitespace { + let mut saw_backslash = false; + for mut c in chars::graphemes(needle) { + if saw_backslash { + if c == ' ' { + needle_.push(' '); + saw_backslash = false; + continue; + } else { + needle_.push('\\'); + } + } + saw_backslash = c == '\\'; + match case { + #[cfg(feature = "unicode-casefold")] + CaseMatching::Ignore => c = chars::to_lower_case(c), + #[cfg(feature = "unicode-casefold")] + CaseMatching::Smart => { + ignore_case = ignore_case && !chars::is_upper_case(c) + } + CaseMatching::Respect => (), + } + match normalization { + #[cfg(feature = "unicode-normalization")] + Normalization::Smart => { + normalize = normalize && chars::normalize(c) == c; + } + Normalization::Never => (), + } + needle_.push(c); + } + } else { + let chars = chars::graphemes(needle).map(|mut c| { + match case { + #[cfg(feature = "unicode-casefold")] + CaseMatching::Ignore => c = chars::to_lower_case(c), + #[cfg(feature = "unicode-casefold")] + CaseMatching::Smart => { + ignore_case = ignore_case && !chars::is_upper_case(c); + } + CaseMatching::Respect => (), + } + match normalization { + #[cfg(feature = "unicode-normalization")] + Normalization::Smart => { + normalize = normalize && chars::normalize(c) == c; + } + Normalization::Never => (), + } + c + }); + needle_.extend(chars); + }; + if append_dollar { + needle_.push('$'); + } + Utf32String::Unicode(needle_.into_boxed_slice()) + }; + Atom { + kind, + needle, + negative: false, + ignore_case, + normalize, + } + } + + /// Parse a pattern atom from a string. Some special trailing and leading + /// characters can be used to control the atom kind. See [`AtomKind`] for + /// details. + pub fn parse(raw: &str, case: CaseMatching, normalize: Normalization) -> Atom { + let mut atom = raw; + let invert = match atom.as_bytes() { + [b'!', ..] => { + atom = &atom[1..]; + true + } + [b'\\', b'!', ..] => { + atom = &atom[1..]; + false + } + _ => false, + }; + + let mut kind = match atom.as_bytes() { + [b'^', ..] => { + atom = &atom[1..]; + AtomKind::Prefix + } + [b'\'', ..] => { + atom = &atom[1..]; + AtomKind::Substring + } + [b'\\', b'^' | b'\'', ..] => { + atom = &atom[1..]; + AtomKind::Fuzzy + } + _ => AtomKind::Fuzzy, + }; + + let mut append_dollar = false; + match atom.as_bytes() { + [.., b'\\', b'$'] => { + append_dollar = true; + atom = &atom[..atom.len() - 2] + } + [.., b'$'] => { + kind = if kind == AtomKind::Fuzzy { + AtomKind::Postfix + } else { + AtomKind::Exact + }; + atom = &atom[..atom.len() - 1] + } + _ => (), + } + + if invert && kind == AtomKind::Fuzzy { + kind = AtomKind::Substring + } + + let mut pattern = Atom::new_inner(atom, case, normalize, kind, true, append_dollar); + pattern.negative = invert; + pattern + } + + /// Matches this pattern against `haystack` (using the allocation and configuration + /// from `matcher`) and calculates a ranking score. See the [`Matcher`]. + /// Documentation for more details. + /// + /// *Note:* The `ignore_case` setting is overwritten to match the casing of + /// each pattern atom. + pub fn score(&self, haystack: Utf32Str<'_>, matcher: &mut Matcher) -> Option<u16> { + matcher.config.ignore_case = self.ignore_case; + matcher.config.normalize = self.normalize; + let pattern_score = match self.kind { + AtomKind::Exact => matcher.exact_match(haystack, self.needle.slice(..)), + AtomKind::Fuzzy => matcher.fuzzy_match(haystack, self.needle.slice(..)), + AtomKind::Substring => matcher.substring_match(haystack, self.needle.slice(..)), + AtomKind::Prefix => matcher.prefix_match(haystack, self.needle.slice(..)), + AtomKind::Postfix => matcher.postfix_match(haystack, self.needle.slice(..)), + }; + if self.negative { + if pattern_score.is_some() { + return None; + } + Some(0) + } else { + pattern_score + } + } + + /// Matches this pattern against `haystack` (using the allocation and + /// configuration from `matcher`), calculates a ranking score and the match + /// indices. See the [`Matcher`]. Documentation for more + /// details. + /// + /// *Note:* The `ignore_case` setting is overwritten to match the casing of + /// each pattern atom. + /// + /// *Note:* The `indices` vector is not cleared by this function. + pub fn indices( + &self, + haystack: Utf32Str<'_>, + matcher: &mut Matcher, + indices: &mut Vec<u32>, + ) -> Option<u16> { + matcher.config.ignore_case = self.ignore_case; + matcher.config.normalize = self.normalize; + if self.negative { + let pattern_score = match self.kind { + AtomKind::Exact => matcher.exact_match(haystack, self.needle.slice(..)), + AtomKind::Fuzzy => matcher.fuzzy_match(haystack, self.needle.slice(..)), + AtomKind::Substring => matcher.substring_match(haystack, self.needle.slice(..)), + AtomKind::Prefix => matcher.prefix_match(haystack, self.needle.slice(..)), + AtomKind::Postfix => matcher.postfix_match(haystack, self.needle.slice(..)), + }; + pattern_score.is_none().then_some(0) + } else { + match self.kind { + AtomKind::Exact => matcher.exact_indices(haystack, self.needle.slice(..), indices), + AtomKind::Fuzzy => matcher.fuzzy_indices(haystack, self.needle.slice(..), indices), + AtomKind::Substring => { + matcher.substring_indices(haystack, self.needle.slice(..), indices) + } + AtomKind::Prefix => { + matcher.prefix_indices(haystack, self.needle.slice(..), indices) + } + AtomKind::Postfix => { + matcher.postfix_indices(haystack, self.needle.slice(..), indices) + } + } + } + } + + /// Returns the needle text that is passed to the matcher. All indices + /// produced by the `indices` functions produce char indices used to index + /// this text + pub fn needle_text(&self) -> Utf32Str<'_> { + self.needle.slice(..) + } + /// Convenience function to easily match (and sort) a (relatively small) + /// list of inputs. + /// + /// *Note* This function is not recommended for building a full fuzzy + /// matching application that can match large numbers of matches (like all + /// files in a directory) as all matching is done on the current thread, + /// effectively blocking the UI. For such applications the high level + /// `nucleo` crate can be used instead. + pub fn match_list<T: AsRef<str>>( + &self, + items: impl IntoIterator<Item = T>, + matcher: &mut Matcher, + ) -> Vec<(T, u16)> { + if self.needle.is_empty() { + return items.into_iter().map(|item| (item, 0)).collect(); + } + let mut buf = Vec::new(); + let mut items: Vec<_> = items + .into_iter() + .filter_map(|item| { + self.score(Utf32Str::new(item.as_ref(), &mut buf), matcher) + .map(|score| (item, score)) + }) + .collect(); + items.sort_by_key(|(_, score)| Reverse(*score)); + items + } +} + +fn pattern_atoms(pattern: &str) -> impl Iterator<Item = &str> + '_ { + let mut saw_backslash = false; + pattern.split(move |c| { + saw_backslash = match c { + c if c.is_whitespace() && !saw_backslash => return true, + '\\' => true, + _ => false, + }; + false + }) +} + +#[derive(Debug, Default)] +/// A text pattern made up of (potentially multiple) [atoms](crate::pattern::Atom). +#[non_exhaustive] +pub struct Pattern { + /// The individual pattern (words) in this pattern + pub atoms: Vec<Atom>, +} + +impl Pattern { + /// Creates a pattern where each word is matched individually (whitespaces + /// can be escaped with `\`). Otherwise no parsing is performed (so `$`, `!`, + /// `'` and `^` don't receive special treatment). If you want to match the entire + /// pattern as a single needle use a single [`Atom`] instead. + pub fn new( + pattern: &str, + case_matching: CaseMatching, + normalize: Normalization, + kind: AtomKind, + ) -> Pattern { + let atoms = pattern_atoms(pattern) + .filter_map(|pat| { + let pat = Atom::new(pat, case_matching, normalize, kind, true); + (!pat.needle.is_empty()).then_some(pat) + }) + .collect(); + Pattern { atoms } + } + /// Creates a pattern where each word is matched individually (whitespaces + /// can be escaped with `\`). And `$`, `!`, `'` and `^` at word boundaries will + /// cause different matching behaviour (see [`AtomKind`]). These can be + /// escaped with backslash. + pub fn parse(pattern: &str, case_matching: CaseMatching, normalize: Normalization) -> Pattern { + let atoms = pattern_atoms(pattern) + .filter_map(|pat| { + let pat = Atom::parse(pat, case_matching, normalize); + (!pat.needle.is_empty()).then_some(pat) + }) + .collect(); + Pattern { atoms } + } + + /// Convenience function to easily match (and sort) a (relatively small) + /// list of inputs. + /// + /// *Note* This function is not recommended for building a full fuzzy + /// matching application that can match large numbers of matches (like all + /// files in a directory) as all matching is done on the current thread, + /// effectively blocking the UI. For such applications the high level + /// `nucleo` crate can be used instead. + pub fn match_list<T: AsRef<str>>( + &self, + items: impl IntoIterator<Item = T>, + matcher: &mut Matcher, + ) -> Vec<(T, u32)> { + if self.atoms.is_empty() { + return items.into_iter().map(|item| (item, 0)).collect(); + } + let mut buf = Vec::new(); + let mut items: Vec<_> = items + .into_iter() + .filter_map(|item| { + self.score(Utf32Str::new(item.as_ref(), &mut buf), matcher) + .map(|score| (item, score)) + }) + .collect(); + items.sort_by_key(|(_, score)| Reverse(*score)); + items + } + + /// Matches this pattern against `haystack` (using the allocation and configuration + /// from `matcher`) and calculates a ranking score. See the [`Matcher`] + /// documentation for more details. + /// + /// *Note:* The `ignore_case` setting is overwritten to match the casing of + /// each pattern atom. + pub fn score(&self, haystack: Utf32Str<'_>, matcher: &mut Matcher) -> Option<u32> { + if self.atoms.is_empty() { + return Some(0); + } + let mut score = 0; + for pattern in &self.atoms { + score += pattern.score(haystack, matcher)? as u32; + } + Some(score) + } + + /// Matches this pattern against `haystack` (using the allocation and + /// configuration from `matcher`), calculates a ranking score and the match + /// indices. See the [`Matcher`] documentation for more + /// details. + /// + /// *Note:* The `ignore_case` setting is overwritten to match the casing of + /// each pattern atom. + /// + /// *Note:* The indices for each pattern are calculated individually + /// and simply appended to the `indices` vector and not deduplicated/sorted. + /// This allows associating the match indices to their source pattern. If + /// required (like for highlighting) unique/sorted indices can be obtained + /// as follows: + /// + /// ``` + /// # let mut indices: Vec<u32> = Vec::new(); + /// indices.sort_unstable(); + /// indices.dedup(); + /// ``` + pub fn indices( + &self, + haystack: Utf32Str<'_>, + matcher: &mut Matcher, + indices: &mut Vec<u32>, + ) -> Option<u32> { + if self.atoms.is_empty() { + return Some(0); + } + let mut score = 0; + for pattern in &self.atoms { + score += pattern.indices(haystack, matcher, indices)? as u32; + } + Some(score) + } + + /// Refreshes this pattern by reparsing it from a string. This is mostly + /// equivalent to just constructing a new pattern using [`Pattern::parse`] + /// but is slightly more efficient by reusing some allocations + pub fn reparse( + &mut self, + pattern: &str, + case_matching: CaseMatching, + normalize: Normalization, + ) { + self.atoms.clear(); + let atoms = pattern_atoms(pattern).filter_map(|atom| { + let atom = Atom::parse(atom, case_matching, normalize); + if atom.needle.is_empty() { + return None; + } + Some(atom) + }); + self.atoms.extend(atoms); + } +} + +impl Clone for Pattern { + fn clone(&self) -> Self { + Self { + atoms: self.atoms.clone(), + } + } + + fn clone_from(&mut self, source: &Self) { + self.atoms.clone_from(&source.atoms); + } +} diff --git a/crates/atuin-nucleo/matcher/src/pattern/tests.rs b/crates/atuin-nucleo/matcher/src/pattern/tests.rs new file mode 100644 index 00000000..88880ba9 --- /dev/null +++ b/crates/atuin-nucleo/matcher/src/pattern/tests.rs @@ -0,0 +1,149 @@ +use crate::pattern::{Atom, AtomKind, CaseMatching, Normalization, Pattern}; + +#[test] +fn negative() { + let pat = Atom::parse("!foo", CaseMatching::Smart, Normalization::Smart); + assert!(pat.negative); + assert_eq!(pat.kind, AtomKind::Substring); + assert_eq!(pat.needle.to_string(), "foo"); + let pat = Atom::parse("!^foo", CaseMatching::Smart, Normalization::Smart); + assert!(pat.negative); + assert_eq!(pat.kind, AtomKind::Prefix); + assert_eq!(pat.needle.to_string(), "foo"); + let pat = Atom::parse("!foo$", CaseMatching::Smart, Normalization::Smart); + assert!(pat.negative); + assert_eq!(pat.kind, AtomKind::Postfix); + assert_eq!(pat.needle.to_string(), "foo"); + let pat = Atom::parse("!^foo$", CaseMatching::Smart, Normalization::Smart); + assert!(pat.negative); + assert_eq!(pat.kind, AtomKind::Exact); + assert_eq!(pat.needle.to_string(), "foo"); +} + +#[test] +fn pattern_kinds() { + let pat = Atom::parse("foo", CaseMatching::Smart, Normalization::Smart); + assert!(!pat.negative); + assert_eq!(pat.kind, AtomKind::Fuzzy); + assert_eq!(pat.needle.to_string(), "foo"); + let pat = Atom::parse("'foo", CaseMatching::Smart, Normalization::Smart); + assert!(!pat.negative); + assert_eq!(pat.kind, AtomKind::Substring); + assert_eq!(pat.needle.to_string(), "foo"); + let pat = Atom::parse("^foo", CaseMatching::Smart, Normalization::Smart); + assert!(!pat.negative); + assert_eq!(pat.kind, AtomKind::Prefix); + assert_eq!(pat.needle.to_string(), "foo"); + let pat = Atom::parse("foo$", CaseMatching::Smart, Normalization::Smart); + assert!(!pat.negative); + assert_eq!(pat.kind, AtomKind::Postfix); + assert_eq!(pat.needle.to_string(), "foo"); + let pat = Atom::parse("^foo$", CaseMatching::Smart, Normalization::Smart); + assert!(!pat.negative); + assert_eq!(pat.kind, AtomKind::Exact); + assert_eq!(pat.needle.to_string(), "foo"); +} + +#[test] +fn case_matching() { + let pat = Atom::parse("foo", CaseMatching::Smart, Normalization::Smart); + assert!(pat.ignore_case); + assert_eq!(pat.needle.to_string(), "foo"); + let pat = Atom::parse("Foo", CaseMatching::Smart, Normalization::Smart); + assert!(!pat.ignore_case); + assert_eq!(pat.needle.to_string(), "Foo"); + let pat = Atom::parse("Foo", CaseMatching::Ignore, Normalization::Smart); + assert!(pat.ignore_case); + assert_eq!(pat.needle.to_string(), "foo"); + let pat = Atom::parse("Foo", CaseMatching::Respect, Normalization::Smart); + assert!(!pat.ignore_case); + assert_eq!(pat.needle.to_string(), "Foo"); + let pat = Atom::parse("Foo", CaseMatching::Respect, Normalization::Smart); + assert!(!pat.ignore_case); + assert_eq!(pat.needle.to_string(), "Foo"); + let pat = Atom::parse("Äxx", CaseMatching::Ignore, Normalization::Smart); + assert!(pat.ignore_case); + assert_eq!(pat.needle.to_string(), "äxx"); + let pat = Atom::parse("Äxx", CaseMatching::Respect, Normalization::Smart); + assert!(!pat.ignore_case); + let pat = Atom::parse("Axx", CaseMatching::Smart, Normalization::Smart); + assert!(!pat.ignore_case); + assert_eq!(pat.needle.to_string(), "Axx"); + let pat = Atom::parse("你xx", CaseMatching::Smart, Normalization::Smart); + assert!(pat.ignore_case); + assert_eq!(pat.needle.to_string(), "你xx"); + let pat = Atom::parse("你xx", CaseMatching::Ignore, Normalization::Smart); + assert!(pat.ignore_case); + assert_eq!(pat.needle.to_string(), "你xx"); + let pat = Atom::parse("Ⲽxx", CaseMatching::Smart, Normalization::Smart); + assert!(!pat.ignore_case); + assert_eq!(pat.needle.to_string(), "Ⲽxx"); + let pat = Atom::parse("Ⲽxx", CaseMatching::Ignore, Normalization::Smart); + assert!(pat.ignore_case); + assert_eq!(pat.needle.to_string(), "ⲽxx"); +} + +#[test] +fn escape() { + let pat = Atom::parse("foo\\ bar", CaseMatching::Smart, Normalization::Smart); + assert_eq!(pat.needle.to_string(), "foo bar"); + let pat = Atom::parse("\\!foo", CaseMatching::Smart, Normalization::Smart); + assert_eq!(pat.needle.to_string(), "!foo"); + assert_eq!(pat.kind, AtomKind::Fuzzy); + let pat = Atom::parse("\\'foo", CaseMatching::Smart, Normalization::Smart); + assert_eq!(pat.needle.to_string(), "'foo"); + assert_eq!(pat.kind, AtomKind::Fuzzy); + let pat = Atom::parse("\\^foo", CaseMatching::Smart, Normalization::Smart); + assert_eq!(pat.needle.to_string(), "^foo"); + assert_eq!(pat.kind, AtomKind::Fuzzy); + let pat = Atom::parse("foo\\$", CaseMatching::Smart, Normalization::Smart); + assert_eq!(pat.needle.to_string(), "foo$"); + assert_eq!(pat.kind, AtomKind::Fuzzy); + let pat = Atom::parse("^foo\\$", CaseMatching::Smart, Normalization::Smart); + assert_eq!(pat.needle.to_string(), "foo$"); + assert_eq!(pat.kind, AtomKind::Prefix); + let pat = Atom::parse("\\^foo\\$", CaseMatching::Smart, Normalization::Smart); + assert_eq!(pat.needle.to_string(), "^foo$"); + assert_eq!(pat.kind, AtomKind::Fuzzy); + let pat = Atom::parse("\\!^foo\\$", CaseMatching::Smart, Normalization::Smart); + assert_eq!(pat.needle.to_string(), "!^foo$"); + assert_eq!(pat.kind, AtomKind::Fuzzy); + let pat = Atom::parse("!\\^foo\\$", CaseMatching::Smart, Normalization::Smart); + assert_eq!(pat.needle.to_string(), "^foo$"); + assert_eq!(pat.kind, AtomKind::Substring); +} + +#[test] +fn pattern_atoms() { + assert_eq!( + Pattern::parse("a b", CaseMatching::Ignore, Normalization::Smart).atoms, + vec![ + Atom::parse("a", CaseMatching::Ignore, Normalization::Smart), + Atom::parse("b", CaseMatching::Ignore, Normalization::Smart), + ] + ); + + assert_eq!( + Pattern::parse("a\n b", CaseMatching::Ignore, Normalization::Smart).atoms, + vec![ + Atom::parse("a", CaseMatching::Ignore, Normalization::Smart), + Atom::parse("b", CaseMatching::Ignore, Normalization::Smart), + ] + ); + + assert_eq!( + Pattern::parse(" a b\r\n", CaseMatching::Ignore, Normalization::Smart).atoms, + vec![ + Atom::parse("a", CaseMatching::Ignore, Normalization::Smart), + Atom::parse("b", CaseMatching::Ignore, Normalization::Smart), + ] + ); + + assert_eq!( + Pattern::parse("ほ げ", CaseMatching::Smart, Normalization::Smart).atoms, + vec![ + Atom::parse("ほ", CaseMatching::Smart, Normalization::Smart), + Atom::parse("げ", CaseMatching::Smart, Normalization::Smart), + ], + ) +} diff --git a/crates/atuin-nucleo/matcher/src/prefilter.rs b/crates/atuin-nucleo/matcher/src/prefilter.rs new file mode 100644 index 00000000..1e79cc2c --- /dev/null +++ b/crates/atuin-nucleo/matcher/src/prefilter.rs @@ -0,0 +1,98 @@ +use ::memchr::{memchr, memchr2, memrchr, memrchr2}; + +use crate::chars::Char; +use crate::utf32_str::Utf32Str; +use crate::Matcher; + +#[inline(always)] +fn find_ascii_ignore_case(c: u8, haystack: &[u8]) -> Option<usize> { + if c >= b'a' && c <= b'z' { + memchr2(c, c - 32, haystack) + } else { + memchr(c, haystack) + } +} + +#[inline(always)] +fn find_ascii_ignore_case_rev(c: u8, haystack: &[u8]) -> Option<usize> { + if c >= b'a' && c <= b'z' { + memrchr2(c, c - 32, haystack) + } else { + memrchr(c, haystack) + } +} + +impl Matcher { + pub(crate) fn prefilter_ascii( + &self, + mut haystack: &[u8], + needle: &[u8], + only_greedy: bool, + ) -> Option<(usize, usize, usize)> { + if self.config.ignore_case { + let start = + find_ascii_ignore_case(needle[0], &haystack[..haystack.len() - needle.len() + 1])?; + let mut greedy_end = start + 1; + haystack = &haystack[greedy_end..]; + for &c in &needle[1..] { + let idx = find_ascii_ignore_case(c, haystack)? + 1; + greedy_end += idx; + haystack = &haystack[idx..]; + } + if only_greedy { + Some((start, greedy_end, greedy_end)) + } else { + let end = greedy_end + + find_ascii_ignore_case_rev(*needle.last().unwrap(), haystack) + .map_or(0, |i| i + 1); + Some((start, greedy_end, end)) + } + } else { + let start = memchr(needle[0], &haystack[..haystack.len() - needle.len() + 1])?; + let mut greedy_end = start + 1; + haystack = &haystack[greedy_end..]; + for &c in &needle[1..] { + let idx = memchr(c, haystack)? + 1; + greedy_end += idx; + haystack = &haystack[idx..]; + } + if only_greedy { + Some((start, greedy_end, greedy_end)) + } else { + let end = + greedy_end + memrchr(*needle.last().unwrap(), haystack).map_or(0, |i| i + 1); + Some((start, greedy_end, end)) + } + } + } + + pub(crate) fn prefilter_non_ascii( + &self, + haystack: &[char], + needle: Utf32Str<'_>, + only_greedy: bool, + ) -> Option<(usize, usize)> { + let needle_char = needle.get(0); + let start = haystack[..haystack.len() - needle.len() + 1] + .iter() + .position(|c| c.normalize(&self.config) == needle_char)?; + let needle_char = needle.last(); + if only_greedy { + if haystack.len() - start < needle.len() { + return None; + } + Some((start, start + 1)) + } else { + let end = haystack.len() + - haystack[start + 1..] + .iter() + .rev() + .position(|c| c.normalize(&self.config) == needle_char)?; + if end - start < needle.len() { + return None; + } + + Some((start, end)) + } + } +} diff --git a/crates/atuin-nucleo/matcher/src/score.rs b/crates/atuin-nucleo/matcher/src/score.rs new file mode 100644 index 00000000..c934a8ef --- /dev/null +++ b/crates/atuin-nucleo/matcher/src/score.rs @@ -0,0 +1,158 @@ +use std::cmp::max; + +use crate::chars::{Char, CharClass}; +use crate::{Config, Matcher}; + +pub(crate) const SCORE_MATCH: u16 = 16; +pub(crate) const PENALTY_GAP_START: u16 = 3; +pub(crate) const PENALTY_GAP_EXTENSION: u16 = 1; +/// If the prefer_prefix option is enabled we want to penalize +/// the initial gap. The prefix should not be too much +pub(crate) const PREFIX_BONUS_SCALE: u16 = 2; +pub(crate) const MAX_PREFIX_BONUS: u16 = BONUS_BOUNDARY; + +// We prefer matches at the beginning of a word, but the bonus should not be +// too great to prevent the longer acronym matches from always winning over +// shorter fuzzy matches. The bonus point here was specifically chosen that +// the bonus is cancelled when the gap between the acronyms grows over +// 8 characters, which is approximately the average length of the words found +// in web2 dictionary and my file system. +pub(crate) const BONUS_BOUNDARY: u16 = SCORE_MATCH / 2; + +// Edge-triggered bonus for matches in camelCase words. +// Their value should be BONUS_BOUNDARY - PENALTY_GAP_EXTENSION = 7. +// However, this priporitzes camel case over non-camel case. +// In fzf/skim this is not a problem since they score off the max +// consecutive bonus. However, we don't do that (because its incorrect) +// so to avoids prioritizing camel we use a lower bonus. I think that's fine +// usually camel case is wekaer boundary than actual wourd boundaries anyway +// This also has the nice sideeffect of perfectly balancing out +// camel case, snake case and the consecutive version of the word +pub(crate) const BONUS_CAMEL123: u16 = BONUS_BOUNDARY - PENALTY_GAP_START; + +/// Although bonus point for non-word characters is non-contextual, we need it +/// for computing bonus points for consecutive chunks starting with a non-word +/// character. +pub(crate) const BONUS_NON_WORD: u16 = BONUS_BOUNDARY; + +// Minimum bonus point given to characters in consecutive chunks. +// Note that bonus points for consecutive matches shouldn't have needed if we +// used fixed match score as in the original algorithm. +pub(crate) const BONUS_CONSECUTIVE: u16 = PENALTY_GAP_START + PENALTY_GAP_EXTENSION; + +// The first character in the typed pattern usually has more significance +// than the rest so it's important that it appears at special positions where +// bonus points are given, e.g. "to-go" vs. "ongoing" on "og" or on "ogo". +// The amount of the extra bonus should be limited so that the gap penalty is +// still respected. +pub(crate) const BONUS_FIRST_CHAR_MULTIPLIER: u16 = 2; + +impl Config { + #[inline] + pub(crate) fn bonus_for(&self, prev_class: CharClass, class: CharClass) -> u16 { + if class > CharClass::Delimiter { + // transition from non word to word + match prev_class { + CharClass::Whitespace => return self.bonus_boundary_white, + CharClass::Delimiter => return self.bonus_boundary_delimiter, + CharClass::NonWord => return BONUS_BOUNDARY, + _ => (), + } + } + if prev_class == CharClass::Lower && class == CharClass::Upper + || prev_class != CharClass::Number && class == CharClass::Number + { + // camelCase letter123 + BONUS_CAMEL123 + } else if class == CharClass::Whitespace { + self.bonus_boundary_white + } else if class == CharClass::NonWord { + return BONUS_NON_WORD; + } else { + 0 + } + } +} +impl Matcher { + #[inline(always)] + pub(crate) fn bonus_for(&self, prev_class: CharClass, class: CharClass) -> u16 { + self.config.bonus_for(prev_class, class) + } + + pub(crate) fn calculate_score<const INDICES: bool, H: Char + PartialEq<N>, N: Char>( + &mut self, + haystack: &[H], + needle: &[N], + start: usize, + end: usize, + indices: &mut Vec<u32>, + ) -> u16 { + if INDICES { + indices.reserve(needle.len()); + } + + let mut prev_class = start + .checked_sub(1) + .map(|i| haystack[i].char_class(&self.config)) + .unwrap_or(self.config.initial_char_class); + let mut needle_iter = needle.iter(); + let mut needle_char = *needle_iter.next().unwrap(); + + let mut in_gap = false; + let mut consecutive = 1; + + // unrolled the first iteration to make applying the first char multiplier less awkward + if INDICES { + indices.push(start as u32) + } + let class = haystack[start].char_class(&self.config); + let mut first_bonus = self.bonus_for(prev_class, class); + let mut score = SCORE_MATCH + first_bonus * BONUS_FIRST_CHAR_MULTIPLIER; + prev_class = class; + needle_char = *needle_iter.next().unwrap_or(&needle_char); + + for (i, c) in haystack[start + 1..end].iter().enumerate() { + let (c, class) = c.char_class_and_normalize(&self.config); + if c == needle_char { + if INDICES { + indices.push(i as u32 + start as u32 + 1) + } + let mut bonus = self.bonus_for(prev_class, class); + if consecutive != 0 { + if bonus >= BONUS_BOUNDARY && bonus > first_bonus { + first_bonus = bonus + } + bonus = max(max(bonus, first_bonus), BONUS_CONSECUTIVE); + } else { + first_bonus = bonus; + } + score += SCORE_MATCH + bonus; + in_gap = false; + consecutive += 1; + if let Some(&next) = needle_iter.next() { + needle_char = next; + } + } else { + let penalty = if in_gap { + PENALTY_GAP_EXTENSION + } else { + PENALTY_GAP_START + }; + score = score.saturating_sub(penalty); + in_gap = true; + consecutive = 0; + } + prev_class = class; + } + if self.config.prefer_prefix { + if start != 0 { + let penalty = PENALTY_GAP_START + + PENALTY_GAP_START * (start - 1).min(u16::MAX as usize) as u16; + score += MAX_PREFIX_BONUS.saturating_sub(penalty / PREFIX_BONUS_SCALE); + } else { + score += MAX_PREFIX_BONUS; + } + } + score + } +} diff --git a/crates/atuin-nucleo/matcher/src/tests.rs b/crates/atuin-nucleo/matcher/src/tests.rs new file mode 100644 index 00000000..32a02403 --- /dev/null +++ b/crates/atuin-nucleo/matcher/src/tests.rs @@ -0,0 +1,771 @@ +use crate::chars::Char; +use crate::pattern::{CaseMatching, Normalization, Pattern}; +use crate::score::{ + BONUS_BOUNDARY, BONUS_CAMEL123, BONUS_CONSECUTIVE, BONUS_FIRST_CHAR_MULTIPLIER, BONUS_NON_WORD, + MAX_PREFIX_BONUS, PENALTY_GAP_EXTENSION, PENALTY_GAP_START, SCORE_MATCH, +}; +use crate::utf32_str::Utf32Str; +use crate::{Config, Matcher}; + +use Algorithm::*; + +#[derive(Debug)] +enum Algorithm { + FuzzyOptimal, + FuzzyGreedy, + Substring, + Prefix, + Postfix, + Exact, +} + +fn assert_matches( + algorithm: &[Algorithm], + normalize: bool, + case_sensitive: bool, + path: bool, + prefer_prefix: bool, + cases: &[(&str, &str, &[u32], u16)], +) { + let mut config = Config { + normalize, + ignore_case: !case_sensitive, + prefer_prefix, + ..Config::DEFAULT + }; + if path { + config.set_match_paths(); + } + let mut matcher = Matcher::new(config); + let mut matched_indices = Vec::new(); + let mut needle_buf = Vec::new(); + let mut haystack_buf = Vec::new(); + for &(haystack, needle, indices, mut score) in cases { + let needle = if !case_sensitive { + needle.to_lowercase() + } else { + needle.to_owned() + }; + let needle = Utf32Str::new(&needle, &mut needle_buf); + let haystack = Utf32Str::new(haystack, &mut haystack_buf); + score += needle.len() as u16 * SCORE_MATCH; + for algo in algorithm { + println!("xx {matched_indices:?} {algo:?}"); + matched_indices.clear(); + let res = match algo { + FuzzyOptimal => matcher.fuzzy_indices(haystack, needle, &mut matched_indices), + FuzzyGreedy => matcher.fuzzy_indices_greedy(haystack, needle, &mut matched_indices), + Substring => matcher.substring_indices(haystack, needle, &mut matched_indices), + Prefix => matcher.prefix_indices(haystack, needle, &mut matched_indices), + Postfix => matcher.postfix_indices(haystack, needle, &mut matched_indices), + Exact => matcher.exact_indices(haystack, needle, &mut matched_indices), + }; + println!("{matched_indices:?}"); + let match_chars: Vec<_> = matched_indices + .iter() + .map(|&i| haystack.get(i).normalize(&matcher.config)) + .collect(); + let needle_chars: Vec<_> = needle.chars().collect(); + + assert_eq!( + res, + Some(score), + "{needle:?} did not match {haystack:?}: matched {match_chars:?} {matched_indices:?} {algo:?}" + ); + assert_eq!( + matched_indices, indices, + "{needle:?} match {haystack:?} {algo:?}" + ); + assert_eq!( + match_chars, needle_chars, + "{needle:?} match {haystack:?} indices are incorrect {matched_indices:?} {algo:?}" + ); + } + } +} + +fn assert_not_matches_with( + normalize: bool, + case_sensitive: bool, + algorithm: &[Algorithm], + cases: &[(&str, &str)], +) { + let config = Config { + normalize, + ignore_case: !case_sensitive, + ..Config::DEFAULT + }; + let mut matcher = Matcher::new(config); + let mut needle_buf = Vec::new(); + let mut haystack_buf = Vec::new(); + for &(haystack, needle) in cases { + let needle = if !case_sensitive { + needle.to_lowercase() + } else { + needle.to_owned() + }; + let needle = Utf32Str::new(&needle, &mut needle_buf); + let haystack = Utf32Str::new(haystack, &mut haystack_buf); + + for algo in algorithm { + let res = match algo { + FuzzyOptimal => matcher.fuzzy_match(haystack, needle), + FuzzyGreedy => matcher.fuzzy_match_greedy(haystack, needle), + Substring => matcher.substring_match(haystack, needle), + Prefix => matcher.prefix_match(haystack, needle), + Postfix => matcher.postfix_match(haystack, needle), + Exact => matcher.exact_match(haystack, needle), + }; + assert_eq!( + res, None, + "{needle:?} should not match {haystack:?} {algo:?}" + ); + } + } +} + +pub fn assert_not_matches(normalize: bool, case_sensitive: bool, cases: &[(&str, &str)]) { + assert_not_matches_with( + normalize, + case_sensitive, + &[FuzzyOptimal, FuzzyGreedy, Substring, Prefix, Postfix, Exact], + cases, + ) +} + +const BONUS_BOUNDARY_WHITE: u16 = Config::DEFAULT.bonus_boundary_white; +const BONUS_BOUNDARY_DELIMITER: u16 = Config::DEFAULT.bonus_boundary_delimiter; + +#[test] +fn test_fuzzy() { + assert_matches( + &[FuzzyGreedy, FuzzyOptimal], + false, + false, + false, + false, + &[ + ( + "fooBarbaz1", + "obr", + &[2, 3, 5], + BONUS_CAMEL123 - PENALTY_GAP_START, + ), + ( + "/usr/share/doc/at/ChangeLog", + "changelog", + &[18, 19, 20, 21, 22, 23, 24, 25, 26], + (BONUS_FIRST_CHAR_MULTIPLIER + 8) * BONUS_BOUNDARY_DELIMITER, + ), + ( + "fooBarbaz1", + "br", + &[3, 5], + BONUS_CAMEL123 * BONUS_FIRST_CHAR_MULTIPLIER - PENALTY_GAP_START, + ), + ( + "foo bar baz", + "fbb", + &[0, 4, 8], + BONUS_BOUNDARY_WHITE * BONUS_FIRST_CHAR_MULTIPLIER + BONUS_BOUNDARY_WHITE * 2 + - 2 * PENALTY_GAP_START + - 4 * PENALTY_GAP_EXTENSION, + ), + ( + "/AutomatorDocument.icns", + "rdoc", + &[9, 10, 11, 12], + BONUS_CAMEL123 + 2 * BONUS_CONSECUTIVE, + ), + ( + "/man1/zshcompctl.1", + "zshc", + &[6, 7, 8, 9], + BONUS_BOUNDARY_DELIMITER * (BONUS_FIRST_CHAR_MULTIPLIER + 3), + ), + ( + "/.oh-my-zsh/cache", + "zshc", + &[8, 9, 10, 12], + BONUS_BOUNDARY * (BONUS_FIRST_CHAR_MULTIPLIER + 2) - PENALTY_GAP_START + + BONUS_BOUNDARY_DELIMITER, + ), + ( + "ab0123 456", + "12356", + &[3, 4, 5, 8, 9], + BONUS_CONSECUTIVE * 3 - PENALTY_GAP_START - PENALTY_GAP_EXTENSION, + ), + ( + "abc123 456", + "12356", + &[3, 4, 5, 8, 9], + BONUS_CAMEL123 * (BONUS_FIRST_CHAR_MULTIPLIER + 2) + - PENALTY_GAP_START + - PENALTY_GAP_EXTENSION + + BONUS_CONSECUTIVE, + ), + ( + "foo/bar/baz", + "fbb", + &[0, 4, 8], + BONUS_BOUNDARY_WHITE * BONUS_FIRST_CHAR_MULTIPLIER + BONUS_BOUNDARY_DELIMITER * 2 + - 2 * PENALTY_GAP_START + - 4 * PENALTY_GAP_EXTENSION, + ), + ( + "fooBarBaz", + "fbb", + &[0, 3, 6], + BONUS_BOUNDARY_WHITE * BONUS_FIRST_CHAR_MULTIPLIER + BONUS_CAMEL123 * 2 + - 2 * PENALTY_GAP_START + - 2 * PENALTY_GAP_EXTENSION, + ), + ( + "foo barbaz", + "fbb", + &[0, 4, 7], + BONUS_BOUNDARY_WHITE * BONUS_FIRST_CHAR_MULTIPLIER + BONUS_BOUNDARY_WHITE + - PENALTY_GAP_START * 2 + - PENALTY_GAP_EXTENSION * 3, + ), + ( + "fooBar Baz", + "foob", + &[0, 1, 2, 3], + BONUS_BOUNDARY_WHITE * (BONUS_FIRST_CHAR_MULTIPLIER + 3), + ), + ( + "xFoo-Bar Baz", + "foo-b", + &[1, 2, 3, 4, 5], + BONUS_CAMEL123 * (BONUS_FIRST_CHAR_MULTIPLIER + 2) + 2 * BONUS_NON_WORD, + ), + ], + ); +} + +#[test] +fn empty_needle() { + assert_matches( + &[Substring, Prefix, Postfix, FuzzyGreedy, FuzzyOptimal, Exact], + false, + false, + false, + false, + &[("foo bar baz", "", &[], 0)], + ); +} + +#[test] +fn test_substring() { + assert_matches( + &[Substring, Prefix], + false, + false, + false, + false, + &[ + ( + "foo bar baz", + "foo", + &[0, 1, 2], + BONUS_BOUNDARY_WHITE * (BONUS_FIRST_CHAR_MULTIPLIER + 2), + ), + ( + " foo bar baz", + "FOO", + &[1, 2, 3], + BONUS_BOUNDARY_WHITE * (BONUS_FIRST_CHAR_MULTIPLIER + 2), + ), + ( + " foo bar baz", + " FOO", + &[0, 1, 2, 3], + BONUS_BOUNDARY_WHITE * (BONUS_FIRST_CHAR_MULTIPLIER + 3), + ), + ], + ); + assert_matches( + &[Substring, Postfix], + false, + false, + false, + false, + &[ + ( + "foo bar baz", + "baz", + &[8, 9, 10], + BONUS_BOUNDARY_WHITE * (BONUS_FIRST_CHAR_MULTIPLIER + 2), + ), + ( + "foo bar baz ", + "baz", + &[8, 9, 10], + BONUS_BOUNDARY_WHITE * (BONUS_FIRST_CHAR_MULTIPLIER + 2), + ), + ( + "foo bar baz ", + "baz ", + &[8, 9, 10, 11], + BONUS_BOUNDARY_WHITE * (BONUS_FIRST_CHAR_MULTIPLIER + 3), + ), + ], + ); + assert_matches( + &[Substring, Prefix, Postfix, Exact, FuzzyGreedy, FuzzyOptimal], + false, + false, + false, + false, + &[ + ( + "foo", + "foo", + &[0, 1, 2], + BONUS_BOUNDARY_WHITE * (BONUS_FIRST_CHAR_MULTIPLIER + 2), + ), + ( + " foo", + "foo", + &[1, 2, 3], + BONUS_BOUNDARY_WHITE * (BONUS_FIRST_CHAR_MULTIPLIER + 2), + ), + ( + " foo", + " foo", + &[0, 1, 2, 3], + BONUS_BOUNDARY_WHITE * (BONUS_FIRST_CHAR_MULTIPLIER + 3), + ), + ], + ); + assert_matches( + &[Substring], + false, + false, + false, + false, + &[ + ( + "fooBarbaz1", + "oba", + &[2, 3, 4], + BONUS_CAMEL123 + BONUS_CONSECUTIVE, + ), + ( + "/AutomatorDocument.icns", + "rdoc", + &[9, 10, 11, 12], + BONUS_CAMEL123 + 2 * BONUS_CONSECUTIVE, + ), + ( + "/man1/zshcompctl.1", + "zshc", + &[6, 7, 8, 9], + BONUS_BOUNDARY_DELIMITER * (BONUS_FIRST_CHAR_MULTIPLIER + 3), + ), + ( + "/.oh-my-zsh/cache", + "zsh/c", + &[8, 9, 10, 11, 12], + BONUS_BOUNDARY * (BONUS_FIRST_CHAR_MULTIPLIER + 2) + + BONUS_NON_WORD + + BONUS_BOUNDARY_DELIMITER, + ), + ], + ); + assert_not_matches_with( + true, + false, + &[Prefix, Substring, Postfix, Exact], + &[( + "At the Road’s End - Seeming - SOL: A Self-Banishment Ritual", + "adi", + )], + ) +} + +#[test] +fn test_substring_case_sensitive() { + assert_matches( + &[Substring, Prefix], + false, + true, + false, + false, + &[ + ( + "Foo bar baz", + "Foo", + &[0, 1, 2], + BONUS_BOUNDARY_WHITE * (BONUS_FIRST_CHAR_MULTIPLIER + 2), + ), + ( + "Fȫô bar baz", + "Fȫô", + &[0, 1, 2], + BONUS_BOUNDARY_WHITE * (BONUS_FIRST_CHAR_MULTIPLIER + 2), + ), + ( + "Foo ฿ar baz", + "Foo", + &[0, 1, 2], + BONUS_BOUNDARY_WHITE * (BONUS_FIRST_CHAR_MULTIPLIER + 2), + ), + ], + ); + assert_not_matches_with(false, true, &[Substring, Prefix], &[("foo bar baz", "Foo")]); +} + +#[test] +fn test_fuzzy_case_sensitive() { + assert_matches( + &[FuzzyGreedy, FuzzyOptimal], + false, + true, + false, + false, + &[ + ( + "fooBarbaz1", + "oBr", + &[2, 3, 5], + BONUS_CAMEL123 - PENALTY_GAP_START, + ), + ( + "Foo/Bar/Baz", + "FBB", + &[0, 4, 8], + BONUS_BOUNDARY_WHITE * BONUS_FIRST_CHAR_MULTIPLIER + BONUS_BOUNDARY_DELIMITER * 2 + - 2 * PENALTY_GAP_START + - 4 * PENALTY_GAP_EXTENSION, + ), + ( + "FooBarBaz", + "FBB", + &[0, 3, 6], + BONUS_BOUNDARY_WHITE * BONUS_FIRST_CHAR_MULTIPLIER + BONUS_CAMEL123 * 2 + - 2 * PENALTY_GAP_START + - 2 * PENALTY_GAP_EXTENSION, + ), + ( + "FooBar Baz", + "FooB", + &[0, 1, 2, 3], + BONUS_BOUNDARY_WHITE * (BONUS_FIRST_CHAR_MULTIPLIER + 3), + ), + ("foo-bar", "o-ba", &[2, 3, 4, 5], BONUS_NON_WORD * 3), + ], + ); +} + +#[test] +fn test_normalize() { + assert_matches( + &[FuzzyGreedy, FuzzyOptimal], + true, + false, + false, + false, + &[ + ( + "Só Danço Samba", + "So", + &[0, 1], + BONUS_BOUNDARY_WHITE * (BONUS_FIRST_CHAR_MULTIPLIER + 1), + ), + ( + "Só Danço Samba", + "sodc", + &[0, 1, 3, 6], + BONUS_BOUNDARY_WHITE * (BONUS_FIRST_CHAR_MULTIPLIER + 1) - PENALTY_GAP_START + + BONUS_BOUNDARY_WHITE + - PENALTY_GAP_START + - PENALTY_GAP_EXTENSION, + ), + ( + "Danço", + "danco", + &[0, 1, 2, 3, 4], + BONUS_BOUNDARY_WHITE * (BONUS_FIRST_CHAR_MULTIPLIER + 4), + ), + ( + "DanÇo", + "danco", + &[0, 1, 2, 3, 4], + BONUS_BOUNDARY_WHITE * (BONUS_FIRST_CHAR_MULTIPLIER + 4), + ), + ( + "xÇando", + "cando", + &[1, 2, 3, 4, 5], + BONUS_CAMEL123 * (BONUS_FIRST_CHAR_MULTIPLIER + 4), + ), + ("ۂ(GCGɴCG", "n", &[5], 0), + ], + ) +} + +#[test] +fn test_unicode() { + assert_matches( + &[FuzzyGreedy, FuzzyOptimal, Substring], + true, + false, + false, + false, + &[ + ( + "你好世界", + "你好", + &[0, 1], + BONUS_BOUNDARY_WHITE * (BONUS_FIRST_CHAR_MULTIPLIER + 1), + ), + ( + " 你好世界", + "你好", + &[1, 2], + BONUS_BOUNDARY_WHITE * (BONUS_FIRST_CHAR_MULTIPLIER + 1), + ), + ], + ); + assert_matches( + &[FuzzyGreedy, FuzzyOptimal], + true, + false, + false, + false, + &[( + "你好世界", + "你世", + &[0, 2], + BONUS_BOUNDARY_WHITE * BONUS_FIRST_CHAR_MULTIPLIER - PENALTY_GAP_START, + )], + ); + assert_not_matches( + false, + false, + &[("Flibbertigibbet / イタズラっ子たち", "lying")], + ); +} + +#[test] +fn test_long_str() { + assert_matches( + &[FuzzyGreedy, FuzzyOptimal], + false, + false, + false, + false, + &[( + &"x".repeat(u16::MAX as usize + 1), + "xx", + &[0, 1], + BONUS_BOUNDARY_WHITE * (BONUS_FIRST_CHAR_MULTIPLIER + 1), + )], + ); +} + +#[test] +fn test_casing() { + assert_matches( + &[FuzzyGreedy, FuzzyOptimal], + false, + false, + false, + false, + &[ + // these two have the same score + ( + "fooBar", + "foobar", + &[0, 1, 2, 3, 4, 5], + BONUS_BOUNDARY_WHITE * (BONUS_FIRST_CHAR_MULTIPLIER + 5), + ), + ( + "foobar", + "foobar", + &[0, 1, 2, 3, 4, 5], + BONUS_BOUNDARY_WHITE * (BONUS_FIRST_CHAR_MULTIPLIER + 5), + ), + // these two have the same score (slightly lower than the other two: 60 instead of 70) + ( + "foo-bar", + "foobar", + &[0, 1, 2, 4, 5, 6], + BONUS_BOUNDARY_WHITE * (BONUS_FIRST_CHAR_MULTIPLIER + 2) - PENALTY_GAP_START + + BONUS_BOUNDARY * 3, + ), + ( + "foo_bar", + "foobar", + &[0, 1, 2, 4, 5, 6], + BONUS_BOUNDARY_WHITE * (BONUS_FIRST_CHAR_MULTIPLIER + 2) - PENALTY_GAP_START + + BONUS_BOUNDARY * 3, + ), + ], + ) +} + +#[test] +fn test_optimal() { + assert_matches( + &[FuzzyOptimal], + false, + false, + false, + false, + &[ + ( + "axxx xx ", + "xx", + &[5, 6], + BONUS_BOUNDARY_WHITE * (BONUS_FIRST_CHAR_MULTIPLIER + 1), + ), + ( + "SS!H", + "S!", + &[0, 2], + BONUS_BOUNDARY_WHITE * BONUS_FIRST_CHAR_MULTIPLIER - PENALTY_GAP_START + + BONUS_NON_WORD, + ), + // this case is a cool example of why our algorithm is more than fzf + // we handle this corretly detect that it's better to match + // the second f instead of the third yielding a higher score + // (despite using the same scoring function!) + ( + "xf.foo", + "xfoo", + &[0, 3, 4, 5], + BONUS_BOUNDARY_WHITE * BONUS_FIRST_CHAR_MULTIPLIER + - PENALTY_GAP_START + - PENALTY_GAP_EXTENSION + + BONUS_BOUNDARY * 3, + ), + ( + "xf fo", + "xfo", + &[0, 3, 4], + BONUS_BOUNDARY_WHITE * (BONUS_FIRST_CHAR_MULTIPLIER + 2) + - PENALTY_GAP_START + - PENALTY_GAP_EXTENSION, + ), + ], + ); +} + +#[test] +fn test_reject() { + assert_not_matches( + true, + false, + &[ + ("你好界", "abc"), + ("你好界", "a"), + ("你好世界", "富"), + ("Só Danço Samba", "sox"), + ("fooBarbaz", "fooBarbazz"), + ("fooBarbaz", "c"), + ], + ); + assert_not_matches( + true, + true, + &[ + ("你好界", "abc"), + ("abc", "你"), + ("abc", "A"), + ("abc", "d"), + ("你好世界", "富"), + ("Só Danço Samba", "sox"), + ("fooBarbaz", "oBZ"), + ("Foo Bar Baz", "fbb"), + ("fooBarbaz", "fooBarbazz"), + ], + ); + assert_not_matches( + false, + true, + &[ + ("Só Danço Samba", "sod"), + ("Só Danço Samba", "soc"), + ("Só Danç", "So"), + ], + ); + assert_not_matches(false, false, &[("ۂۂfoۂۂ", "foo")]); +} + +#[test] +fn test_prefer_prefix() { + assert_matches( + &[FuzzyOptimal, FuzzyGreedy], + false, + false, + false, + true, + &[ + ( + "Moby Dick", + "md", + &[0, 5], + BONUS_BOUNDARY_WHITE * (BONUS_FIRST_CHAR_MULTIPLIER + 1) + MAX_PREFIX_BONUS + - PENALTY_GAP_START + - 3 * PENALTY_GAP_EXTENSION, + ), + ( + "Though I cannot tell why it was exactly that those stage managers, the Fates, put me down for this shabby part of a whaling voyage", + "md", + &[82, 85], + BONUS_BOUNDARY_WHITE * (BONUS_FIRST_CHAR_MULTIPLIER + 1) + - PENALTY_GAP_START + - PENALTY_GAP_EXTENSION, + ), + ], + ); +} + +#[test] +fn test_single_char_needle() { + assert_matches( + &[FuzzyOptimal], + false, + false, + false, + false, + &[( + "foO", + "o", + &[2], + BONUS_FIRST_CHAR_MULTIPLIER * BONUS_CAMEL123, + )], + ); + assert_matches( + &[FuzzyOptimal], + false, + false, + false, + false, + &[( + "föÖ", + "ö", + &[2], + BONUS_FIRST_CHAR_MULTIPLIER * BONUS_CAMEL123, + )], + ); +} + +#[test] +fn umlaut() { + let paths = ["be", "bë"]; + let mut matcher = Matcher::new(Config::DEFAULT); + let matches = Pattern::parse("ë", CaseMatching::Ignore, Normalization::Smart) + .match_list(paths, &mut matcher); + assert_eq!(matches.len(), 1); + let matches = Pattern::parse("e", CaseMatching::Ignore, Normalization::Never) + .match_list(paths, &mut matcher); + assert_eq!(matches.len(), 1); + let matches = Pattern::parse("e", CaseMatching::Ignore, Normalization::Smart) + .match_list(paths, &mut matcher); + assert_eq!(matches.len(), 2); +} diff --git a/crates/atuin-nucleo/matcher/src/utf32_str.rs b/crates/atuin-nucleo/matcher/src/utf32_str.rs new file mode 100644 index 00000000..664dae7a --- /dev/null +++ b/crates/atuin-nucleo/matcher/src/utf32_str.rs @@ -0,0 +1,428 @@ +#[cfg(test)] +mod tests; + +use std::borrow::Cow; +use std::ops::{Bound, RangeBounds}; +use std::{fmt, slice}; + +use memchr::memmem; + +use crate::chars; + +/// Check if a given string can be represented internally as the `Ascii` variant in a +/// [`Utf32String`] or a [`Utf32Str`]. +/// +/// This returns true if the string is ASCII and does not contain a windows-style newline +/// `'\r'`. +/// The additional carriage return check is required since even for strings consisting only +/// of ASCII, the windows-style newline `\r\n` is treated as a single grapheme. +#[inline] +fn has_ascii_graphemes(string: &str) -> bool { + string.is_ascii() && memmem::find(string.as_bytes(), b"\r\n").is_none() +} + +/// A UTF-32 encoded (char array) string that is used as an input to (fuzzy) matching. +/// +/// This is mostly intended as an internal string type, but some methods are exposed for +/// convenience. We make the following API guarantees for `Utf32Str(ing)`s produced from a string +/// using one of its `From<T>` constructors for string types `T` or from the +/// [`Utf32Str::new`] method. +/// +/// 1. The `Ascii` variant contains a byte buffer which is guaranteed to be a valid string +/// slice. +/// 2. It is guaranteed that the string slice internal to the `Ascii` variant is identical +/// to the original string. +/// 3. The length of a `Utf32Str(ing)` is exactly the number of graphemes in the original string. +/// +/// Since `Utf32Str(ing)`s variants may be constructed directly, you **must not** make these +/// assumptions when handling `Utf32Str(ing)`s of unknown origin. +/// +/// ## Caveats +/// Despite the name, this type is quite far from being a true string type. Here are some +/// examples demonstrating this. +/// +/// ### String conversions are not round-trip +/// In the presence of a multi-codepoint grapheme (e.g. `"u\u{0308}"` which is `u + +/// COMBINING_DIAERESIS`), the trailing codepoints are truncated. +/// ``` +/// # use nucleo_matcher::Utf32String; +/// assert_eq!(Utf32String::from("u\u{0308}").to_string(), "u"); +/// ``` +/// +/// ### Indexing is done by grapheme +/// Indexing into a string is done by grapheme rather than by codepoint. +/// ``` +/// # use nucleo_matcher::Utf32String; +/// assert!(Utf32String::from("au\u{0308}").len() == 2); +/// ``` +/// +/// ### A `Unicode` variant may be produced by all-ASCII characters. +/// Since the windows-style newline `\r\n` is ASCII only but considered to be a single grapheme, +/// strings containing `\r\n` will still result in a `Unicode` variant. +/// ``` +/// # use nucleo_matcher::Utf32String; +/// let s = Utf32String::from("\r\n"); +/// assert!(!s.slice(..).is_ascii()); +/// assert!(s.len() == 1); +/// assert!(s.slice(..).get(0) == '\n'); +/// ``` +/// +/// ## Design rationale +/// Usually Rust's UTF-8 encoded strings are great. However, since fuzzy matching +/// operates on codepoints (ideally, it should operate on graphemes but that's too +/// much hassle to deal with), we want to quickly iterate over codepoints (up to 5 +/// times) during matching. +/// +/// Doing codepoint segmentation on the fly not only blows trough the cache +/// (lookup tables and I-cache) but also has nontrivial runtime compared to the +/// matching itself. Furthermore there are many extra optimizations available +/// for ASCII only text, but checking each match has too much overhead. +/// +/// Of course, this comes at extra memory cost as we usually still need the UTF-8 +/// encoded variant for rendering. In the (dominant) case of ASCII-only text +/// we don't require a copy. Furthermore fuzzy matching usually is applied while +/// the user is typing on the fly so the same item is potentially matched many +/// times (making the the up-front cost more worth it). That means that its +/// basically always worth it to pre-segment the string. +/// +/// For usecases that only match (a lot of) strings once its possible to keep +/// char buffer around that is filled with the presegmented chars. +/// +/// Another advantage of this approach is that the matcher will naturally +/// produce grapheme indices (instead of utf8 offsets) anyway. With a +/// codepoint basic representation like this the indices can be used +/// directly +#[derive(PartialEq, Eq, PartialOrd, Ord, Clone, Copy, Hash)] +pub enum Utf32Str<'a> { + /// A string represented as ASCII encoded bytes. + /// Correctness invariant: must only contain valid ASCII (`<= 127`) + Ascii(&'a [u8]), + /// A string represented as an array of unicode codepoints (basically UTF-32). + Unicode(&'a [char]), +} + +impl<'a> Utf32Str<'a> { + /// Convenience method to construct a `Utf32Str` from a normal UTF-8 str + pub fn new(str: &'a str, buf: &'a mut Vec<char>) -> Self { + if has_ascii_graphemes(str) { + Utf32Str::Ascii(str.as_bytes()) + } else { + buf.clear(); + buf.extend(crate::chars::graphemes(str)); + Utf32Str::Unicode(buf) + } + } + + /// Returns the number of characters in this string. + #[inline] + pub fn len(self) -> usize { + match self { + Utf32Str::Unicode(codepoints) => codepoints.len(), + Utf32Str::Ascii(ascii_bytes) => ascii_bytes.len(), + } + } + + /// Returns whether this string is empty. + #[inline] + pub fn is_empty(self) -> bool { + match self { + Utf32Str::Unicode(codepoints) => codepoints.is_empty(), + Utf32Str::Ascii(ascii_bytes) => ascii_bytes.is_empty(), + } + } + + /// Creates a slice with a string that contains the characters in + /// the specified **character range**. + #[inline] + pub fn slice(self, range: impl RangeBounds<usize>) -> Utf32Str<'a> { + let start = match range.start_bound() { + Bound::Included(&start) => start, + Bound::Excluded(&start) => start + 1, + Bound::Unbounded => 0, + }; + let end = match range.end_bound() { + Bound::Included(&end) => end + 1, + Bound::Excluded(&end) => end, + Bound::Unbounded => self.len(), + }; + match self { + Utf32Str::Ascii(bytes) => Utf32Str::Ascii(&bytes[start..end]), + Utf32Str::Unicode(codepoints) => Utf32Str::Unicode(&codepoints[start..end]), + } + } + + /// Returns the number of leading whitespaces in this string + #[inline] + pub(crate) fn leading_white_space(self) -> usize { + match self { + Utf32Str::Ascii(bytes) => bytes + .iter() + .position(|b| !b.is_ascii_whitespace()) + .unwrap_or(0), + Utf32Str::Unicode(codepoints) => codepoints + .iter() + .position(|c| !c.is_whitespace()) + .unwrap_or(0), + } + } + + /// Returns the number of trailing whitespaces in this string + #[inline] + pub(crate) fn trailing_white_space(self) -> usize { + match self { + Utf32Str::Ascii(bytes) => bytes + .iter() + .rev() + .position(|b| !b.is_ascii_whitespace()) + .unwrap_or(0), + Utf32Str::Unicode(codepoints) => codepoints + .iter() + .rev() + .position(|c| !c.is_whitespace()) + .unwrap_or(0), + } + } + + /// Same as `slice` but accepts a u32 range for convenience since + /// those are the indices returned by the matcher. + #[inline] + pub fn slice_u32(self, range: impl RangeBounds<u32>) -> Utf32Str<'a> { + let start = match range.start_bound() { + Bound::Included(&start) => start as usize, + Bound::Excluded(&start) => start as usize + 1, + Bound::Unbounded => 0, + }; + let end = match range.end_bound() { + Bound::Included(&end) => end as usize + 1, + Bound::Excluded(&end) => end as usize, + Bound::Unbounded => self.len(), + }; + match self { + Utf32Str::Ascii(bytes) => Utf32Str::Ascii(&bytes[start..end]), + Utf32Str::Unicode(codepoints) => Utf32Str::Unicode(&codepoints[start..end]), + } + } + + /// Returns whether this string only contains graphemes which are single ASCII chars. + /// + /// This is almost equivalent to the string being ASCII, except with the additional requirement + /// that the string cannot contain a windows-style newline `\r\n` which is treated as a single + /// grapheme. + pub fn is_ascii(self) -> bool { + matches!(self, Utf32Str::Ascii(_)) + } + + /// Returns the `n`th character in this string, zero-indexed + pub fn get(self, n: u32) -> char { + match self { + Utf32Str::Ascii(bytes) => bytes[n as usize] as char, + Utf32Str::Unicode(codepoints) => codepoints[n as usize], + } + } + + /// Returns the last character in this string. + /// + /// Panics if the string is empty. + pub(crate) fn last(self) -> char { + match self { + Utf32Str::Ascii(bytes) => bytes[bytes.len() - 1] as char, + Utf32Str::Unicode(codepoints) => codepoints[codepoints.len() - 1], + } + } + + /// Returns the first character in this string. + /// + /// Panics if the string is empty. + pub(crate) fn first(self) -> char { + match self { + Utf32Str::Ascii(bytes) => bytes[0] as char, + Utf32Str::Unicode(codepoints) => codepoints[0], + } + } + + /// Returns an iterator over the characters in this string + pub fn chars(self) -> Chars<'a> { + match self { + Utf32Str::Ascii(bytes) => Chars::Ascii(bytes.iter()), + Utf32Str::Unicode(codepoints) => Chars::Unicode(codepoints.iter()), + } + } +} + +impl fmt::Debug for Utf32Str<'_> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "\"")?; + for c in self.chars() { + for c in c.escape_debug() { + write!(f, "{c}")? + } + } + write!(f, "\"") + } +} + +impl fmt::Display for Utf32Str<'_> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + for c in self.chars() { + write!(f, "{c}")? + } + Ok(()) + } +} + +pub enum Chars<'a> { + Ascii(slice::Iter<'a, u8>), + Unicode(slice::Iter<'a, char>), +} + +impl Iterator for Chars<'_> { + type Item = char; + + fn next(&mut self) -> Option<Self::Item> { + match self { + Chars::Ascii(iter) => iter.next().map(|&c| c as char), + Chars::Unicode(iter) => iter.next().copied(), + } + } +} + +impl DoubleEndedIterator for Chars<'_> { + fn next_back(&mut self) -> Option<Self::Item> { + match self { + Chars::Ascii(iter) => iter.next_back().map(|&c| c as char), + Chars::Unicode(iter) => iter.next_back().copied(), + } + } +} + +#[derive(PartialEq, Eq, PartialOrd, Ord, Clone, Hash)] +/// An owned version of [`Utf32Str`]. +/// +/// See the API documentation for [`Utf32Str`] for more detail. +pub enum Utf32String { + /// A string represented as ASCII encoded bytes. + /// Correctness invariant: must only contain valid ASCII (<=127) + Ascii(Box<str>), + /// A string represented as an array of unicode codepoints (basically UTF-32). + Unicode(Box<[char]>), +} + +impl Default for Utf32String { + fn default() -> Self { + Self::Ascii(String::new().into_boxed_str()) + } +} + +impl Utf32String { + /// Returns the number of characters in this string. + #[inline] + pub fn len(&self) -> usize { + match self { + Utf32String::Unicode(codepoints) => codepoints.len(), + Utf32String::Ascii(ascii_bytes) => ascii_bytes.len(), + } + } + + /// Returns whether this string is empty. + #[inline] + pub fn is_empty(&self) -> bool { + match self { + Utf32String::Unicode(codepoints) => codepoints.is_empty(), + Utf32String::Ascii(ascii_bytes) => ascii_bytes.is_empty(), + } + } + + /// Creates a slice with a string that contains the characters in + /// the specified **character range**. + #[inline] + pub fn slice(&self, range: impl RangeBounds<usize>) -> Utf32Str { + let start = match range.start_bound() { + Bound::Included(&start) => start, + Bound::Excluded(&start) => start + 1, + Bound::Unbounded => 0, + }; + let end = match range.end_bound() { + Bound::Included(&end) => end + 1, + Bound::Excluded(&end) => end, + Bound::Unbounded => self.len(), + }; + match self { + Utf32String::Ascii(bytes) => Utf32Str::Ascii(&bytes.as_bytes()[start..end]), + Utf32String::Unicode(codepoints) => Utf32Str::Unicode(&codepoints[start..end]), + } + } + + /// Same as `slice` but accepts a u32 range for convenience since + /// those are the indices returned by the matcher. + #[inline] + pub fn slice_u32(&self, range: impl RangeBounds<u32>) -> Utf32Str { + let start = match range.start_bound() { + Bound::Included(&start) => start, + Bound::Excluded(&start) => start + 1, + Bound::Unbounded => 0, + }; + let end = match range.end_bound() { + Bound::Included(&end) => end + 1, + Bound::Excluded(&end) => end, + Bound::Unbounded => self.len() as u32, + }; + match self { + Utf32String::Ascii(bytes) => { + Utf32Str::Ascii(&bytes.as_bytes()[start as usize..end as usize]) + } + Utf32String::Unicode(codepoints) => { + Utf32Str::Unicode(&codepoints[start as usize..end as usize]) + } + } + } +} + +impl From<&str> for Utf32String { + #[inline] + fn from(value: &str) -> Self { + if has_ascii_graphemes(value) { + Self::Ascii(value.to_owned().into_boxed_str()) + } else { + Self::Unicode(chars::graphemes(value).collect()) + } + } +} + +impl From<Box<str>> for Utf32String { + fn from(value: Box<str>) -> Self { + if has_ascii_graphemes(&value) { + Self::Ascii(value) + } else { + Self::Unicode(chars::graphemes(&value).collect()) + } + } +} + +impl From<String> for Utf32String { + #[inline] + fn from(value: String) -> Self { + value.into_boxed_str().into() + } +} + +impl<'a> From<Cow<'a, str>> for Utf32String { + #[inline] + fn from(value: Cow<'a, str>) -> Self { + match value { + Cow::Borrowed(value) => value.into(), + Cow::Owned(value) => value.into(), + } + } +} + +impl fmt::Debug for Utf32String { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self.slice(..)) + } +} + +impl fmt::Display for Utf32String { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{}", self.slice(..)) + } +} diff --git a/crates/atuin-nucleo/matcher/src/utf32_str/tests.rs b/crates/atuin-nucleo/matcher/src/utf32_str/tests.rs new file mode 100644 index 00000000..a38c8875 --- /dev/null +++ b/crates/atuin-nucleo/matcher/src/utf32_str/tests.rs @@ -0,0 +1,44 @@ +use crate::{Utf32Str, Utf32String}; + +#[test] +fn test_utf32str_ascii() { + /// Helper function for testing + fn expect_ascii(src: &str, is_ascii: bool) { + let mut buffer = Vec::new(); + assert!(Utf32Str::new(src, &mut buffer).is_ascii() == is_ascii); + assert!(Utf32String::from(src).slice(..).is_ascii() == is_ascii); + assert!(Utf32String::from(src.to_owned()).slice(..).is_ascii() == is_ascii); + } + + // ascii + expect_ascii("", true); + expect_ascii("a", true); + expect_ascii("a\nb", true); + expect_ascii("\n\r", true); + + // not ascii + expect_ascii("aü", false); + expect_ascii("au\u{0308}", false); + + // windows-style newline + expect_ascii("a\r\nb", false); + expect_ascii("ü\r\n", false); + expect_ascii("\r\n", false); +} + +#[test] +fn test_grapheme_truncation() { + // ascii is preserved + let s = Utf32String::from("ab"); + assert_eq!(s.slice(..).get(0), 'a'); + assert_eq!(s.slice(..).get(1), 'b'); + + // windows-style newline is truncated to '\n' + let s = Utf32String::from("\r\n"); + assert_eq!(s.slice(..).get(0), '\n'); + + // normal graphemes are truncated to the first character + let s = Utf32String::from("u\u{0308}\r\n"); + assert_eq!(s.slice(..).get(0), 'u'); + assert_eq!(s.slice(..).get(1), '\n'); +} diff --git a/crates/atuin-nucleo/src/boxcar.rs b/crates/atuin-nucleo/src/boxcar.rs new file mode 100644 index 00000000..9b48809d --- /dev/null +++ b/crates/atuin-nucleo/src/boxcar.rs @@ -0,0 +1,786 @@ +//! Adapted from the `boxcar` crate at <https://github.com/ibraheemdev/boxcar/blob/master/src/raw.rs> +//! under MIT licenes: +//! +//! Copyright (c) 2022 Ibraheem Ahmed +//! +//! Permission is hereby granted, free of charge, to any person obtaining a copy +//! of this software and associated documentation files (the "Software"), to deal +//! in the Software without restriction, including without limitation the rights +//! to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +//! copies of the Software, and to permit persons to whom the Software is +//! furnished to do so, subject to the following conditions: +//! +//! The above copyright notice and this permission notice shall be included in all +//! copies or substantial portions of the Software. +//! +//! THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +//! IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +//! FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +//! AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +//! LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +//! OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +//! SOFTWARE. + +use std::alloc::Layout; +use std::cell::UnsafeCell; +use std::fmt::Debug; +use std::mem::MaybeUninit; +use std::sync::atomic::{AtomicBool, AtomicPtr, AtomicU64, Ordering}; +use std::{ptr, slice}; + +use crate::{Item, Utf32String}; + +const BUCKETS: u32 = u32::BITS - SKIP_BUCKET; +const MAX_ENTRIES: u32 = u32::MAX - SKIP; + +/// A lock-free, append-only vector. +pub(crate) struct Vec<T> { + /// a counter used to retrieve a unique index to push to. + /// + /// this value may be more than the true length as it will + /// be incremented before values are actually stored. + inflight: AtomicU64, + /// buckets of length 32, 64 .. 2^31 + buckets: [Bucket<T>; BUCKETS as usize], + /// the number of matcher columns in this vector, its absolutely critical that + /// this remains constant and after initilaziaton (safety invariant) since + /// it is used to calculate the Entry layout + columns: u32, +} + +impl<T> Vec<T> { + /// Constructs a new, empty `Vec<T>` with the specified capacity and matcher columns. + pub fn with_capacity(capacity: u32, columns: u32) -> Vec<T> { + assert_ne!(columns, 0, "there must be atleast one matcher column"); + let init = match capacity { + 0 => 0, + // initialize enough buckets for `capacity` elements + n => Location::of(n).bucket, + }; + + let mut buckets = [ptr::null_mut(); BUCKETS as usize]; + + for (i, bucket) in buckets[..=init as usize].iter_mut().enumerate() { + let len = Location::bucket_len(i as u32); + *bucket = unsafe { Bucket::alloc(len, columns) }; + } + + Vec { + buckets: buckets.map(Bucket::new), + inflight: AtomicU64::new(0), + columns, + } + } + pub fn columns(&self) -> u32 { + self.columns + } + + /// Returns the number of elements in the vector. + #[inline] + pub fn count(&self) -> u32 { + self.inflight + .load(Ordering::Acquire) + .min(MAX_ENTRIES as u64) as u32 + } + + // Returns a reference to the element at the given index. + // + // # Safety + // + // Entry at `index` must be initialized. + #[inline] + pub unsafe fn get_unchecked(&self, index: u32) -> Item<'_, T> { + let location = Location::of(index); + + unsafe { + let entries = self + .buckets + .get_unchecked(location.bucket as usize) + .entries + .load(Ordering::Relaxed); + debug_assert!(!entries.is_null()); + let entry = Bucket::<T>::get(entries, location.entry, self.columns); + // this looks odd but is necessary to ensure cross + // thread synchronization (essentially acting as a memory barrier) + // since the caller must only guarantee that he has observed active on any thread + // but the current thread might still have an old value cached (although unlikely) + let _ = (*entry).active.load(Ordering::Acquire); + Entry::read(entry, self.columns) + } + } + + /// Returns a reference to the element at the given index. + pub fn get(&self, index: u32) -> Option<Item<'_, T>> { + let location = Location::of(index); + + unsafe { + // safety: `location.bucket` is always in bounds + let entries = self + .buckets + .get_unchecked(location.bucket as usize) + .entries + .load(Ordering::Relaxed); + + // bucket is uninitialized + if entries.is_null() { + return None; + } + + // safety: `location.entry` is always in bounds for it's bucket + let entry = Bucket::<T>::get(entries, location.entry, self.columns); + + // safety: the entry is active + (*entry) + .active + .load(Ordering::Acquire) + .then(|| Entry::read(entry, self.columns)) + } + } + + /// Appends an element to the back of the vector. + pub fn push(&self, value: T, fill_columns: impl FnOnce(&T, &mut [Utf32String])) -> u32 { + let index = self.inflight.fetch_add(1, Ordering::Release); + // the inflight counter is a `u64` to catch overflows of the vector'scapacity + let index: u32 = index.try_into().expect("overflowed maximum capacity"); + let location = Location::of(index); + + // eagerly allocate the next bucket if we are close to the end of this one + if index == (location.bucket_len - (location.bucket_len >> 3)) { + if let Some(next_bucket) = self.buckets.get(location.bucket as usize + 1) { + Vec::get_or_alloc(next_bucket, location.bucket_len << 1, self.columns); + } + } + + // safety: `location.bucket` is always in bounds + let bucket = unsafe { self.buckets.get_unchecked(location.bucket as usize) }; + let mut entries = bucket.entries.load(Ordering::Acquire); + + // the bucket has not been allocated yet + if entries.is_null() { + entries = Vec::get_or_alloc(bucket, location.bucket_len, self.columns); + } + + unsafe { + // safety: `location.entry` is always in bounds for it's bucket + let entry = Bucket::get(entries, location.entry, self.columns); + + // safety: we have unique access to this entry. + // + // 1. it is impossible for another thread to attempt a `push` + // to this location as we retrieved it from `inflight.fetch_add` + // + // 2. any thread trying to `get` this entry will see `active == false`, + // and will not try to access it + for col in Entry::matcher_cols_raw(entry, self.columns) { + col.get().write(MaybeUninit::new(Utf32String::default())) + } + fill_columns(&value, Entry::matcher_cols_mut(entry, self.columns)); + (*entry).slot.get().write(MaybeUninit::new(value)); + // let other threads know that this entry is active + (*entry).active.store(true, Ordering::Release); + } + + index + } + + /// Extends the vector by appending multiple elements at once. + pub fn extend<I>(&self, values: I, fill_columns: impl Fn(&T, &mut [Utf32String])) + where + I: IntoIterator<Item = T> + ExactSizeIterator, + { + let count: u32 = values + .len() + .try_into() + .expect("overflowed maximum capacity"); + if count == 0 { + assert!( + values.into_iter().next().is_none(), + "The `values` variable reported incorrect length." + ); + return; + } + + // Reserve all indices at once + let start_index: u32 = self + .inflight + .fetch_add(u64::from(count), Ordering::Release) + .try_into() + .expect("overflowed maximum capacity"); + + // Compute first and last locations + let start_location = Location::of(start_index); + let end_location = Location::of(start_index + count); + + // Eagerly allocate the next bucket if the last entry is close to the end of its next bucket + let alloc_entry = end_location.alloc_next_bucket_entry(); + if end_location.entry >= alloc_entry + && (start_location.bucket != end_location.bucket || start_location.entry <= alloc_entry) + { + // This might be the last bucket, hence the check + if let Some(next_bucket) = self.buckets.get(end_location.bucket as usize + 1) { + Vec::get_or_alloc(next_bucket, end_location.bucket_len << 1, self.columns); + } + } + + let mut bucket = unsafe { self.buckets.get_unchecked(start_location.bucket as usize) }; + let mut entries = bucket.entries.load(Ordering::Acquire); + if entries.is_null() { + entries = Vec::get_or_alloc( + bucket, + Location::bucket_len(start_location.bucket), + self.columns, + ); + } + // Route each value to its corresponding bucket + let mut location; + let count = count as usize; + for (i, v) in values.into_iter().enumerate() { + // ExactSizeIterator is a safe trait that can have bugs/lie about it's size. + // Unsafe code cannot rely on the reported length being correct. + assert!(i < count); + + location = + Location::of(start_index + u32::try_from(i).expect("overflowed maximum capacity")); + + // if we're starting to insert into a different bucket, allocate it beforehand + if location.entry == 0 && i != 0 { + // safety: `location.bucket` is always in bounds + bucket = unsafe { self.buckets.get_unchecked(location.bucket as usize) }; + entries = bucket.entries.load(Ordering::Acquire); + + if entries.is_null() { + entries = Vec::get_or_alloc( + bucket, + Location::bucket_len(location.bucket), + self.columns, + ); + } + } + + unsafe { + let entry = Bucket::get(entries, location.entry, self.columns); + + // Initialize matcher columns + for col in Entry::matcher_cols_raw(entry, self.columns) { + col.get().write(MaybeUninit::new(Utf32String::default())); + } + fill_columns(&v, Entry::matcher_cols_mut(entry, self.columns)); + (*entry).slot.get().write(MaybeUninit::new(v)); + (*entry).active.store(true, Ordering::Release); + } + } + } + + /// race to initialize a bucket + fn get_or_alloc(bucket: &Bucket<T>, len: u32, cols: u32) -> *mut Entry<T> { + let entries = unsafe { Bucket::alloc(len, cols) }; + match bucket.entries.compare_exchange( + ptr::null_mut(), + entries, + Ordering::Release, + Ordering::Acquire, + ) { + Ok(_) => entries, + Err(found) => unsafe { + Bucket::dealloc(entries, len, cols); + found + }, + } + } + + /// Returns an iterator over the vector starting at `start` + /// the iterator is deterministically sized and will not grow + /// as more elements are pushed + pub unsafe fn snapshot(&self, start: u32) -> Iter<'_, T> { + let end = self + .inflight + .load(Ordering::Acquire) + .min(MAX_ENTRIES as u64) as u32; + assert!(start <= end, "index {start} is out of bounds!"); + Iter { + location: Location::of(start), + vec: self, + idx: start, + end, + } + } + + /// Returns an iterator over the vector starting at `start` + /// the iterator is deterministically sized and will not grow + /// as more elements are pushed + pub unsafe fn par_snapshot(&self, start: u32) -> ParIter<'_, T> { + let end = self + .inflight + .load(Ordering::Acquire) + .min(MAX_ENTRIES as u64) as u32; + assert!(start <= end, "index {start} is out of bounds!"); + + ParIter { + start, + end, + vec: self, + } + } +} + +impl<T> Drop for Vec<T> { + fn drop(&mut self) { + for (i, bucket) in self.buckets.iter_mut().enumerate() { + let entries = *bucket.entries.get_mut(); + + if entries.is_null() { + break; + } + + let len = Location::bucket_len(i as u32); + // safety: in drop + unsafe { Bucket::dealloc(entries, len, self.columns) } + } + } +} +type SnapshotItem<'v, T> = (u32, Option<Item<'v, T>>); + +pub struct Iter<'v, T> { + location: Location, + idx: u32, + end: u32, + vec: &'v Vec<T>, +} +impl<T> Iter<'_, T> { + pub fn end(&self) -> u32 { + self.end + } +} + +impl<'v, T> Iterator for Iter<'v, T> { + type Item = SnapshotItem<'v, T>; + fn size_hint(&self) -> (usize, Option<usize>) { + ( + (self.end - self.idx) as usize, + Some((self.end - self.idx) as usize), + ) + } + + fn next(&mut self) -> Option<SnapshotItem<'v, T>> { + if self.end == self.idx { + return None; + } + debug_assert!(self.idx < self.end, "huh {} {}", self.idx, self.end); + debug_assert!(self.end as u64 <= self.vec.inflight.load(Ordering::Relaxed)); + + loop { + let entries = unsafe { + self.vec + .buckets + .get_unchecked(self.location.bucket as usize) + .entries + .load(Ordering::Relaxed) + }; + debug_assert!(self.location.bucket < BUCKETS); + + if self.location.entry < self.location.bucket_len { + if entries.is_null() { + // we still want to yield these + let index = self.idx; + self.location.entry += 1; + self.idx += 1; + return Some((index, None)); + } + // safety: bounds and null checked above + let entry = unsafe { Bucket::get(entries, self.location.entry, self.vec.columns) }; + let index = self.idx; + self.location.entry += 1; + self.idx += 1; + + let entry = unsafe { + (*entry) + .active + .load(Ordering::Acquire) + .then(|| Entry::read(entry, self.vec.columns)) + }; + return Some((index, entry)); + } + + self.location.entry = 0; + self.location.bucket += 1; + + if self.location.bucket < BUCKETS { + self.location.bucket_len = Location::bucket_len(self.location.bucket); + } + } + } +} +impl<T> ExactSizeIterator for Iter<'_, T> {} +impl<T> DoubleEndedIterator for Iter<'_, T> { + fn next_back(&mut self) -> Option<Self::Item> { + unimplemented!() + } +} + +pub struct ParIter<'v, T> { + end: u32, + start: u32, + vec: &'v Vec<T>, +} +impl<T> ParIter<'_, T> { + pub fn end(&self) -> u32 { + self.end + } +} + +impl<'v, T: Send + Sync> rayon::iter::ParallelIterator for ParIter<'v, T> { + type Item = SnapshotItem<'v, T>; + + fn drive_unindexed<C>(self, consumer: C) -> C::Result + where + C: rayon::iter::plumbing::UnindexedConsumer<Self::Item>, + { + rayon::iter::plumbing::bridge(self, consumer) + } + + fn opt_len(&self) -> Option<usize> { + Some((self.end - self.start) as usize) + } +} + +impl<T: Send + Sync> rayon::iter::IndexedParallelIterator for ParIter<'_, T> { + fn len(&self) -> usize { + (self.end - self.start) as usize + } + + fn drive<C: rayon::iter::plumbing::Consumer<Self::Item>>(self, consumer: C) -> C::Result { + rayon::iter::plumbing::bridge(self, consumer) + } + + fn with_producer<CB>(self, callback: CB) -> CB::Output + where + CB: rayon::iter::plumbing::ProducerCallback<Self::Item>, + { + callback.callback(ParIterProducer { + start: self.start, + end: self.end, + vec: self.vec, + }) + } +} + +struct ParIterProducer<'v, T: Send> { + start: u32, + end: u32, + vec: &'v Vec<T>, +} + +impl<'v, T: 'v + Send + Sync> rayon::iter::plumbing::Producer for ParIterProducer<'v, T> { + type Item = SnapshotItem<'v, T>; + type IntoIter = Iter<'v, T>; + + fn into_iter(self) -> Self::IntoIter { + debug_assert!(self.start <= self.end); + Iter { + location: Location::of(self.start), + idx: self.start, + end: self.end, + vec: self.vec, + } + } + + fn split_at(self, index: usize) -> (Self, Self) { + assert!(index <= (self.end - self.start) as usize); + let index = index as u32; + ( + ParIterProducer { + start: self.start, + end: self.start + index, + vec: self.vec, + }, + ParIterProducer { + start: self.start + index, + end: self.end, + vec: self.vec, + }, + ) + } +} + +struct Bucket<T> { + entries: AtomicPtr<Entry<T>>, +} + +impl<T> Bucket<T> { + fn layout(len: u32, layout: Layout) -> Layout { + Layout::from_size_align(layout.size() * len as usize, layout.align()) + .expect("exceeded maximum allocation size") + } + + unsafe fn alloc(len: u32, cols: u32) -> *mut Entry<T> { + let layout = Entry::<T>::layout(cols); + let arr_layout = Self::layout(len, layout); + let entries = std::alloc::alloc(arr_layout); + if entries.is_null() { + std::alloc::handle_alloc_error(arr_layout) + } + + for i in 0..len { + let active = entries.add(i as usize * layout.size()) as *mut AtomicBool; + active.write(AtomicBool::new(false)) + } + entries as *mut Entry<T> + } + + unsafe fn dealloc(entries: *mut Entry<T>, len: u32, cols: u32) { + let layout = Entry::<T>::layout(cols); + let arr_layout = Self::layout(len, layout); + for i in 0..len { + let entry = Bucket::get(entries, i, cols); + if *(*entry).active.get_mut() { + ptr::drop_in_place((*(*entry).slot.get()).as_mut_ptr()); + for matcher_col in Entry::matcher_cols_raw(entry, cols) { + ptr::drop_in_place((*matcher_col.get()).as_mut_ptr()); + } + } + } + std::alloc::dealloc(entries as *mut u8, arr_layout) + } + + unsafe fn get(entries: *mut Entry<T>, idx: u32, cols: u32) -> *mut Entry<T> { + let layout = Entry::<T>::layout(cols); + let ptr = entries as *mut u8; + ptr.add(layout.size() * idx as usize) as *mut Entry<T> + } + + fn new(entries: *mut Entry<T>) -> Bucket<T> { + Bucket { + entries: AtomicPtr::new(entries), + } + } +} + +#[repr(C)] +struct Entry<T> { + active: AtomicBool, + slot: UnsafeCell<MaybeUninit<T>>, + tail: [UnsafeCell<MaybeUninit<Utf32String>>; 0], +} + +impl<T> Entry<T> { + fn layout(cols: u32) -> Layout { + let head = Layout::new::<Self>(); + let tail = Layout::array::<Utf32String>(cols as usize).expect("invalid memory layout"); + head.extend(tail) + .expect("invalid memory layout") + .0 + .pad_to_align() + } + + unsafe fn matcher_cols_raw<'a>( + ptr: *mut Entry<T>, + cols: u32, + ) -> &'a [UnsafeCell<MaybeUninit<Utf32String>>] { + // this whole thing looks weird. The reason we do this is that + // we must make sure the pointer retains its provenance which may (or may not?) + // be lost if we used tail.as_ptr() + let tail = std::ptr::addr_of!((*ptr).tail) as *const u8; + let offset = tail.offset_from(ptr as *mut u8) as usize; + let ptr = (ptr as *mut u8).add(offset) as *mut _; + slice::from_raw_parts(ptr, cols as usize) + } + + unsafe fn matcher_cols_mut<'a>(ptr: *mut Entry<T>, cols: u32) -> &'a mut [Utf32String] { + // this whole thing looks weird. The reason we do this is that + // we must make sure the pointer retains its provenance which may (or may not?) + // be lost if we used tail.as_ptr() + let tail = std::ptr::addr_of!((*ptr).tail) as *const u8; + let offset = tail.offset_from(ptr as *mut u8) as usize; + let ptr = (ptr as *mut u8).add(offset) as *mut _; + slice::from_raw_parts_mut(ptr, cols as usize) + } + // # Safety + // + // Value must be initialized. + unsafe fn read<'a>(ptr: *mut Entry<T>, cols: u32) -> Item<'a, T> { + // this whole thing looks weird. The reason we do this is that + // we must make sure the pointer retains its provenance which may (or may not?) + // be lost if we used tail.as_ptr() + let data = (*(*ptr).slot.get()).assume_init_ref(); + let tail = std::ptr::addr_of!((*ptr).tail) as *const u8; + let offset = tail.offset_from(ptr as *mut u8) as usize; + let ptr = (ptr as *mut u8).add(offset) as *mut _; + let matcher_columns = slice::from_raw_parts(ptr, cols as usize); + Item { + data, + matcher_columns, + } + } +} + +#[derive(Debug)] +struct Location { + // the index of the bucket + bucket: u32, + // the length of `bucket` + bucket_len: u32, + // the index of the entry in `bucket` + entry: u32, +} + +// skip the shorter buckets to avoid unnecessary allocations. +// this also reduces the maximum capacity of a vector. +const SKIP: u32 = 32; +const SKIP_BUCKET: u32 = (u32::BITS - SKIP.leading_zeros()) - 1; + +impl Location { + fn of(index: u32) -> Location { + let skipped = index.checked_add(SKIP).expect("exceeded maximum length"); + let bucket = u32::BITS - skipped.leading_zeros(); + let bucket = bucket - (SKIP_BUCKET + 1); + let bucket_len = Location::bucket_len(bucket); + let entry = skipped ^ bucket_len; + + Location { + bucket, + bucket_len, + entry, + } + } + + fn bucket_len(bucket: u32) -> u32 { + 1 << (bucket + SKIP_BUCKET) + } + + /// The entry index at which the next bucket should be pre-allocated. + fn alloc_next_bucket_entry(&self) -> u32 { + self.bucket_len - (self.bucket_len >> 3) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn location() { + assert_eq!(Location::bucket_len(0), 32); + for i in 0..32 { + let loc = Location::of(i); + assert_eq!(loc.bucket_len, 32); + assert_eq!(loc.bucket, 0); + assert_eq!(loc.entry, i); + } + + assert_eq!(Location::bucket_len(1), 64); + for i in 33..96 { + let loc = Location::of(i); + assert_eq!(loc.bucket_len, 64); + assert_eq!(loc.bucket, 1); + assert_eq!(loc.entry, i - 32); + } + + assert_eq!(Location::bucket_len(2), 128); + for i in 96..224 { + let loc = Location::of(i); + assert_eq!(loc.bucket_len, 128); + assert_eq!(loc.bucket, 2); + assert_eq!(loc.entry, i - 96); + } + + let max = Location::of(MAX_ENTRIES); + assert_eq!(max.bucket, BUCKETS - 1); + assert_eq!(max.bucket_len, 1 << 31); + assert_eq!(max.entry, (1 << 31) - 1); + } + + #[test] + fn extend_unique_bucket() { + let vec = Vec::<u32>::with_capacity(1, 1); + vec.extend(0..10, |_, _| {}); + assert_eq!(vec.count(), 10); + for i in 0..10 { + assert_eq!(*vec.get(i).unwrap().data, i); + } + assert!(vec.get(10).is_none()); + } + + #[test] + fn extend_over_two_buckets() { + let vec = Vec::<u32>::with_capacity(1, 1); + vec.extend(0..100, |_, _| {}); + assert_eq!(vec.count(), 100); + for i in 0..100 { + assert_eq!(*vec.get(i).unwrap().data, i); + } + assert!(vec.get(100).is_none()); + } + + #[test] + fn extend_over_more_than_two_buckets() { + let vec = Vec::<u32>::with_capacity(1, 1); + vec.extend(0..1000, |_, _| {}); + assert_eq!(vec.count(), 1000); + for i in 0..1000 { + assert_eq!(*vec.get(i).unwrap().data, i); + } + assert!(vec.get(1000).is_none()); + } + + #[test] + /// test that ExactSizeIterator returning incorrect length is caught (0 AND more than reported) + fn extend_with_incorrect_reported_len_is_caught() { + struct IncorrectLenIter { + len: usize, + iter: std::ops::Range<u32>, + } + + impl Iterator for IncorrectLenIter { + type Item = u32; + + fn next(&mut self) -> Option<Self::Item> { + self.iter.next() + } + } + + impl ExactSizeIterator for IncorrectLenIter { + fn len(&self) -> usize { + self.len + } + } + + let vec = Vec::<u32>::with_capacity(1, 1); + let iter = IncorrectLenIter { + len: 10, + iter: (0..12), + }; + // this should panic + assert!(std::panic::catch_unwind(|| vec.extend(iter, |_, _| {})).is_err()); + + let vec = Vec::<u32>::with_capacity(1, 1); + let iter = IncorrectLenIter { + len: 12, + iter: (0..10), + }; + // this shouldn't panic and should just ignore the extra elements + assert!(std::panic::catch_unwind(|| vec.extend(iter, |_, _| {})).is_ok()); + // we should reserve 12 elements but only 10 should be present + assert_eq!(vec.count(), 12); + for i in 0..10 { + assert_eq!(*vec.get(i).unwrap().data, i); + } + assert!(vec.get(10).is_none()); + + let vec = Vec::<u32>::with_capacity(1, 1); + let iter = IncorrectLenIter { + len: 0, + iter: (0..2), + }; + // this should panic + assert!(std::panic::catch_unwind(|| vec.extend(iter, |_, _| {})).is_err()); + } + + // test |values| does not fit in the boxcar + #[test] + fn extend_over_max_capacity() { + let vec = Vec::<u32>::with_capacity(1, 1); + let count = MAX_ENTRIES as usize + 2; + let iter = std::iter::repeat(0).take(count); + assert!(std::panic::catch_unwind(|| vec.extend(iter, |_, _| {})).is_err()); + } +} diff --git a/crates/atuin-nucleo/src/lib.rs b/crates/atuin-nucleo/src/lib.rs new file mode 100644 index 00000000..7ddb7407 --- /dev/null +++ b/crates/atuin-nucleo/src/lib.rs @@ -0,0 +1,462 @@ +/*! +`nucleo` is a high level crate that provides a high level matcher API that +provides a highly effective (parallel) matcher worker. It's designed to allow +quickly plugging a fully featured (and faster) fzf/skim like fuzzy matcher into +your TUI application. + +It's designed to run matching on a background threadpool while providing a +snapshot of the last complete match. That means the matcher can update the +results live while the user is typing while never blocking the main UI thread +(beyond a user provided timeout). Nucleo also supports fully concurrent lock-free +(and wait-free) streaming of input items. + +The [`Nucleo`] struct serves as the main API entrypoint for this crate. + +# Status + +Nucleo is used in the helix-editor and therefore has a large user base with lots +or real world testing. The core matcher implementation is considered complete +and is unlikely to see major changes. The `nucleo-matcher` crate is finished and +ready for widespread use, breaking changes should be very rare (a 1.0 release +should not be far away). + +While the high level `nucleo` crate also works well (and is also used in helix), +there are still additional features that will be added in the future. The high +level crate also need better documentation and will likely see a few minor API +changes in the future. + +*/ +use std::ops::{Bound, RangeBounds}; +use std::sync::atomic::{self, AtomicBool, Ordering}; +use std::sync::Arc; +use std::time::Duration; + +use parking_lot::Mutex; +use rayon::ThreadPool; + +use crate::pattern::MultiPattern; +use crate::worker::Worker; +pub use nucleo_matcher::{chars, Config, Matcher, Utf32Str, Utf32String}; + +mod boxcar; +mod par_sort; +pub mod pattern; +mod worker; + +#[cfg(test)] +mod tests; + +/// A match candidate stored in a [`Nucleo`] worker. +pub struct Item<'a, T> { + pub data: &'a T, + pub matcher_columns: &'a [Utf32String], +} + +/// A handle that allows adding new items to a [`Nucleo`] worker. +/// +/// It's internally reference counted and can be cheaply cloned +/// and sent across threads. +pub struct Injector<T> { + items: Arc<boxcar::Vec<T>>, + notify: Arc<(dyn Fn() + Sync + Send)>, +} + +impl<T> Clone for Injector<T> { + fn clone(&self) -> Self { + Injector { + items: self.items.clone(), + notify: self.notify.clone(), + } + } +} + +impl<T> Injector<T> { + /// Appends an element to the list of matched items. + /// This function is lock-free and wait-free. + pub fn push(&self, value: T, fill_columns: impl FnOnce(&T, &mut [Utf32String])) -> u32 { + let idx = self.items.push(value, fill_columns); + (self.notify)(); + idx + } + + /// Appends multiple elements to the list of matched items. + /// This function is lock-free and wait-free. + /// + /// You should favor this function over `push` if at least one of the following is true: + /// - the number of items you're adding can be computed beforehand and is typically larger + /// than 1k + /// - you're able to batch incoming items + /// - you're adding items from multiple threads concurrently (this function results in less + /// contention) + pub fn extend<I>(&self, values: I, fill_columns: impl Fn(&T, &mut [Utf32String])) + where + I: IntoIterator<Item = T> + ExactSizeIterator, + { + self.items.extend(values, fill_columns); + (self.notify)(); + } + + /// Returns the total number of items injected in the matcher. This might + /// not match the number of items in the match snapshot (if the matcher + /// is still running) + pub fn injected_items(&self) -> u32 { + self.items.count() + } + + /// Returns a reference to the item at the given index. + /// + /// # Safety + /// + /// Item at `index` must be initialized. That means you must have observed + /// `push` returning this value or `get` returning `Some` for this value. + /// Just because a later index is initialized doesn't mean that this index + /// is initialized + pub unsafe fn get_unchecked(&self, index: u32) -> Item<'_, T> { + self.items.get_unchecked(index) + } + + /// Returns a reference to the element at the given index. + pub fn get(&self, index: u32) -> Option<Item<'_, T>> { + self.items.get(index) + } +} + +/// An [item](crate::Item) that was successfully matched by a [`Nucleo`] worker. +#[derive(PartialEq, Eq, Debug, Clone, Copy)] +pub struct Match { + pub score: u32, + pub idx: u32, +} + +/// That status of a [`Nucleo`] worker after a match. +#[derive(PartialEq, Eq, Debug, Clone, Copy)] +pub struct Status { + /// Whether the current snapshot has changed. + pub changed: bool, + /// Whether the matcher is still processing in the background. + pub running: bool, +} + +/// A snapshot represent the results of a [`Nucleo`] worker after +/// finishing a [`tick`](Nucleo::tick). +pub struct Snapshot<T: Sync + Send + 'static> { + item_count: u32, + matches: Vec<Match>, + pattern: MultiPattern, + items: Arc<boxcar::Vec<T>>, +} + +impl<T: Sync + Send + 'static> Snapshot<T> { + fn clear(&mut self, new_items: Arc<boxcar::Vec<T>>) { + self.item_count = 0; + self.matches.clear(); + self.items = new_items + } + + fn update(&mut self, worker: &Worker<T>) { + self.item_count = worker.item_count(); + self.pattern.clone_from(&worker.pattern); + self.matches.clone_from(&worker.matches); + if !Arc::ptr_eq(&worker.items, &self.items) { + self.items = worker.items.clone() + } + } + + /// Returns that total number of items + pub fn item_count(&self) -> u32 { + self.item_count + } + + /// Returns the pattern which items were matched against + pub fn pattern(&self) -> &MultiPattern { + &self.pattern + } + + /// Returns that number of items that matched the pattern + pub fn matched_item_count(&self) -> u32 { + self.matches.len() as u32 + } + + /// Returns an iterator over the items that correspond to a subrange of + /// all the matches in this snapshot. + /// + /// # Panics + /// Panics if `range` has a range bound that is larger than + /// the matched item count + pub fn matched_items( + &self, + range: impl RangeBounds<u32>, + ) -> impl ExactSizeIterator<Item = Item<'_, T>> + DoubleEndedIterator + '_ { + // TODO: use TAIT + let start = match range.start_bound() { + Bound::Included(&start) => start as usize, + Bound::Excluded(&start) => start as usize + 1, + Bound::Unbounded => 0, + }; + let end = match range.end_bound() { + Bound::Included(&end) => end as usize + 1, + Bound::Excluded(&end) => end as usize, + Bound::Unbounded => self.matches.len(), + }; + self.matches[start..end] + .iter() + .map(|&m| unsafe { self.items.get_unchecked(m.idx) }) + } + + /// Returns a reference to the item at the given index. + /// + /// # Safety + /// + /// Item at `index` must be initialized. That means you must have observed a + /// match with the corresponding index in this exact snapshot. Observing + /// a higher index is not enough as item indices can be non-contigously + /// initialized + #[inline] + pub unsafe fn get_item_unchecked(&self, index: u32) -> Item<'_, T> { + self.items.get_unchecked(index) + } + + /// Returns a reference to the item at the given index. + /// + /// Returns `None` if the given `index` is not initialized. This function + /// is only guarteed to return `Some` for item indices that can be found in + /// the `matches` of this struct. Both smaller and larger indices may return + /// `None`. + #[inline] + pub fn get_item(&self, index: u32) -> Option<Item<'_, T>> { + self.items.get(index) + } + + /// Return the matches corresponding to this snapshot. + #[inline] + pub fn matches(&self) -> &[Match] { + &self.matches + } + + /// A convenience function to return the [`Item`] corresponding to the + /// `n`th match. + /// + /// Returns `None` if `n` is greater than or equal to the match count. + #[inline] + pub fn get_matched_item(&self, n: u32) -> Option<Item<'_, T>> { + // SAFETY: A match index is guaranteed to corresponding to a valid global index in this + // snapshot. + unsafe { Some(self.get_item_unchecked(self.matches.get(n as usize)?.idx)) } + } +} + +#[repr(u8)] +#[derive(Clone, Copy, PartialEq, Eq)] +enum State { + Init, + /// items have been cleared but snapshot and items are still outdated + Cleared, + /// items are fresh + Fresh, +} + +impl State { + fn matcher_item_refs(self) -> usize { + match self { + State::Cleared => 1, + State::Init | State::Fresh => 2, + } + } + + fn canceled(self) -> bool { + self != State::Fresh + } + + fn cleared(self) -> bool { + self != State::Fresh + } +} + +/// A high level matcher worker that quickly computes matches in a background +/// threadpool. +pub struct Nucleo<T: Sync + Send + 'static> { + // the way the API is build we totally don't actually need these to be Arcs + // but this lets us avoid some unsafe + canceled: Arc<AtomicBool>, + should_notify: Arc<AtomicBool>, + worker: Arc<Mutex<Worker<T>>>, + pool: ThreadPool, + state: State, + items: Arc<boxcar::Vec<T>>, + notify: Arc<(dyn Fn() + Sync + Send)>, + snapshot: Snapshot<T>, + /// The pattern matched by this matcher. To update the match pattern + /// [`MultiPattern::reparse`](`pattern::MultiPattern::reparse`) should be used. + /// Note that the matcher worker will only become aware of the new pattern + /// after a call to [`tick`](Nucleo::tick). + pub pattern: MultiPattern, +} + +impl<T: Sync + Send + 'static> Nucleo<T> { + /// Constructs a new `nucleo` worker threadpool with the provided `config`. + /// + /// `notify` is called everytime new information is available and + /// [`tick`](Nucleo::tick) should be called. Note that `notify` is not + /// debounced, that should be handled by the downstream crate (for example + /// debouncing to only redraw at most every 1/60 seconds). + /// + /// If `None` is passed for the number of worker threads, nucleo will use + /// one thread per hardware thread. + /// + /// Nucleo can match items with multiple orthogonal properties. `columns` + /// indicates how many matching columns each item (and the pattern) has. The + /// number of columns cannot be changed after construction. + pub fn new( + config: Config, + notify: Arc<(dyn Fn() + Sync + Send)>, + num_threads: Option<usize>, + columns: u32, + ) -> Self { + let (pool, worker) = Worker::new(num_threads, config, notify.clone(), columns); + Self { + canceled: worker.canceled.clone(), + should_notify: worker.should_notify.clone(), + items: worker.items.clone(), + pool, + pattern: MultiPattern::new(columns as usize), + snapshot: Snapshot { + matches: Vec::with_capacity(2 * 1024), + pattern: MultiPattern::new(columns as usize), + item_count: 0, + items: worker.items.clone(), + }, + worker: Arc::new(Mutex::new(worker)), + state: State::Init, + notify, + } + } + + /// Returns the total number of active injectors + pub fn active_injectors(&self) -> usize { + Arc::strong_count(&self.items) + - self.state.matcher_item_refs() + - (Arc::ptr_eq(&self.snapshot.items, &self.items)) as usize + } + + /// Returns a snapshot of the current matcher state. + pub fn snapshot(&self) -> &Snapshot<T> { + &self.snapshot + } + + /// Returns an injector that can be used for adding candidates to the matcher. + pub fn injector(&self) -> Injector<T> { + Injector { + items: self.items.clone(), + notify: self.notify.clone(), + } + } + + /// Restart the the item stream. Removes all items and disconnects all + /// previously created injectors from this instance. If `clear_snapshot` + /// is `true` then all items and matched are removed from the [`Snapshot`] + /// immediately. Otherwise the snapshot will keep the current matches until + /// the matcher has run again. + /// + /// # Note + /// + /// The injectors will continue to function but they will not affect this + /// instance anymore. The old items will only be dropped when all injectors + /// were dropped. + pub fn restart(&mut self, clear_snapshot: bool) { + self.canceled.store(true, Ordering::Relaxed); + self.items = Arc::new(boxcar::Vec::with_capacity(1024, self.items.columns())); + self.state = State::Cleared; + if clear_snapshot { + self.snapshot.clear(self.items.clone()); + } + } + + /// Update the internal configuration. + pub fn update_config(&mut self, config: Config) { + self.worker.lock().update_config(config) + } + + // Set whether the matcher should sort search results by score after + // matching. Defaults to true. + pub fn sort_results(&mut self, sort_results: bool) { + self.worker.lock().sort_results(sort_results) + } + + // Set whether the matcher should reverse the order of the input. + // Defaults to false. + pub fn reverse_items(&mut self, reverse_items: bool) { + self.worker.lock().reverse_items(reverse_items) + } + + /// The main way to interact with the matcher, this should be called + /// regularly (for example each time a frame is rendered). To avoid + /// excessive redraws this method will wait `timeout` milliseconds for the + /// worker thread to finish. It is recommend to set the timeout to 10ms. + pub fn tick(&mut self, timeout: u64) -> Status { + self.should_notify.store(false, atomic::Ordering::Relaxed); + let status = self.pattern.status(); + let canceled = status != pattern::Status::Unchanged || self.state.canceled(); + let mut res = self.tick_inner(timeout, canceled, status); + if !canceled { + return res; + } + self.state = State::Fresh; + let status2 = self.tick_inner(timeout, false, pattern::Status::Unchanged); + res.changed |= status2.changed; + res.running = status2.running; + res + } + + fn tick_inner(&mut self, timeout: u64, canceled: bool, status: pattern::Status) -> Status { + let mut inner = if canceled { + self.pattern.reset_status(); + self.canceled.store(true, atomic::Ordering::Relaxed); + self.worker.lock_arc() + } else { + let Some(worker) = self.worker.try_lock_arc_for(Duration::from_millis(timeout)) else { + self.should_notify.store(true, Ordering::Release); + return Status { + changed: false, + running: true, + }; + }; + worker + }; + + let changed = inner.running; + + let running = canceled || self.items.count() > inner.item_count(); + if inner.running { + inner.running = false; + if !inner.was_canceled && !self.state.canceled() { + self.snapshot.update(&inner) + } + } + if running { + inner.pattern.clone_from(&self.pattern); + self.canceled.store(false, atomic::Ordering::Relaxed); + if !canceled { + self.should_notify.store(true, atomic::Ordering::Release); + } + let cleared = self.state.cleared(); + if cleared { + inner.items = self.items.clone(); + } + self.pool + .spawn(move || unsafe { inner.run(status, cleared) }) + } + Status { changed, running } + } +} + +impl<T: Sync + Send> Drop for Nucleo<T> { + fn drop(&mut self) { + // we ensure the worker quits before dropping items to ensure that + // the worker can always assume the items outlive it + self.canceled.store(true, atomic::Ordering::Relaxed); + let lock = self.worker.try_lock_for(Duration::from_secs(1)); + if lock.is_none() { + unreachable!("thread pool failed to shutdown properly") + } + } +} diff --git a/crates/atuin-nucleo/src/par_sort.rs b/crates/atuin-nucleo/src/par_sort.rs new file mode 100644 index 00000000..92f716cc --- /dev/null +++ b/crates/atuin-nucleo/src/par_sort.rs @@ -0,0 +1,895 @@ +//! Parallel quicksort. +//! +//! This implementation is copied verbatim from `std::slice::sort_unstable` and then parallelized. +//! The only difference from the original is that calls to `recurse` are executed in parallel using +//! `rayon_core::join`. +//! Further modified for nucleo to allow canceling the sort + +// Copyright (c) 2010 The Rust Project Developers +// +// Permission is hereby granted, free of charge, to any +// person obtaining a copy of this software and associated +// documentation files (the "Software"), to deal in the +// Software without restriction, including without +// limitation the rights to use, copy, modify, merge, +// publish, distribute, sublicense, and/or sell copies of +// the Software, and to permit persons to whom the Software +// is furnished to do so, subject to the following +// conditions: +// +// The above copyright notice and this permission notice +// shall be included in all copies or substantial portions +// of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF +// ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED +// TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A +// PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT +// SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY +// CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION +// OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR +// IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER +// DEALINGS IN THE SOFTWARE. + +use std::cmp; +use std::mem::{self, MaybeUninit}; +use std::ptr; +use std::sync::atomic::{self, AtomicBool}; + +/// When dropped, copies from `src` into `dest`. +struct CopyOnDrop<T> { + src: *const T, + dest: *mut T, +} + +impl<T> Drop for CopyOnDrop<T> { + fn drop(&mut self) { + // SAFETY: This is a helper class. + // Please refer to its usage for correctness. + // Namely, one must be sure that `src` and `dst` does not overlap as required by `ptr::copy_nonoverlapping`. + unsafe { + ptr::copy_nonoverlapping(self.src, self.dest, 1); + } + } +} + +/// Shifts the first element to the right until it encounters a greater or equal element. +fn shift_head<T, F>(v: &mut [T], is_less: &F) +where + F: Fn(&T, &T) -> bool, +{ + let len = v.len(); + // SAFETY: The unsafe operations below involves indexing without a bounds check (by offsetting a + // pointer) and copying memory (`ptr::copy_nonoverlapping`). + // + // a. Indexing: + // 1. We checked the size of the array to >=2. + // 2. All the indexing that we will do is always between {0 <= index < len} at most. + // + // b. Memory copying + // 1. We are obtaining pointers to references which are guaranteed to be valid. + // 2. They cannot overlap because we obtain pointers to difference indices of the slice. + // Namely, `i` and `i-1`. + // 3. If the slice is properly aligned, the elements are properly aligned. + // It is the caller's responsibility to make sure the slice is properly aligned. + // + // See comments below for further detail. + unsafe { + // If the first two elements are out-of-order... + if len >= 2 && is_less(v.get_unchecked(1), v.get_unchecked(0)) { + // Read the first element into a stack-allocated variable. If a following comparison + // operation panics, `hole` will get dropped and automatically write the element back + // into the slice. + let tmp = mem::ManuallyDrop::new(ptr::read(v.get_unchecked(0))); + let v = v.as_mut_ptr(); + let mut hole = CopyOnDrop { + src: &*tmp, + dest: v.add(1), + }; + ptr::copy_nonoverlapping(v.add(1), v.add(0), 1); + + for i in 2..len { + if !is_less(&*v.add(i), &*tmp) { + break; + } + + // Move `i`-th element one place to the left, thus shifting the hole to the right. + ptr::copy_nonoverlapping(v.add(i), v.add(i - 1), 1); + hole.dest = v.add(i); + } + // `hole` gets dropped and thus copies `tmp` into the remaining hole in `v`. + } + } +} + +/// Shifts the last element to the left until it encounters a smaller or equal element. +fn shift_tail<T, F>(v: &mut [T], is_less: &F) +where + F: Fn(&T, &T) -> bool, +{ + let len = v.len(); + // SAFETY: The unsafe operations below involves indexing without a bound check (by offsetting a + // pointer) and copying memory (`ptr::copy_nonoverlapping`). + // + // a. Indexing: + // 1. We checked the size of the array to >= 2. + // 2. All the indexing that we will do is always between `0 <= index < len-1` at most. + // + // b. Memory copying + // 1. We are obtaining pointers to references which are guaranteed to be valid. + // 2. They cannot overlap because we obtain pointers to difference indices of the slice. + // Namely, `i` and `i+1`. + // 3. If the slice is properly aligned, the elements are properly aligned. + // It is the caller's responsibility to make sure the slice is properly aligned. + // + // See comments below for further detail. + unsafe { + // If the last two elements are out-of-order... + if len >= 2 && is_less(v.get_unchecked(len - 1), v.get_unchecked(len - 2)) { + // Read the last element into a stack-allocated variable. If a following comparison + // operation panics, `hole` will get dropped and automatically write the element back + // into the slice. + let tmp = mem::ManuallyDrop::new(ptr::read(v.get_unchecked(len - 1))); + let v = v.as_mut_ptr(); + let mut hole = CopyOnDrop { + src: &*tmp, + dest: v.add(len - 2), + }; + ptr::copy_nonoverlapping(v.add(len - 2), v.add(len - 1), 1); + + for i in (0..len - 2).rev() { + if !is_less(&*tmp, &*v.add(i)) { + break; + } + + // Move `i`-th element one place to the right, thus shifting the hole to the left. + ptr::copy_nonoverlapping(v.add(i), v.add(i + 1), 1); + hole.dest = v.add(i); + } + // `hole` gets dropped and thus copies `tmp` into the remaining hole in `v`. + } + } +} + +/// Partially sorts a slice by shifting several out-of-order elements around. +/// +/// Returns `true` if the slice is sorted at the end. This function is *O*(*n*) worst-case. +#[cold] +fn partial_insertion_sort<T, F>(v: &mut [T], is_less: &F) -> bool +where + F: Fn(&T, &T) -> bool, +{ + // Maximum number of adjacent out-of-order pairs that will get shifted. + const MAX_STEPS: usize = 5; + // If the slice is shorter than this, don't shift any elements. + const SHORTEST_SHIFTING: usize = 50; + + let len = v.len(); + let mut i = 1; + + for _ in 0..MAX_STEPS { + // SAFETY: We already explicitly did the bound checking with `i < len`. + // All our subsequent indexing is only in the range `0 <= index < len` + unsafe { + // Find the next pair of adjacent out-of-order elements. + while i < len && !is_less(v.get_unchecked(i), v.get_unchecked(i - 1)) { + i += 1; + } + } + + // Are we done? + if i == len { + return true; + } + + // Don't shift elements on short arrays, that has a performance cost. + if len < SHORTEST_SHIFTING { + return false; + } + + // Swap the found pair of elements. This puts them in correct order. + v.swap(i - 1, i); + + // Shift the smaller element to the left. + shift_tail(&mut v[..i], is_less); + // Shift the greater element to the right. + shift_head(&mut v[i..], is_less); + } + + // Didn't manage to sort the slice in the limited number of steps. + false +} + +/// Sorts a slice using insertion sort, which is *O*(*n*^2) worst-case. +fn insertion_sort<T, F>(v: &mut [T], is_less: &F) +where + F: Fn(&T, &T) -> bool, +{ + for i in 1..v.len() { + shift_tail(&mut v[..i + 1], is_less); + } +} + +/// Sorts `v` using heapsort, which guarantees *O*(*n* \* log(*n*)) worst-case. +#[cold] +fn heapsort<T, F>(v: &mut [T], is_less: &F) +where + F: Fn(&T, &T) -> bool, +{ + // This binary heap respects the invariant `parent >= child`. + let sift_down = |v: &mut [T], mut node| { + loop { + // Children of `node`. + let mut child = 2 * node + 1; + if child >= v.len() { + break; + } + + // Choose the greater child. + if child + 1 < v.len() && is_less(&v[child], &v[child + 1]) { + child += 1; + } + + // Stop if the invariant holds at `node`. + if !is_less(&v[node], &v[child]) { + break; + } + + // Swap `node` with the greater child, move one step down, and continue sifting. + v.swap(node, child); + node = child; + } + }; + + // Build the heap in linear time. + for i in (0..v.len() / 2).rev() { + sift_down(v, i); + } + + // Pop maximal elements from the heap. + for i in (1..v.len()).rev() { + v.swap(0, i); + sift_down(&mut v[..i], 0); + } +} + +/// Partitions `v` into elements smaller than `pivot`, followed by elements greater than or equal +/// to `pivot`. +/// +/// Returns the number of elements smaller than `pivot`. +/// +/// Partitioning is performed block-by-block in order to minimize the cost of branching operations. +/// This idea is presented in the [BlockQuicksort][pdf] paper. +/// +/// [pdf]: https://drops.dagstuhl.de/opus/volltexte/2016/6389/pdf/LIPIcs-ESA-2016-38.pdf +fn partition_in_blocks<T, F>(v: &mut [T], pivot: &T, is_less: &F) -> usize +where + F: Fn(&T, &T) -> bool, +{ + // Number of elements in a typical block. + const BLOCK: usize = 128; + + // The partitioning algorithm repeats the following steps until completion: + // + // 1. Trace a block from the left side to identify elements greater than or equal to the pivot. + // 2. Trace a block from the right side to identify elements smaller than the pivot. + // 3. Exchange the identified elements between the left and right side. + // + // We keep the following variables for a block of elements: + // + // 1. `block` - Number of elements in the block. + // 2. `start` - Start pointer into the `offsets` array. + // 3. `end` - End pointer into the `offsets` array. + // 4. `offsets - Indices of out-of-order elements within the block. + + // The current block on the left side (from `l` to `l.add(block_l)`). + let mut l = v.as_mut_ptr(); + let mut block_l = BLOCK; + let mut start_l = ptr::null_mut(); + let mut end_l = ptr::null_mut(); + let mut offsets_l = [MaybeUninit::<u8>::uninit(); BLOCK]; + + // The current block on the right side (from `r.sub(block_r)` to `r`). + // SAFETY: The documentation for .add() specifically mention that `vec.as_ptr().add(vec.len())` is always safe` + let mut r = unsafe { l.add(v.len()) }; + let mut block_r = BLOCK; + let mut start_r = ptr::null_mut(); + let mut end_r = ptr::null_mut(); + let mut offsets_r = [MaybeUninit::<u8>::uninit(); BLOCK]; + + // FIXME: When we get VLAs, try creating one array of length `min(v.len(), 2 * BLOCK)` rather + // than two fixed-size arrays of length `BLOCK`. VLAs might be more cache-efficient. + + // Returns the number of elements between pointers `l` (inclusive) and `r` (exclusive). + fn width<T>(l: *mut T, r: *mut T) -> usize { + assert!(mem::size_of::<T>() > 0); + // FIXME: this should *likely* use `offset_from`, but more + // investigation is needed (including running tests in miri). + // TODO unstable: (r.addr() - l.addr()) / mem::size_of::<T>() + (r as usize - l as usize) / mem::size_of::<T>() + } + + loop { + // We are done with partitioning block-by-block when `l` and `r` get very close. Then we do + // some patch-up work in order to partition the remaining elements in between. + let is_done = width(l, r) <= 2 * BLOCK; + + if is_done { + // Number of remaining elements (still not compared to the pivot). + let mut rem = width(l, r); + if start_l < end_l || start_r < end_r { + rem -= BLOCK; + } + + // Adjust block sizes so that the left and right block don't overlap, but get perfectly + // aligned to cover the whole remaining gap. + if start_l < end_l { + block_r = rem; + } else if start_r < end_r { + block_l = rem; + } else { + // There were the same number of elements to switch on both blocks during the last + // iteration, so there are no remaining elements on either block. Cover the remaining + // items with roughly equally-sized blocks. + block_l = rem / 2; + block_r = rem - block_l; + } + debug_assert!(block_l <= BLOCK && block_r <= BLOCK); + debug_assert!(width(l, r) == block_l + block_r); + } + + if start_l == end_l { + // Trace `block_l` elements from the left side. + // TODO unstable: start_l = MaybeUninit::slice_as_mut_ptr(&mut offsets_l); + start_l = offsets_l.as_mut_ptr() as *mut u8; + end_l = start_l; + let mut elem = l; + + for i in 0..block_l { + // SAFETY: The unsafety operations below involve the usage of the `offset`. + // According to the conditions required by the function, we satisfy them because: + // 1. `offsets_l` is stack-allocated, and thus considered separate allocated object. + // 2. The function `is_less` returns a `bool`. + // Casting a `bool` will never overflow `isize`. + // 3. We have guaranteed that `block_l` will be `<= BLOCK`. + // Plus, `end_l` was initially set to the begin pointer of `offsets_` which was declared on the stack. + // Thus, we know that even in the worst case (all invocations of `is_less` returns false) we will only be at most 1 byte pass the end. + // Another unsafety operation here is dereferencing `elem`. + // However, `elem` was initially the begin pointer to the slice which is always valid. + unsafe { + // Branchless comparison. + *end_l = i as u8; + end_l = end_l.offset(!is_less(&*elem, pivot) as isize); + elem = elem.offset(1); + } + } + } + + if start_r == end_r { + // Trace `block_r` elements from the right side. + // TODO unstable: start_r = MaybeUninit::slice_as_mut_ptr(&mut offsets_r); + start_r = offsets_r.as_mut_ptr() as *mut u8; + end_r = start_r; + let mut elem = r; + + for i in 0..block_r { + // SAFETY: The unsafety operations below involve the usage of the `offset`. + // According to the conditions required by the function, we satisfy them because: + // 1. `offsets_r` is stack-allocated, and thus considered separate allocated object. + // 2. The function `is_less` returns a `bool`. + // Casting a `bool` will never overflow `isize`. + // 3. We have guaranteed that `block_r` will be `<= BLOCK`. + // Plus, `end_r` was initially set to the begin pointer of `offsets_` which was declared on the stack. + // Thus, we know that even in the worst case (all invocations of `is_less` returns true) we will only be at most 1 byte pass the end. + // Another unsafety operation here is dereferencing `elem`. + // However, `elem` was initially `1 * sizeof(T)` past the end and we decrement it by `1 * sizeof(T)` before accessing it. + // Plus, `block_r` was asserted to be less than `BLOCK` and `elem` will therefore at most be pointing to the beginning of the slice. + unsafe { + // Branchless comparison. + elem = elem.offset(-1); + *end_r = i as u8; + end_r = end_r.offset(is_less(&*elem, pivot) as isize); + } + } + } + + // Number of out-of-order elements to swap between the left and right side. + let count = cmp::min(width(start_l, end_l), width(start_r, end_r)); + + if count > 0 { + macro_rules! left { + () => { + l.offset(*start_l as isize) + }; + } + macro_rules! right { + () => { + r.offset(-(*start_r as isize) - 1) + }; + } + + // Instead of swapping one pair at the time, it is more efficient to perform a cyclic + // permutation. This is not strictly equivalent to swapping, but produces a similar + // result using fewer memory operations. + + // SAFETY: The use of `ptr::read` is valid because there is at least one element in + // both `offsets_l` and `offsets_r`, so `left!` is a valid pointer to read from. + // + // The uses of `left!` involve calls to `offset` on `l`, which points to the + // beginning of `v`. All the offsets pointed-to by `start_l` are at most `block_l`, so + // these `offset` calls are safe as all reads are within the block. The same argument + // applies for the uses of `right!`. + // + // The calls to `start_l.offset` are valid because there are at most `count-1` of them, + // plus the final one at the end of the unsafe block, where `count` is the minimum number + // of collected offsets in `offsets_l` and `offsets_r`, so there is no risk of there not + // being enough elements. The same reasoning applies to the calls to `start_r.offset`. + // + // The calls to `copy_nonoverlapping` are safe because `left!` and `right!` are guaranteed + // not to overlap, and are valid because of the reasoning above. + unsafe { + let tmp = ptr::read(left!()); + ptr::copy_nonoverlapping(right!(), left!(), 1); + + for _ in 1..count { + start_l = start_l.offset(1); + ptr::copy_nonoverlapping(left!(), right!(), 1); + start_r = start_r.offset(1); + ptr::copy_nonoverlapping(right!(), left!(), 1); + } + + ptr::copy_nonoverlapping(&tmp, right!(), 1); + mem::forget(tmp); + start_l = start_l.offset(1); + start_r = start_r.offset(1); + } + } + + if start_l == end_l { + // All out-of-order elements in the left block were moved. Move to the next block. + + // block-width-guarantee + // SAFETY: if `!is_done` then the slice width is guaranteed to be at least `2*BLOCK` wide. There + // are at most `BLOCK` elements in `offsets_l` because of its size, so the `offset` operation is + // safe. Otherwise, the debug assertions in the `is_done` case guarantee that + // `width(l, r) == block_l + block_r`, namely, that the block sizes have been adjusted to account + // for the smaller number of remaining elements. + l = unsafe { l.add(block_l) }; + } + + if start_r == end_r { + // All out-of-order elements in the right block were moved. Move to the previous block. + + // SAFETY: Same argument as [block-width-guarantee]. Either this is a full block `2*BLOCK`-wide, + // or `block_r` has been adjusted for the last handful of elements. + r = unsafe { r.offset(-(block_r as isize)) }; + } + + if is_done { + break; + } + } + + // All that remains now is at most one block (either the left or the right) with out-of-order + // elements that need to be moved. Such remaining elements can be simply shifted to the end + // within their block. + + if start_l < end_l { + // The left block remains. + // Move its remaining out-of-order elements to the far right. + debug_assert_eq!(width(l, r), block_l); + while start_l < end_l { + // remaining-elements-safety + // SAFETY: while the loop condition holds there are still elements in `offsets_l`, so it + // is safe to point `end_l` to the previous element. + // + // The `ptr::swap` is safe if both its arguments are valid for reads and writes: + // - Per the debug assert above, the distance between `l` and `r` is `block_l` + // elements, so there can be at most `block_l` remaining offsets between `start_l` + // and `end_l`. This means `r` will be moved at most `block_l` steps back, which + // makes the `r.offset` calls valid (at that point `l == r`). + // - `offsets_l` contains valid offsets into `v` collected during the partitioning of + // the last block, so the `l.offset` calls are valid. + unsafe { + end_l = end_l.offset(-1); + ptr::swap(l.offset(*end_l as isize), r.offset(-1)); + r = r.offset(-1); + } + } + width(v.as_mut_ptr(), r) + } else if start_r < end_r { + // The right block remains. + // Move its remaining out-of-order elements to the far left. + debug_assert_eq!(width(l, r), block_r); + while start_r < end_r { + // SAFETY: See the reasoning in [remaining-elements-safety]. + unsafe { + end_r = end_r.offset(-1); + ptr::swap(l, r.offset(-(*end_r as isize) - 1)); + l = l.offset(1); + } + } + width(v.as_mut_ptr(), l) + } else { + // Nothing else to do, we're done. + width(v.as_mut_ptr(), l) + } +} + +/// Partitions `v` into elements smaller than `v[pivot]`, followed by elements greater than or +/// equal to `v[pivot]`. +/// +/// Returns a tuple of: +/// +/// 1. Number of elements smaller than `v[pivot]`. +/// 2. True if `v` was already partitioned. +fn partition<T, F>(v: &mut [T], pivot: usize, is_less: &F) -> (usize, bool) +where + F: Fn(&T, &T) -> bool, +{ + let (mid, was_partitioned) = { + // Place the pivot at the beginning of slice. + v.swap(0, pivot); + let (pivot, v) = v.split_at_mut(1); + let pivot = &mut pivot[0]; + + // Read the pivot into a stack-allocated variable for efficiency. If a following comparison + // operation panics, the pivot will be automatically written back into the slice. + + // SAFETY: `pivot` is a reference to the first element of `v`, so `ptr::read` is safe. + let tmp = mem::ManuallyDrop::new(unsafe { ptr::read(pivot) }); + let _pivot_guard = CopyOnDrop { + src: &*tmp, + dest: pivot, + }; + let pivot = &*tmp; + + // Find the first pair of out-of-order elements. + let mut l = 0; + let mut r = v.len(); + + // SAFETY: The unsafety below involves indexing an array. + // For the first one: We already do the bounds checking here with `l < r`. + // For the second one: We initially have `l == 0` and `r == v.len()` and we checked that `l < r` at every indexing operation. + // From here we know that `r` must be at least `r == l` which was shown to be valid from the first one. + unsafe { + // Find the first element greater than or equal to the pivot. + while l < r && is_less(v.get_unchecked(l), pivot) { + l += 1; + } + + // Find the last element smaller that the pivot. + while l < r && !is_less(v.get_unchecked(r - 1), pivot) { + r -= 1; + } + } + + ( + l + partition_in_blocks(&mut v[l..r], pivot, is_less), + l >= r, + ) + + // `_pivot_guard` goes out of scope and writes the pivot (which is a stack-allocated + // variable) back into the slice where it originally was. This step is critical in ensuring + // safety! + }; + + // Place the pivot between the two partitions. + v.swap(0, mid); + + (mid, was_partitioned) +} + +/// Partitions `v` into elements equal to `v[pivot]` followed by elements greater than `v[pivot]`. +/// +/// Returns the number of elements equal to the pivot. It is assumed that `v` does not contain +/// elements smaller than the pivot. +fn partition_equal<T, F>(v: &mut [T], pivot: usize, is_less: &F) -> usize +where + F: Fn(&T, &T) -> bool, +{ + // Place the pivot at the beginning of slice. + v.swap(0, pivot); + let (pivot, v) = v.split_at_mut(1); + let pivot = &mut pivot[0]; + + // Read the pivot into a stack-allocated variable for efficiency. If a following comparison + // operation panics, the pivot will be automatically written back into the slice. + // SAFETY: The pointer here is valid because it is obtained from a reference to a slice. + let tmp = mem::ManuallyDrop::new(unsafe { ptr::read(pivot) }); + let _pivot_guard = CopyOnDrop { + src: &*tmp, + dest: pivot, + }; + let pivot = &*tmp; + + // Now partition the slice. + let mut l = 0; + let mut r = v.len(); + loop { + // SAFETY: The unsafety below involves indexing an array. + // For the first one: We already do the bounds checking here with `l < r`. + // For the second one: We initially have `l == 0` and `r == v.len()` and we checked that `l < r` at every indexing operation. + // From here we know that `r` must be at least `r == l` which was shown to be valid from the first one. + unsafe { + // Find the first element greater than the pivot. + while l < r && !is_less(pivot, v.get_unchecked(l)) { + l += 1; + } + + // Find the last element equal to the pivot. + while l < r && is_less(pivot, v.get_unchecked(r - 1)) { + r -= 1; + } + + // Are we done? + if l >= r { + break; + } + + // Swap the found pair of out-of-order elements. + r -= 1; + let ptr = v.as_mut_ptr(); + ptr::swap(ptr.add(l), ptr.add(r)); + l += 1; + } + } + + // We found `l` elements equal to the pivot. Add 1 to account for the pivot itself. + l + 1 + + // `_pivot_guard` goes out of scope and writes the pivot (which is a stack-allocated variable) + // back into the slice where it originally was. This step is critical in ensuring safety! +} + +/// Scatters some elements around in an attempt to break patterns that might cause imbalanced +/// partitions in quicksort. +#[cold] +fn break_patterns<T>(v: &mut [T]) { + let len = v.len(); + if len >= 8 { + // Pseudorandom number generator from the "Xorshift RNGs" paper by George Marsaglia. + let mut random = len as u32; + let mut gen_u32 = || { + random ^= random << 13; + random ^= random >> 17; + random ^= random << 5; + random + }; + let mut gen_usize = || { + if usize::BITS <= 32 { + gen_u32() as usize + } else { + (((gen_u32() as u64) << 32) | (gen_u32() as u64)) as usize + } + }; + + // Take random numbers modulo this number. + // The number fits into `usize` because `len` is not greater than `isize::MAX`. + let modulus = len.next_power_of_two(); + + // Some pivot candidates will be in the nearby of this index. Let's randomize them. + let pos = len / 4 * 2; + + for i in 0..3 { + // Generate a random number modulo `len`. However, in order to avoid costly operations + // we first take it modulo a power of two, and then decrease by `len` until it fits + // into the range `[0, len - 1]`. + let mut other = gen_usize() & (modulus - 1); + + // `other` is guaranteed to be less than `2 * len`. + if other >= len { + other -= len; + } + + v.swap(pos - 1 + i, other); + } + } +} + +/// Chooses a pivot in `v` and returns the index and `true` if the slice is likely already sorted. +/// +/// Elements in `v` might be reordered in the process. +fn choose_pivot<T, F>(v: &mut [T], is_less: &F) -> (usize, bool) +where + F: Fn(&T, &T) -> bool, +{ + // Minimum length to choose the median-of-medians method. + // Shorter slices use the simple median-of-three method. + const SHORTEST_MEDIAN_OF_MEDIANS: usize = 50; + // Maximum number of swaps that can be performed in this function. + const MAX_SWAPS: usize = 4 * 3; + + let len = v.len(); + + // Three indices near which we are going to choose a pivot. + #[allow(clippy::identity_op)] + let mut a = len / 4 * 1; + let mut b = len / 4 * 2; + let mut c = len / 4 * 3; + + // Counts the total number of swaps we are about to perform while sorting indices. + let mut swaps = 0; + + if len >= 8 { + // Swaps indices so that `v[a] <= v[b]`. + // SAFETY: `len >= 8` so there are at least two elements in the neighborhoods of + // `a`, `b` and `c`. This means the three calls to `sort_adjacent` result in + // corresponding calls to `sort3` with valid 3-item neighborhoods around each + // pointer, which in turn means the calls to `sort2` are done with valid + // references. Thus the `v.get_unchecked` calls are safe, as is the `ptr::swap` + // call. + let mut sort2 = |a: &mut usize, b: &mut usize| unsafe { + if is_less(v.get_unchecked(*b), v.get_unchecked(*a)) { + ptr::swap(a, b); + swaps += 1; + } + }; + + // Swaps indices so that `v[a] <= v[b] <= v[c]`. + let mut sort3 = |a: &mut usize, b: &mut usize, c: &mut usize| { + sort2(a, b); + sort2(b, c); + sort2(a, b); + }; + + if len >= SHORTEST_MEDIAN_OF_MEDIANS { + // Finds the median of `v[a - 1], v[a], v[a + 1]` and stores the index into `a`. + let mut sort_adjacent = |a: &mut usize| { + let tmp = *a; + sort3(&mut (tmp - 1), a, &mut (tmp + 1)); + }; + + // Find medians in the neighborhoods of `a`, `b`, and `c`. + sort_adjacent(&mut a); + sort_adjacent(&mut b); + sort_adjacent(&mut c); + } + + // Find the median among `a`, `b`, and `c`. + sort3(&mut a, &mut b, &mut c); + } + + if swaps < MAX_SWAPS { + (b, swaps == 0) + } else { + // The maximum number of swaps was performed. Chances are the slice is descending or mostly + // descending, so reversing will probably help sort it faster. + v.reverse(); + (len - 1 - b, true) + } +} + +/// Sorts `v` recursively. +/// +/// If the slice had a predecessor in the original array, it is specified as `pred`. +/// +/// `limit` is the number of allowed imbalanced partitions before switching to `heapsort`. If zero, +/// this function will immediately switch to heapsort. +fn recurse<'a, T, F>( + mut v: &'a mut [T], + is_less: &F, + mut pred: Option<&'a mut T>, + mut limit: u32, + canceled: &AtomicBool, +) -> bool +where + T: Send, + F: Fn(&T, &T) -> bool + Sync, +{ + // Slices of up to this length get sorted using insertion sort. + const MAX_INSERTION: usize = 20; + // If both partitions are up to this length, we continue sequentially. This number is as small + // as possible but so that the overhead of Rayon's task scheduling is still negligible. + const MAX_SEQUENTIAL: usize = 2000; + + // True if the last partitioning was reasonably balanced. + let mut was_balanced = true; + // True if the last partitioning didn't shuffle elements (the slice was already partitioned). + let mut was_partitioned = true; + + loop { + let len = v.len(); + + // Very short slices get sorted using insertion sort. + if len <= MAX_INSERTION { + insertion_sort(v, is_less); + return false; + } + + // If too many bad pivot choices were made, simply fall back to heapsort in order to + // guarantee `O(n * log(n))` worst-case. + if limit == 0 { + heapsort(v, is_less); + return false; + } + + // If the last partitioning was imbalanced, try breaking patterns in the slice by shuffling + // some elements around. Hopefully we'll choose a better pivot this time. + if !was_balanced { + break_patterns(v); + limit -= 1; + } + + // Choose a pivot and try guessing whether the slice is already sorted. + let (pivot, likely_sorted) = choose_pivot(v, is_less); + + // If the last partitioning was decently balanced and didn't shuffle elements, and if pivot + // selection predicts the slice is likely already sorted... + if was_balanced && was_partitioned && likely_sorted { + // Try identifying several out-of-order elements and shifting them to correct + // positions. If the slice ends up being completely sorted, we're done. + if partial_insertion_sort(v, is_less) { + return false; + } + } + + // If the chosen pivot is equal to the predecessor, then it's the smallest element in the + // slice. Partition the slice into elements equal to and elements greater than the pivot. + // This case is usually hit when the slice contains many duplicate elements. + if let Some(ref p) = pred { + if !is_less(p, &v[pivot]) { + let mid = partition_equal(v, pivot, is_less); + + // Continue sorting elements greater than the pivot. + v = &mut v[mid..]; + continue; + } + } + + // Partition the slice. + let (mid, was_p) = partition(v, pivot, is_less); + was_balanced = cmp::min(mid, len - mid) >= len / 8; + was_partitioned = was_p; + + // Split the slice into `left`, `pivot`, and `right`. + let (left, right) = v.split_at_mut(mid); + let (pivot, right) = right.split_at_mut(1); + let pivot = &mut pivot[0]; + + if cmp::max(left.len(), right.len()) <= MAX_SEQUENTIAL { + // Recurse into the shorter side only in order to minimize the total number of recursive + // calls and consume less stack space. Then just continue with the longer side (this is + // akin to tail recursion). + if left.len() < right.len() { + recurse(left, is_less, pred, limit, canceled); + v = right; + pred = Some(pivot); + } else { + recurse(right, is_less, Some(pivot), limit, canceled); + v = left; + } + } else if canceled.load(atomic::Ordering::Relaxed) { + break true; + } else { + // Sort the left and right half in parallel. + let (canceled1, canceled2) = rayon::join( + || recurse(left, is_less, pred, limit, canceled), + || recurse(right, is_less, Some(pivot), limit, canceled), + ); + break canceled1 | canceled2; + } + } +} + +/// Sorts `v` using pattern-defeating quicksort in parallel. +/// +/// The algorithm is unstable, in-place, and *O*(*n* \* log(*n*)) worst-case. +pub(crate) fn par_quicksort<T, F>(v: &mut [T], is_less: F, canceled: &AtomicBool) -> bool +where + T: Send, + F: Fn(&T, &T) -> bool + Sync, +{ + // Sorting has no meaningful behavior on zero-sized types. + if mem::size_of::<T>() == 0 { + return false; + } + if canceled.load(atomic::Ordering::Relaxed) { + return true; + } + + // Limit the number of imbalanced partitions to `floor(log2(len)) + 1`. + let limit = usize::BITS - v.len().leading_zeros(); + + recurse(v, &is_less, None, limit, canceled) +} diff --git a/crates/atuin-nucleo/src/pattern.rs b/crates/atuin-nucleo/src/pattern.rs new file mode 100644 index 00000000..816b0a31 --- /dev/null +++ b/crates/atuin-nucleo/src/pattern.rs @@ -0,0 +1,100 @@ +pub use nucleo_matcher::pattern::{Atom, AtomKind, CaseMatching, Normalization, Pattern}; +use nucleo_matcher::{Matcher, Utf32String}; + +#[cfg(test)] +mod tests; + +#[derive(Debug, PartialEq, Eq, Clone, Copy, PartialOrd, Ord, Default)] +pub(crate) enum Status { + #[default] + Unchanged, + Update, + Rescore, +} + +#[derive(Debug)] +pub struct MultiPattern { + cols: Vec<(Pattern, Status)>, +} + +impl Clone for MultiPattern { + fn clone(&self) -> Self { + Self { + cols: self.cols.clone(), + } + } + + fn clone_from(&mut self, source: &Self) { + self.cols.clone_from(&source.cols) + } +} + +impl MultiPattern { + /// Creates a multi pattern with `columns` empty column patterns. + pub fn new(columns: usize) -> Self { + Self { + cols: vec![Default::default(); columns], + } + } + + /// Reparses a column. By specifying `append` the caller promises that text passed + /// to the previous `reparse` invocation is a prefix of `new_text`. This enables + /// additional optimizations but can lead to missing matches if an incorrect value + /// is passed. + pub fn reparse( + &mut self, + column: usize, + new_text: &str, + case_matching: CaseMatching, + normalization: Normalization, + append: bool, + ) { + let old_status = self.cols[column].1; + if append + && old_status != Status::Rescore + && self.cols[column] + .0 + .atoms + .last() + .map_or(true, |last| !last.negative) + { + self.cols[column].1 = Status::Update; + } else { + self.cols[column].1 = Status::Rescore; + } + self.cols[column] + .0 + .reparse(new_text, case_matching, normalization); + } + + pub fn column_pattern(&self, column: usize) -> &Pattern { + &self.cols[column].0 + } + + pub(crate) fn status(&self) -> Status { + self.cols + .iter() + .map(|&(_, status)| status) + .max() + .unwrap_or(Status::Unchanged) + } + + pub(crate) fn reset_status(&mut self) { + for (_, status) in &mut self.cols { + *status = Status::Unchanged + } + } + + pub fn score(&self, haystack: &[Utf32String], matcher: &mut Matcher) -> Option<u32> { + // TODO: wheight columns? + let mut score = 0; + for ((pattern, _), haystack) in self.cols.iter().zip(haystack) { + score += pattern.score(haystack.slice(..), matcher)? + } + Some(score) + } + + pub fn is_empty(&self) -> bool { + self.cols.iter().all(|(pat, _)| pat.atoms.is_empty()) + } +} diff --git a/crates/atuin-nucleo/src/pattern/tests.rs b/crates/atuin-nucleo/src/pattern/tests.rs new file mode 100644 index 00000000..40e8e328 --- /dev/null +++ b/crates/atuin-nucleo/src/pattern/tests.rs @@ -0,0 +1,14 @@ +use nucleo_matcher::pattern::{CaseMatching, Normalization}; + +use crate::pattern::{MultiPattern, Status}; + +#[test] +fn append() { + let mut pat = MultiPattern::new(1); + pat.reparse(0, "!", CaseMatching::Smart, Normalization::Smart, true); + assert_eq!(pat.status(), Status::Update); + pat.reparse(0, "!f", CaseMatching::Smart, Normalization::Smart, true); + assert_eq!(pat.status(), Status::Update); + pat.reparse(0, "!fo", CaseMatching::Smart, Normalization::Smart, true); + assert_eq!(pat.status(), Status::Rescore); +} diff --git a/crates/atuin-nucleo/src/tests.rs b/crates/atuin-nucleo/src/tests.rs new file mode 100644 index 00000000..676c50df --- /dev/null +++ b/crates/atuin-nucleo/src/tests.rs @@ -0,0 +1,27 @@ +use std::sync::Arc; + +use nucleo_matcher::Config; + +use crate::Nucleo; + +#[test] +fn active_injector_count() { + let mut nucleo: Nucleo<()> = Nucleo::new(Config::DEFAULT, Arc::new(|| ()), Some(1), 1); + assert_eq!(nucleo.active_injectors(), 0); + let injector = nucleo.injector(); + assert_eq!(nucleo.active_injectors(), 1); + let injector2 = nucleo.injector(); + assert_eq!(nucleo.active_injectors(), 2); + drop(injector2); + assert_eq!(nucleo.active_injectors(), 1); + nucleo.restart(false); + assert_eq!(nucleo.active_injectors(), 0); + let injector3 = nucleo.injector(); + assert_eq!(nucleo.active_injectors(), 1); + nucleo.tick(0); + assert_eq!(nucleo.active_injectors(), 1); + drop(injector); + assert_eq!(nucleo.active_injectors(), 1); + drop(injector3); + assert_eq!(nucleo.active_injectors(), 0); +} diff --git a/crates/atuin-nucleo/src/worker.rs b/crates/atuin-nucleo/src/worker.rs new file mode 100644 index 00000000..f4077e6e --- /dev/null +++ b/crates/atuin-nucleo/src/worker.rs @@ -0,0 +1,301 @@ +use std::cell::UnsafeCell; +use std::mem::take; +use std::sync::atomic::{self, AtomicBool, AtomicU32}; +use std::sync::Arc; + +use nucleo_matcher::Config; +use parking_lot::Mutex; +use rayon::{prelude::*, ThreadPool}; + +use crate::par_sort::par_quicksort; +use crate::pattern::{self, MultiPattern}; +use crate::{boxcar, Match}; + +struct Matchers(Box<[UnsafeCell<nucleo_matcher::Matcher>]>); + +impl Matchers { + // this is not a true mut from ref, we use a cell here + #[allow(clippy::mut_from_ref)] + unsafe fn get(&self) -> &mut nucleo_matcher::Matcher { + &mut *self.0[rayon::current_thread_index().unwrap()].get() + } +} + +unsafe impl Sync for Matchers {} +unsafe impl Send for Matchers {} + +pub(crate) struct Worker<T: Sync + Send + 'static> { + pub(crate) running: bool, + matchers: Matchers, + pub(crate) matches: Vec<Match>, + pub(crate) pattern: MultiPattern, + pub(crate) sort_results: bool, + pub(crate) reverse_items: bool, + pub(crate) canceled: Arc<AtomicBool>, + pub(crate) should_notify: Arc<AtomicBool>, + pub(crate) was_canceled: bool, + pub(crate) last_snapshot: u32, + notify: Arc<(dyn Fn() + Sync + Send)>, + pub(crate) items: Arc<boxcar::Vec<T>>, + in_flight: Vec<u32>, +} + +impl<T: Sync + Send + 'static> Worker<T> { + pub(crate) fn item_count(&self) -> u32 { + self.last_snapshot - self.in_flight.len() as u32 + } + pub(crate) fn update_config(&mut self, config: Config) { + for matcher in self.matchers.0.iter_mut() { + matcher.get_mut().config = config.clone(); + } + } + pub(crate) fn sort_results(&mut self, sort_results: bool) { + self.sort_results = sort_results; + } + pub(crate) fn reverse_items(&mut self, reverse_items: bool) { + self.reverse_items = reverse_items; + } + + pub(crate) fn new( + worker_threads: Option<usize>, + config: Config, + notify: Arc<(dyn Fn() + Sync + Send)>, + cols: u32, + ) -> (ThreadPool, Self) { + let worker_threads = worker_threads + .unwrap_or_else(|| std::thread::available_parallelism().map_or(4, |it| it.get())); + let pool = rayon::ThreadPoolBuilder::new() + .thread_name(|i| format!("nucleo worker {i}")) + .num_threads(worker_threads) + .build() + .expect("creating threadpool failed"); + let matchers = (0..worker_threads) + .map(|_| UnsafeCell::new(nucleo_matcher::Matcher::new(config.clone()))) + .collect(); + let worker = Worker { + running: false, + matchers: Matchers(matchers), + last_snapshot: 0, + matches: Vec::new(), + // just a placeholder + pattern: MultiPattern::new(cols as usize), + sort_results: true, + reverse_items: false, + canceled: Arc::new(AtomicBool::new(false)), + should_notify: Arc::new(AtomicBool::new(false)), + was_canceled: false, + notify, + items: Arc::new(boxcar::Vec::with_capacity(2 * 1024, cols)), + in_flight: Vec::with_capacity(64), + }; + (pool, worker) + } + + unsafe fn process_new_items(&mut self, unmatched: &AtomicU32) { + let matchers = &self.matchers; + let pattern = &self.pattern; + self.matches.reserve(self.in_flight.len()); + self.in_flight.retain(|&idx| { + let Some(item) = self.items.get(idx) else { + return true; + }; + if let Some(score) = pattern.score(item.matcher_columns, matchers.get()) { + self.matches.push(Match { score, idx }); + }; + false + }); + let new_snapshot = self.items.par_snapshot(self.last_snapshot); + if new_snapshot.end() != self.last_snapshot { + let end = new_snapshot.end(); + let in_flight = Mutex::new(&mut self.in_flight); + let items = new_snapshot.map(|(idx, item)| { + let Some(item) = item else { + in_flight.lock().push(idx); + unmatched.fetch_add(1, atomic::Ordering::Relaxed); + return Match { + score: 0, + idx: u32::MAX, + }; + }; + if self.canceled.load(atomic::Ordering::Relaxed) { + return Match { score: 0, idx }; + } + let Some(score) = pattern.score(item.matcher_columns, matchers.get()) else { + unmatched.fetch_add(1, atomic::Ordering::Relaxed); + return Match { + score: 0, + idx: u32::MAX, + }; + }; + Match { score, idx } + }); + self.matches.par_extend(items); + self.last_snapshot = end; + } + } + + fn remove_in_flight_matches(&mut self) { + let mut off = 0; + self.in_flight.retain(|&i| { + let is_in_flight = self.items.get(i).is_none(); + if is_in_flight { + self.matches.remove((i - off) as usize); + off += 1; + } + is_in_flight + }); + } + + unsafe fn process_new_items_trivial(&mut self) { + let new_snapshot = self.items.snapshot(self.last_snapshot); + if new_snapshot.end() != self.last_snapshot { + let end = new_snapshot.end(); + let items = new_snapshot.filter_map(|(idx, item)| { + if item.is_none() { + self.in_flight.push(idx); + return None; + }; + Some(Match { score: 0, idx }) + }); + self.matches.extend(items); + self.last_snapshot = end; + } + } + + pub(crate) unsafe fn run(&mut self, pattern_status: pattern::Status, cleared: bool) { + self.running = true; + self.was_canceled = false; + + if cleared { + self.last_snapshot = 0; + self.in_flight.clear(); + self.matches.clear(); + } + + // TODO: be smarter around reusing past results for rescoring + if self.pattern.is_empty() { + self.reset_matches(); + self.process_new_items_trivial(); + let canceled = self.sort_matches(); + if canceled { + self.was_canceled = true; + } else if self.should_notify.load(atomic::Ordering::Relaxed) { + (self.notify)(); + } + return; + } + + if pattern_status == pattern::Status::Rescore { + self.reset_matches(); + } + + let mut unmatched = AtomicU32::new(0); + if pattern_status != pattern::Status::Unchanged && !self.matches.is_empty() { + self.process_new_items_trivial(); + let matchers = &self.matchers; + let pattern = &self.pattern; + self.matches + .par_iter_mut() + .take_any_while(|_| !self.canceled.load(atomic::Ordering::Relaxed)) + .for_each(|match_| { + if match_.idx == u32::MAX { + debug_assert_eq!(match_.score, 0); + unmatched.fetch_add(1, atomic::Ordering::Relaxed); + return; + } + // safety: in-flight items are never added to the matches + let item = self.items.get_unchecked(match_.idx); + if let Some(score) = pattern.score(item.matcher_columns, matchers.get()) { + match_.score = score; + } else { + unmatched.fetch_add(1, atomic::Ordering::Relaxed); + match_.score = 0; + match_.idx = u32::MAX; + } + }); + } else { + self.process_new_items(&unmatched); + } + + let canceled = self.sort_matches(); + if canceled { + self.was_canceled = true; + } else { + self.matches + .truncate(self.matches.len() - take(unmatched.get_mut()) as usize); + if self.should_notify.load(atomic::Ordering::Relaxed) { + (self.notify)(); + } + } + } + + unsafe fn sort_matches(&mut self) -> bool { + if self.sort_results { + par_quicksort( + &mut self.matches, + |match1, match2| { + if match1.score != match2.score { + return match1.score > match2.score; + } + if match1.idx == u32::MAX { + return false; + } + if match2.idx == u32::MAX { + return true; + } + // the tie breaker is comparatively rarely needed so we keep it + // in a branch especially because we need to access the items + // array here which involves some pointer chasing + let item1 = self.items.get_unchecked(match1.idx); + let item2 = &self.items.get_unchecked(match2.idx); + let len1: u32 = item1 + .matcher_columns + .iter() + .map(|haystack| haystack.len() as u32) + .sum(); + let len2 = item2 + .matcher_columns + .iter() + .map(|haystack| haystack.len() as u32) + .sum(); + if len1 == len2 { + if self.reverse_items { + match2.idx < match1.idx + } else { + match1.idx < match2.idx + } + } else { + len1 < len2 + } + }, + &self.canceled, + ) + } else { + par_quicksort( + &mut self.matches, + |match1, match2| { + if match1.idx == u32::MAX { + return false; + } + if match2.idx == u32::MAX { + return true; + } + if self.reverse_items { + match2.idx < match1.idx + } else { + match1.idx < match2.idx + } + }, + &self.canceled, + ) + } + } + + fn reset_matches(&mut self) { + self.matches.clear(); + self.matches + .extend((0..self.last_snapshot).map(|idx| Match { score: 0, idx })); + // there are usually only very few in flight items (one for each writer) + self.remove_in_flight_matches(); + } +} diff --git a/crates/atuin-nucleo/tarpaulin.toml b/crates/atuin-nucleo/tarpaulin.toml new file mode 100644 index 00000000..4869f256 --- /dev/null +++ b/crates/atuin-nucleo/tarpaulin.toml @@ -0,0 +1 @@ +exclude = ["matcher/src/tests.rs", "matcher/src/debug.rs", "matcher/src/chars/normalize.rs"] diff --git a/crates/atuin-nucleo/typos.toml b/crates/atuin-nucleo/typos.toml new file mode 100644 index 00000000..14fc504d --- /dev/null +++ b/crates/atuin-nucleo/typos.toml @@ -0,0 +1,3 @@ +default.extend-ignore-re = ["\\\\u\\{[0-9A-Za-z]*\\}"] +[files] +extend-exclude = ["matcher/src/tests.rs","src/pattern/tests.rs", "*.html"] diff --git a/crates/atuin-scripts/Cargo.toml b/crates/atuin-scripts/Cargo.toml new file mode 100644 index 00000000..f8294b79 --- /dev/null +++ b/crates/atuin-scripts/Cargo.toml @@ -0,0 +1,34 @@ +[package] +name = "atuin-scripts" +edition = "2024" +version = { workspace = true } +description = "The scripts crate for Atuin" + +authors.workspace = true +rust-version.workspace = true +license.workspace = true +homepage.workspace = true +repository.workspace = true +readme.workspace = true + +# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html + +[dependencies] +atuin-client = { path = "../atuin-client", version = "18.13.2" } +atuin-common = { path = "../atuin-common", version = "18.13.2" } + +tracing = { workspace = true } +tracing-subscriber = { workspace = true } +rmp = { version = "0.8.14" } +uuid = { workspace = true } +eyre = { workspace = true } +tokio = { workspace = true } +serde = { workspace = true } +typed-builder = { workspace = true } +pretty_assertions = { workspace = true } +sql-builder = { workspace = true } +sqlx = { workspace = true } +tempfile = { workspace = true } +minijinja = { workspace = true } +serde_json = { workspace = true } + diff --git a/crates/atuin-scripts/migrations/20250326160051_create_scripts.down.sql b/crates/atuin-scripts/migrations/20250326160051_create_scripts.down.sql new file mode 100644 index 00000000..b2c5a363 --- /dev/null +++ b/crates/atuin-scripts/migrations/20250326160051_create_scripts.down.sql @@ -0,0 +1,2 @@ +DROP TABLE scripts; +DROP TABLE script_tags;
\ No newline at end of file diff --git a/crates/atuin-scripts/migrations/20250326160051_create_scripts.up.sql b/crates/atuin-scripts/migrations/20250326160051_create_scripts.up.sql new file mode 100644 index 00000000..1b2f3688 --- /dev/null +++ b/crates/atuin-scripts/migrations/20250326160051_create_scripts.up.sql @@ -0,0 +1,17 @@ +-- Add up migration script here +CREATE TABLE scripts ( + id TEXT PRIMARY KEY, + name TEXT NOT NULL, + description TEXT NOT NULL, + shebang TEXT NOT NULL, + script TEXT NOT NULL, + inserted_at INTEGER NOT NULL DEFAULT (strftime('%s', 'now')) +); + +CREATE TABLE script_tags ( + id INTEGER PRIMARY KEY, + script_id TEXT NOT NULL, + tag TEXT NOT NULL +); + +CREATE UNIQUE INDEX idx_script_tags ON script_tags (script_id, tag);
\ No newline at end of file diff --git a/crates/atuin-scripts/migrations/20250402170430_unique_names.down.sql b/crates/atuin-scripts/migrations/20250402170430_unique_names.down.sql new file mode 100644 index 00000000..269b8cd9 --- /dev/null +++ b/crates/atuin-scripts/migrations/20250402170430_unique_names.down.sql @@ -0,0 +1,2 @@ +-- Add down migration script here +alter table scripts drop index name_uniq_idx;
\ No newline at end of file diff --git a/crates/atuin-scripts/migrations/20250402170430_unique_names.up.sql b/crates/atuin-scripts/migrations/20250402170430_unique_names.up.sql new file mode 100644 index 00000000..d2cdd02f --- /dev/null +++ b/crates/atuin-scripts/migrations/20250402170430_unique_names.up.sql @@ -0,0 +1,2 @@ +-- Add up migration script here +create unique index name_uniq_idx ON scripts(name);
\ No newline at end of file diff --git a/crates/atuin-scripts/src/database.rs b/crates/atuin-scripts/src/database.rs new file mode 100644 index 00000000..be113526 --- /dev/null +++ b/crates/atuin-scripts/src/database.rs @@ -0,0 +1,371 @@ +use std::{path::Path, str::FromStr, time::Duration}; + +use atuin_common::utils; +use sqlx::{ + Result, Row, + sqlite::{ + SqliteConnectOptions, SqliteJournalMode, SqlitePool, SqlitePoolOptions, SqliteRow, + SqliteSynchronous, + }, +}; +use tokio::fs; +use tracing::debug; +use uuid::Uuid; + +use crate::store::script::Script; + +#[derive(Debug, Clone)] +pub struct Database { + pub pool: SqlitePool, +} + +impl Database { + pub async fn new(path: impl AsRef<Path>, timeout: f64) -> Result<Self> { + let path = path.as_ref(); + debug!("opening script sqlite database at {:?}", path); + + if utils::broken_symlink(path) { + eprintln!( + "Atuin: Script sqlite db path ({path:?}) is a broken symlink. Unable to read or create replacement." + ); + std::process::exit(1); + } + + if !path.exists() + && let Some(dir) = path.parent() + { + fs::create_dir_all(dir).await?; + } + + let opts = SqliteConnectOptions::from_str(path.as_os_str().to_str().unwrap())? + .journal_mode(SqliteJournalMode::Wal) + .optimize_on_close(true, None) + .synchronous(SqliteSynchronous::Normal) + .with_regexp() + .foreign_keys(true) + .create_if_missing(true); + + let pool = SqlitePoolOptions::new() + .acquire_timeout(Duration::from_secs_f64(timeout)) + .connect_with(opts) + .await?; + + Self::setup_db(&pool).await?; + Ok(Self { pool }) + } + + pub async fn sqlite_version(&self) -> Result<String> { + sqlx::query_scalar("SELECT sqlite_version()") + .fetch_one(&self.pool) + .await + } + + async fn setup_db(pool: &SqlitePool) -> Result<()> { + debug!("running sqlite database setup"); + + sqlx::migrate!("./migrations").run(pool).await?; + + Ok(()) + } + + async fn save_raw(tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, s: &Script) -> Result<()> { + sqlx::query( + "insert or ignore into scripts(id, name, description, shebang, script) + values(?1, ?2, ?3, ?4, ?5)", + ) + .bind(s.id.to_string()) + .bind(s.name.as_str()) + .bind(s.description.as_str()) + .bind(s.shebang.as_str()) + .bind(s.script.as_str()) + .execute(&mut **tx) + .await?; + + for tag in s.tags.iter() { + sqlx::query( + "insert or ignore into script_tags(script_id, tag) + values(?1, ?2)", + ) + .bind(s.id.to_string()) + .bind(tag) + .execute(&mut **tx) + .await?; + } + + Ok(()) + } + + pub async fn save(&self, s: &Script) -> Result<()> { + debug!("saving script to sqlite"); + let mut tx = self.pool.begin().await?; + Self::save_raw(&mut tx, s).await?; + tx.commit().await?; + + Ok(()) + } + + pub async fn save_bulk(&self, s: &[Script]) -> Result<()> { + debug!("saving scripts to sqlite"); + + let mut tx = self.pool.begin().await?; + + for i in s { + Self::save_raw(&mut tx, i).await?; + } + + tx.commit().await?; + + Ok(()) + } + + fn query_script(row: SqliteRow) -> Script { + let id = row.get("id"); + let name = row.get("name"); + let description = row.get("description"); + let shebang = row.get("shebang"); + let script = row.get("script"); + + let id = Uuid::parse_str(id).unwrap(); + + Script { + id, + name, + description, + shebang, + script, + tags: vec![], + } + } + + fn query_script_tags(row: SqliteRow) -> String { + row.get("tag") + } + + #[allow(dead_code)] + async fn load(&self, id: &str) -> Result<Option<Script>> { + debug!("loading script item {}", id); + + let res = sqlx::query("select * from scripts where id = ?1") + .bind(id) + .map(Self::query_script) + .fetch_optional(&self.pool) + .await?; + + // intentionally not joining, don't want to duplicate the script data in memory a whole bunch. + if let Some(mut script) = res { + let tags = sqlx::query("select tag from script_tags where script_id = ?1") + .bind(id) + .map(Self::query_script_tags) + .fetch_all(&self.pool) + .await?; + + script.tags = tags; + Ok(Some(script)) + } else { + Ok(None) + } + } + + pub async fn list(&self) -> Result<Vec<Script>> { + debug!("listing scripts"); + + let mut res = sqlx::query("select * from scripts") + .map(Self::query_script) + .fetch_all(&self.pool) + .await?; + + // Fetch all the tags for each script + for script in res.iter_mut() { + let tags = sqlx::query("select tag from script_tags where script_id = ?1") + .bind(script.id.to_string()) + .map(Self::query_script_tags) + .fetch_all(&self.pool) + .await?; + + script.tags = tags; + } + + Ok(res) + } + + pub async fn clear(&self) -> Result<()> { + debug!("clearing all scripts from sqlite"); + + sqlx::query("delete from script_tags") + .execute(&self.pool) + .await?; + sqlx::query("delete from scripts") + .execute(&self.pool) + .await?; + + Ok(()) + } + + pub async fn delete(&self, id: &str) -> Result<()> { + debug!("deleting script {}", id); + + sqlx::query("delete from scripts where id = ?1") + .bind(id) + .execute(&self.pool) + .await?; + + // delete all the tags for the script + sqlx::query("delete from script_tags where script_id = ?1") + .bind(id) + .execute(&self.pool) + .await?; + + Ok(()) + } + + pub async fn update(&self, s: &Script) -> Result<()> { + debug!("updating script {:?}", s); + + let mut tx = self.pool.begin().await?; + + // Update the script's base fields + sqlx::query("update scripts set name = ?1, description = ?2, shebang = ?3, script = ?4 where id = ?5") + .bind(s.name.as_str()) + .bind(s.description.as_str()) + .bind(s.shebang.as_str()) + .bind(s.script.as_str()) + .bind(s.id.to_string()) + .execute(&mut *tx) + .await?; + + // Delete all existing tags for this script + sqlx::query("delete from script_tags where script_id = ?1") + .bind(s.id.to_string()) + .execute(&mut *tx) + .await?; + + // Insert new tags + for tag in s.tags.iter() { + sqlx::query( + "insert or ignore into script_tags(script_id, tag) + values(?1, ?2)", + ) + .bind(s.id.to_string()) + .bind(tag) + .execute(&mut *tx) + .await?; + } + + tx.commit().await?; + + Ok(()) + } + + pub async fn get_by_name(&self, name: &str) -> Result<Option<Script>> { + let res = sqlx::query("select * from scripts where name = ?1") + .bind(name) + .map(Self::query_script) + .fetch_optional(&self.pool) + .await?; + + let script = if let Some(mut script) = res { + let tags = sqlx::query("select tag from script_tags where script_id = ?1") + .bind(script.id.to_string()) + .map(Self::query_script_tags) + .fetch_all(&self.pool) + .await?; + + script.tags = tags; + Some(script) + } else { + None + }; + + Ok(script) + } +} + +#[cfg(test)] +mod test { + use super::*; + + #[tokio::test] + async fn test_list() { + let db = Database::new("sqlite::memory:", 1.0).await.unwrap(); + let scripts = db.list().await.unwrap(); + assert_eq!(scripts.len(), 0); + + let script = Script::builder() + .name("test".to_string()) + .description("test".to_string()) + .shebang("test".to_string()) + .script("test".to_string()) + .build(); + + db.save(&script).await.unwrap(); + + let scripts = db.list().await.unwrap(); + assert_eq!(scripts.len(), 1); + assert_eq!(scripts[0].name, "test"); + } + + #[tokio::test] + async fn test_save_load() { + let db = Database::new("sqlite::memory:", 1.0).await.unwrap(); + + let script = Script::builder() + .name("test name".to_string()) + .description("test description".to_string()) + .shebang("test shebang".to_string()) + .script("test script".to_string()) + .build(); + + db.save(&script).await.unwrap(); + + let loaded = db.load(&script.id.to_string()).await.unwrap().unwrap(); + + assert_eq!(loaded, script); + } + + #[tokio::test] + async fn test_save_bulk() { + let db = Database::new("sqlite::memory:", 1.0).await.unwrap(); + + let scripts = vec![ + Script::builder() + .name("test name".to_string()) + .description("test description".to_string()) + .shebang("test shebang".to_string()) + .script("test script".to_string()) + .build(), + Script::builder() + .name("test name 2".to_string()) + .description("test description 2".to_string()) + .shebang("test shebang 2".to_string()) + .script("test script 2".to_string()) + .build(), + ]; + + db.save_bulk(&scripts).await.unwrap(); + + let loaded = db.list().await.unwrap(); + assert_eq!(loaded.len(), 2); + assert_eq!(loaded[0].name, "test name"); + assert_eq!(loaded[1].name, "test name 2"); + } + + #[tokio::test] + async fn test_delete() { + let db = Database::new("sqlite::memory:", 1.0).await.unwrap(); + + let script = Script::builder() + .name("test name".to_string()) + .description("test description".to_string()) + .shebang("test shebang".to_string()) + .script("test script".to_string()) + .build(); + + db.save(&script).await.unwrap(); + + assert_eq!(db.list().await.unwrap().len(), 1); + db.delete(&script.id.to_string()).await.unwrap(); + + let loaded = db.list().await.unwrap(); + assert_eq!(loaded.len(), 0); + } +} diff --git a/crates/atuin-scripts/src/execution.rs b/crates/atuin-scripts/src/execution.rs new file mode 100644 index 00000000..5bf94aaa --- /dev/null +++ b/crates/atuin-scripts/src/execution.rs @@ -0,0 +1,286 @@ +use crate::store::script::Script; +use eyre::Result; +use std::collections::{HashMap, HashSet}; +use std::process::Stdio; +use tempfile::NamedTempFile; +use tokio::io::{AsyncReadExt, AsyncWriteExt, BufReader}; +use tokio::sync::mpsc; +use tokio::task; +use tracing::debug; + +// Helper function to build a complete script with shebang +pub fn build_executable_script(script: String, shebang: String) -> String { + if shebang.is_empty() { + // Default to bash if no shebang is provided + format!("#!/usr/bin/env bash\n{script}") + } else if script.starts_with("#!") { + format!("{shebang}\n{script}") + } else { + format!("#!{shebang}\n{script}") + } +} + +/// Represents the communication channels for an interactive script +pub struct ScriptSession { + /// Channel to send input to the script + pub stdin_tx: mpsc::Sender<String>, + /// Exit code of the process once it completes + pub exit_code_rx: mpsc::Receiver<i32>, +} + +impl ScriptSession { + /// Send input to the running script + pub async fn send_input(&self, input: String) -> Result<(), mpsc::error::SendError<String>> { + self.stdin_tx.send(input).await + } + + /// Wait for the script to complete and get the exit code + pub async fn wait_for_exit(&mut self) -> Option<i32> { + self.exit_code_rx.recv().await + } +} + +fn setup_template(script: &Script) -> Result<minijinja::Environment<'_>> { + let mut env = minijinja::Environment::new(); + env.set_trim_blocks(true); + env.add_template("script", script.script.as_str())?; + + Ok(env) +} + +/// Template a script with the given context +pub fn template_script( + script: &Script, + context: &HashMap<String, serde_json::Value>, +) -> Result<String> { + let env = setup_template(script)?; + let template = env.get_template("script")?; + let rendered = template.render(context)?; + + Ok(rendered) +} + +/// Get the variables that need to be templated in a script +pub fn template_variables(script: &Script) -> Result<HashSet<String>> { + let env = setup_template(script)?; + let template = env.get_template("script")?; + + Ok(template.undeclared_variables(true)) +} + +/// Execute a script interactively, allowing for ongoing stdin/stdout interaction +pub async fn execute_script_interactive( + script: String, + shebang: String, +) -> Result<ScriptSession, Box<dyn std::error::Error + Send + Sync>> { + // Create a temporary file for the script + let temp_file = NamedTempFile::new()?; + let temp_path = temp_file.path().to_path_buf(); + + debug!("creating temp file at {}", temp_path.display()); + + // Extract interpreter from shebang for fallback execution + let interpreter = if !shebang.is_empty() { + shebang.trim_start_matches("#!").trim().to_string() + } else { + "/usr/bin/env bash".to_string() + }; + + // Write script content to the temp file, including the shebang + let full_script_content = build_executable_script(script.clone(), shebang.clone()); + + debug!("writing script content to temp file"); + tokio::fs::write(&temp_path, &full_script_content).await?; + + // Make it executable on Unix systems + #[cfg(unix)] + { + debug!("making script executable"); + use std::os::unix::fs::PermissionsExt; + let mut perms = std::fs::metadata(&temp_path)?.permissions(); + perms.set_mode(0o755); + std::fs::set_permissions(&temp_path, perms)?; + } + + // Store the temp_file to prevent it from being dropped + // This ensures it won't be deleted while the script is running + let _keep_temp_file = temp_file; + + debug!("attempting direct script execution"); + let mut child_result = tokio::process::Command::new(temp_path.to_str().unwrap()) + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .spawn(); + + // If direct execution fails, try using the interpreter + if let Err(e) = &child_result { + debug!("direct execution failed: {}, trying with interpreter", e); + + // When falling back to interpreter, remove the shebang from the file + // Some interpreters don't handle scripts with shebangs well + debug!("writing script content without shebang for interpreter execution"); + tokio::fs::write(&temp_path, &script).await?; + + // Parse the interpreter command + let parts: Vec<&str> = interpreter.split_whitespace().collect(); + if !parts.is_empty() { + let mut cmd = tokio::process::Command::new(parts[0]); + + // Add any interpreter args + for i in parts.iter().skip(1) { + cmd.arg(i); + } + + // Add the script path + cmd.arg(temp_path.to_str().unwrap()); + + // Try with the interpreter + child_result = cmd + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .spawn(); + } + } + + // If it still fails, return the error + let mut child = match child_result { + Ok(child) => child, + Err(e) => { + return Err(format!("Failed to execute script: {e}").into()); + } + }; + + // Get handles to stdin, stdout, stderr + let mut stdin = child + .stdin + .take() + .ok_or_else(|| "Failed to open child process stdin".to_string())?; + let stdout = child + .stdout + .take() + .ok_or_else(|| "Failed to open child process stdout".to_string())?; + let stderr = child + .stderr + .take() + .ok_or_else(|| "Failed to open child process stderr".to_string())?; + + // Create channels for the interactive session + let (stdin_tx, mut stdin_rx) = mpsc::channel::<String>(32); + let (exit_code_tx, exit_code_rx) = mpsc::channel::<i32>(1); + + // handle user stdin + debug!("spawning stdin handler"); + tokio::spawn(async move { + while let Some(input) = stdin_rx.recv().await { + if let Err(e) = stdin.write_all(input.as_bytes()).await { + eprintln!("Error writing to stdin: {e}"); + break; + } + if let Err(e) = stdin.flush().await { + eprintln!("Error flushing stdin: {e}"); + break; + } + } + // when the channel closes (sender dropped), we let stdin close naturally + }); + + // handle stdout + debug!("spawning stdout handler"); + let stdout_handle = task::spawn(async move { + let mut stdout_reader = BufReader::new(stdout); + let mut buffer = [0u8; 1024]; + let mut stdout_writer = tokio::io::stdout(); + + loop { + match stdout_reader.read(&mut buffer).await { + Ok(0) => break, // End of stdout + Ok(n) => { + if let Err(e) = stdout_writer.write_all(&buffer[0..n]).await { + eprintln!("Error writing to stdout: {e}"); + break; + } + if let Err(e) = stdout_writer.flush().await { + eprintln!("Error flushing stdout: {e}"); + break; + } + } + Err(e) => { + eprintln!("Error reading from process stdout: {e}"); + break; + } + } + } + }); + + // Process stderr in a separate task + debug!("spawning stderr handler"); + let stderr_handle = task::spawn(async move { + let mut stderr_reader = BufReader::new(stderr); + let mut buffer = [0u8; 1024]; + let mut stderr_writer = tokio::io::stderr(); + + loop { + match stderr_reader.read(&mut buffer).await { + Ok(0) => break, // End of stderr + Ok(n) => { + if let Err(e) = stderr_writer.write_all(&buffer[0..n]).await { + eprintln!("Error writing to stderr: {e}"); + break; + } + if let Err(e) = stderr_writer.flush().await { + eprintln!("Error flushing stderr: {e}"); + break; + } + } + Err(e) => { + eprintln!("Error reading from process stderr: {e}"); + break; + } + } + } + }); + + // Spawn a task to wait for the child process to complete + debug!("spawning exit code handler"); + let _keep_temp_file_clone = _keep_temp_file; + tokio::spawn(async move { + // Keep the temp file alive until the process completes + let _temp_file_ref = _keep_temp_file_clone; + + // Wait for the child process to complete + let status = match child.wait().await { + Ok(status) => { + debug!("Process exited with status: {:?}", status); + status + } + Err(e) => { + eprintln!("Error waiting for child process: {e}"); + // Send a default error code + let _ = exit_code_tx.send(-1).await; + return; + } + }; + + // Wait for stdout/stderr tasks to complete + if let Err(e) = stdout_handle.await { + eprintln!("Error joining stdout task: {e}"); + } + + if let Err(e) = stderr_handle.await { + eprintln!("Error joining stderr task: {e}"); + } + + // Send the exit code + let exit_code = status.code().unwrap_or(-1); + debug!("Sending exit code: {}", exit_code); + let _ = exit_code_tx.send(exit_code).await; + }); + + // Return the communication channels as a ScriptSession + Ok(ScriptSession { + stdin_tx, + exit_code_rx, + }) +} diff --git a/crates/atuin-scripts/src/lib.rs b/crates/atuin-scripts/src/lib.rs new file mode 100644 index 00000000..c79c7089 --- /dev/null +++ b/crates/atuin-scripts/src/lib.rs @@ -0,0 +1,4 @@ +pub mod database; +pub mod execution; +pub mod settings; +pub mod store; diff --git a/crates/atuin-scripts/src/settings.rs b/crates/atuin-scripts/src/settings.rs new file mode 100644 index 00000000..8b137891 --- /dev/null +++ b/crates/atuin-scripts/src/settings.rs @@ -0,0 +1 @@ + diff --git a/crates/atuin-scripts/src/store.rs b/crates/atuin-scripts/src/store.rs new file mode 100644 index 00000000..e70f6909 --- /dev/null +++ b/crates/atuin-scripts/src/store.rs @@ -0,0 +1,114 @@ +use eyre::{Result, bail}; + +use atuin_client::record::sqlite_store::SqliteStore; +use atuin_client::record::{encryption::PASETO_V4, store::Store}; +use atuin_common::record::{Host, HostId, Record, RecordId, RecordIdx}; +use record::ScriptRecord; +use script::{SCRIPT_TAG, SCRIPT_VERSION, Script}; + +use crate::database::Database; + +pub mod record; +pub mod script; + +#[derive(Debug, Clone)] +pub struct ScriptStore { + pub store: SqliteStore, + pub host_id: HostId, + pub encryption_key: [u8; 32], +} + +impl ScriptStore { + pub fn new(store: SqliteStore, host_id: HostId, encryption_key: [u8; 32]) -> Self { + ScriptStore { + store, + host_id, + encryption_key, + } + } + + async fn push_record(&self, record: ScriptRecord) -> Result<(RecordId, RecordIdx)> { + let bytes = record.serialize()?; + let idx = self + .store + .last(self.host_id, SCRIPT_TAG) + .await? + .map_or(0, |p| p.idx + 1); + + let record = Record::builder() + .host(Host::new(self.host_id)) + .version(SCRIPT_VERSION.to_string()) + .tag(SCRIPT_TAG.to_string()) + .idx(idx) + .data(bytes) + .build(); + + let id = record.id; + + self.store + .push(&record.encrypt::<PASETO_V4>(&self.encryption_key)) + .await?; + + Ok((id, idx)) + } + + pub async fn create(&self, script: Script) -> Result<()> { + let record = ScriptRecord::Create(script); + self.push_record(record).await?; + Ok(()) + } + + pub async fn update(&self, script: Script) -> Result<()> { + let record = ScriptRecord::Update(script); + self.push_record(record).await?; + Ok(()) + } + + pub async fn delete(&self, script_id: uuid::Uuid) -> Result<()> { + let record = ScriptRecord::Delete(script_id); + self.push_record(record).await?; + Ok(()) + } + + pub async fn scripts(&self) -> Result<Vec<ScriptRecord>> { + let records = self.store.all_tagged(SCRIPT_TAG).await?; + let mut ret = Vec::with_capacity(records.len()); + + for record in records.into_iter() { + let script = match record.version.as_str() { + SCRIPT_VERSION => { + let decrypted = record.decrypt::<PASETO_V4>(&self.encryption_key)?; + + ScriptRecord::deserialize(&decrypted.data, SCRIPT_VERSION) + } + version => bail!("unknown history version {version:?}"), + }?; + + ret.push(script); + } + + Ok(ret) + } + + pub async fn build(&self, database: Database) -> Result<()> { + // Clear existing data before replaying all records from the store. + // Without this, stale rows can cause unique constraint violations + // when records are replayed (eg name conflicts from renamed scripts). + database.clear().await?; + + // Get all the scripts from the store - they are already sorted by timestamp + let scripts = self.scripts().await?; + + for script in scripts { + match script { + ScriptRecord::Create(script) => { + database.save(&script).await?; + } + ScriptRecord::Update(script) => database.update(&script).await?, + ScriptRecord::Delete(id) => database.delete(&id.to_string()).await?, + } + } + + Ok(()) + } +} diff --git a/crates/atuin-scripts/src/store/record.rs b/crates/atuin-scripts/src/store/record.rs new file mode 100644 index 00000000..4c925be3 --- /dev/null +++ b/crates/atuin-scripts/src/store/record.rs @@ -0,0 +1,215 @@ +use atuin_common::record::DecryptedData; +use eyre::{Result, eyre}; +use uuid::Uuid; + +use crate::store::script::SCRIPT_VERSION; + +use super::script::Script; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ScriptRecord { + Create(Script), + Update(Script), + Delete(Uuid), +} + +impl ScriptRecord { + pub fn serialize(&self) -> Result<DecryptedData> { + use rmp::encode; + + let mut output = vec![]; + + match self { + ScriptRecord::Create(script) => { + // 0 -> a script create + encode::write_u8(&mut output, 0)?; + + let bytes = script.serialize()?; + + encode::write_bin(&mut output, &bytes.0)?; + } + + ScriptRecord::Delete(id) => { + // 1 -> a script delete + encode::write_u8(&mut output, 1)?; + encode::write_str(&mut output, id.to_string().as_str())?; + } + + ScriptRecord::Update(script) => { + // 2 -> a script update + encode::write_u8(&mut output, 2)?; + let bytes = script.serialize()?; + encode::write_bin(&mut output, &bytes.0)?; + } + }; + + Ok(DecryptedData(output)) + } + + pub fn deserialize(data: &DecryptedData, version: &str) -> Result<Self> { + use rmp::decode; + + fn error_report<E: std::fmt::Debug>(err: E) -> eyre::Report { + eyre!("{err:?}") + } + + match version { + SCRIPT_VERSION => { + let mut bytes = decode::Bytes::new(&data.0); + + let record_type = decode::read_u8(&mut bytes).map_err(error_report)?; + + match record_type { + // create + 0 => { + // written by encode::write_bin above + let _ = decode::read_bin_len(&mut bytes).map_err(error_report)?; + let script = Script::deserialize(bytes.remaining_slice())?; + Ok(ScriptRecord::Create(script)) + } + + // delete + 1 => { + let bytes = bytes.remaining_slice(); + let (id, _) = decode::read_str_from_slice(bytes).map_err(error_report)?; + Ok(ScriptRecord::Delete(Uuid::parse_str(id)?)) + } + + // update + 2 => { + // written by encode::write_bin above + let _ = decode::read_bin_len(&mut bytes).map_err(error_report)?; + let script = Script::deserialize(bytes.remaining_slice())?; + Ok(ScriptRecord::Update(script)) + } + + _ => Err(eyre!("unknown script record type {record_type}")), + } + } + _ => Err(eyre!("unknown version {version:?}")), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_serialize_create() { + let script = Script::builder() + .id(uuid::Uuid::parse_str("0195c825a35f7982bdb016168881cbc6").unwrap()) + .name("test".to_string()) + .description("test".to_string()) + .shebang("test".to_string()) + .tags(vec!["test".to_string()]) + .script("test".to_string()) + .build(); + + let record = ScriptRecord::Create(script); + + let serialized = record.serialize().unwrap(); + + assert_eq!( + serialized.0, + vec![ + 204, 0, 196, 65, 150, 217, 36, 48, 49, 57, 53, 99, 56, 50, 53, 45, 97, 51, 53, 102, + 45, 55, 57, 56, 50, 45, 98, 100, 98, 48, 45, 49, 54, 49, 54, 56, 56, 56, 49, 99, + 98, 99, 54, 164, 116, 101, 115, 116, 164, 116, 101, 115, 116, 164, 116, 101, 115, + 116, 145, 164, 116, 101, 115, 116, 164, 116, 101, 115, 116 + ] + ); + } + + #[test] + fn test_serialize_delete() { + let record = ScriptRecord::Delete( + uuid::Uuid::parse_str("0195c825a35f7982bdb016168881cbc6").unwrap(), + ); + + let serialized = record.serialize().unwrap(); + + assert_eq!( + serialized.0, + vec![ + 204, 1, 217, 36, 48, 49, 57, 53, 99, 56, 50, 53, 45, 97, 51, 53, 102, 45, 55, 57, + 56, 50, 45, 98, 100, 98, 48, 45, 49, 54, 49, 54, 56, 56, 56, 49, 99, 98, 99, 54 + ] + ); + } + + #[test] + fn test_serialize_update() { + let script = Script::builder() + .id(uuid::Uuid::parse_str("0195c825a35f7982bdb016168881cbc6").unwrap()) + .name(String::from("test")) + .description(String::from("test")) + .shebang(String::from("test")) + .tags(vec![String::from("test"), String::from("test2")]) + .script(String::from("test")) + .build(); + + let record = ScriptRecord::Update(script); + + let serialized = record.serialize().unwrap(); + + assert_eq!( + serialized.0, + vec![ + 204, 2, 196, 71, 150, 217, 36, 48, 49, 57, 53, 99, 56, 50, 53, 45, 97, 51, 53, 102, + 45, 55, 57, 56, 50, 45, 98, 100, 98, 48, 45, 49, 54, 49, 54, 56, 56, 56, 49, 99, + 98, 99, 54, 164, 116, 101, 115, 116, 164, 116, 101, 115, 116, 164, 116, 101, 115, + 116, 146, 164, 116, 101, 115, 116, 165, 116, 101, 115, 116, 50, 164, 116, 101, 115, + 116 + ], + ); + } + + #[test] + fn test_serialize_deserialize_create() { + let script = Script::builder() + .name("test".to_string()) + .description("test".to_string()) + .shebang("test".to_string()) + .tags(vec!["test".to_string()]) + .script("test".to_string()) + .build(); + + let record = ScriptRecord::Create(script); + + let serialized = record.serialize().unwrap(); + let deserialized = ScriptRecord::deserialize(&serialized, SCRIPT_VERSION).unwrap(); + + assert_eq!(record, deserialized); + } + + #[test] + fn test_serialize_deserialize_delete() { + let record = ScriptRecord::Delete( + uuid::Uuid::parse_str("0195c825a35f7982bdb016168881cbc6").unwrap(), + ); + + let serialized = record.serialize().unwrap(); + let deserialized = ScriptRecord::deserialize(&serialized, SCRIPT_VERSION).unwrap(); + + assert_eq!(record, deserialized); + } + + #[test] + fn test_serialize_deserialize_update() { + let script = Script::builder() + .name("test".to_string()) + .description("test".to_string()) + .shebang("test".to_string()) + .tags(vec!["test".to_string()]) + .script("test".to_string()) + .build(); + + let record = ScriptRecord::Update(script); + + let serialized = record.serialize().unwrap(); + let deserialized = ScriptRecord::deserialize(&serialized, SCRIPT_VERSION).unwrap(); + + assert_eq!(record, deserialized); + } +} diff --git a/crates/atuin-scripts/src/store/script.rs b/crates/atuin-scripts/src/store/script.rs new file mode 100644 index 00000000..af180320 --- /dev/null +++ b/crates/atuin-scripts/src/store/script.rs @@ -0,0 +1,151 @@ +use atuin_common::record::DecryptedData; +use eyre::{Result, bail, ensure}; +use uuid::Uuid; + +use rmp::{ + decode::{self, Bytes}, + encode, +}; +use typed_builder::TypedBuilder; + +pub const SCRIPT_VERSION: &str = "v0"; +pub const SCRIPT_TAG: &str = "script"; +pub const SCRIPT_LEN: usize = 20000; // 20kb max total len + +#[derive(Debug, Clone, PartialEq, Eq, TypedBuilder)] +/// A script is a set of commands that can be run, with the specified shebang +pub struct Script { + /// The id of the script + #[builder(default = uuid::Uuid::new_v4())] + pub id: Uuid, + + /// The name of the script + pub name: String, + + /// The description of the script + #[builder(default = String::new())] + pub description: String, + + /// The interpreter of the script + #[builder(default = String::new())] + pub shebang: String, + + /// The tags of the script + #[builder(default = Vec::new())] + pub tags: Vec<String>, + + /// The script content + pub script: String, +} + +impl Script { + pub fn serialize(&self) -> Result<DecryptedData> { + // sort the tags first, to ensure consistent ordering + let mut tags = self.tags.clone(); + tags.sort(); + + let mut output = vec![]; + + encode::write_array_len(&mut output, 6)?; + encode::write_str(&mut output, &self.id.to_string())?; + encode::write_str(&mut output, &self.name)?; + encode::write_str(&mut output, &self.description)?; + encode::write_str(&mut output, &self.shebang)?; + encode::write_array_len(&mut output, self.tags.len() as u32)?; + + for tag in &tags { + encode::write_str(&mut output, tag)?; + } + + encode::write_str(&mut output, &self.script)?; + + Ok(DecryptedData(output)) + } + + pub fn deserialize(bytes: &[u8]) -> Result<Self> { + let mut bytes = decode::Bytes::new(bytes); + let nfields = decode::read_array_len(&mut bytes).unwrap(); + + ensure!(nfields == 6, "too many entries in v0 script record"); + + let bytes = bytes.remaining_slice(); + + let (id, bytes) = decode::read_str_from_slice(bytes).unwrap(); + let (name, bytes) = decode::read_str_from_slice(bytes).unwrap(); + let (description, bytes) = decode::read_str_from_slice(bytes).unwrap(); + let (shebang, bytes) = decode::read_str_from_slice(bytes).unwrap(); + + let mut bytes = Bytes::new(bytes); + let tags_len = decode::read_array_len(&mut bytes).unwrap(); + + let mut bytes = bytes.remaining_slice(); + + let mut tags = Vec::new(); + for _ in 0..tags_len { + let (tag, remaining) = decode::read_str_from_slice(bytes).unwrap(); + tags.push(tag.to_owned()); + bytes = remaining; + } + + let (script, bytes) = decode::read_str_from_slice(bytes).unwrap(); + + if !bytes.is_empty() { + bail!("trailing bytes in encoded script record. malformed") + } + + Ok(Script { + id: Uuid::parse_str(id).unwrap(), + name: name.to_owned(), + description: description.to_owned(), + shebang: shebang.to_owned(), + tags, + script: script.to_owned(), + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_serialize() { + let script = Script { + id: uuid::Uuid::parse_str("0195c825a35f7982bdb016168881cbc6").unwrap(), + name: "test".to_string(), + description: "test".to_string(), + shebang: "test".to_string(), + tags: vec!["test".to_string()], + script: "test".to_string(), + }; + + let serialized = script.serialize().unwrap(); + assert_eq!( + serialized.0, + vec![ + 150, 217, 36, 48, 49, 57, 53, 99, 56, 50, 53, 45, 97, 51, 53, 102, 45, 55, 57, 56, + 50, 45, 98, 100, 98, 48, 45, 49, 54, 49, 54, 56, 56, 56, 49, 99, 98, 99, 54, 164, + 116, 101, 115, 116, 164, 116, 101, 115, 116, 164, 116, 101, 115, 116, 145, 164, + 116, 101, 115, 116, 164, 116, 101, 115, 116 + ] + ); + } + + #[test] + fn test_serialize_deserialize() { + let script = Script { + id: uuid::Uuid::new_v4(), + name: "test".to_string(), + description: "test".to_string(), + shebang: "test".to_string(), + tags: vec!["test".to_string()], + script: "test".to_string(), + }; + + let serialized = script.serialize().unwrap(); + + let deserialized = Script::deserialize(&serialized.0).unwrap(); + + assert_eq!(script, deserialized); + } +} diff --git a/crates/atuin-server-database/Cargo.toml b/crates/atuin-server-database/Cargo.toml new file mode 100644 index 00000000..61c86a49 --- /dev/null +++ b/crates/atuin-server-database/Cargo.toml @@ -0,0 +1,20 @@ +[package] +name = "atuin-server-database" +edition = "2024" +description = "server database library for atuin" + +version = { workspace = true } +authors = { workspace = true } +license = { workspace = true } +homepage = { workspace = true } +repository = { workspace = true } + +[dependencies] +atuin-common = { path = "../atuin-common", version = "18.13.2" } + +tracing = { workspace = true } +time = { workspace = true } +eyre = { workspace = true } +serde = { workspace = true } +async-trait = { workspace = true } +url = "2.5.2" diff --git a/crates/atuin-server-database/src/calendar.rs b/crates/atuin-server-database/src/calendar.rs new file mode 100644 index 00000000..2229667b --- /dev/null +++ b/crates/atuin-server-database/src/calendar.rs @@ -0,0 +1,18 @@ +// Calendar data + +use serde::{Deserialize, Serialize}; +use time::Month; + +pub enum TimePeriod { + Year, + Month { year: i32 }, + Day { year: i32, month: Month }, +} + +#[derive(Debug, Serialize, Deserialize)] +pub struct TimePeriodInfo { + pub count: u64, + + // TODO: Use this for merkle tree magic + pub hash: String, +} diff --git a/crates/atuin-server-database/src/lib.rs b/crates/atuin-server-database/src/lib.rs new file mode 100644 index 00000000..6000a530 --- /dev/null +++ b/crates/atuin-server-database/src/lib.rs @@ -0,0 +1,226 @@ +#![forbid(unsafe_code)] + +pub mod calendar; +pub mod models; + +use std::{ + collections::HashMap, + fmt::{Debug, Display}, + ops::Range, +}; + +use self::{ + calendar::{TimePeriod, TimePeriodInfo}, + models::{History, NewHistory, NewSession, NewUser, Session, User}, +}; +use async_trait::async_trait; +use atuin_common::record::{EncryptedData, HostId, Record, RecordIdx, RecordStatus}; +use serde::{Deserialize, Serialize}; +use time::{Date, Duration, Month, OffsetDateTime, Time, UtcOffset}; +use tracing::instrument; + +#[derive(Debug)] +pub enum DbError { + NotFound, + Other(eyre::Report), +} + +impl Display for DbError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{self:?}") + } +} + +impl<T: std::error::Error + Into<time::error::Error>> From<T> for DbError { + fn from(value: T) -> Self { + DbError::Other(value.into().into()) + } +} + +impl std::error::Error for DbError {} + +pub type DbResult<T> = Result<T, DbError>; + +#[derive(Debug, PartialEq)] +pub enum DbType { + Postgres, + Sqlite, + Unknown, +} + +#[derive(Clone, Deserialize, Serialize)] +pub struct DbSettings { + pub db_uri: String, + /// Optional URI for read replicas. If set, read-only queries will use this connection. + pub read_db_uri: Option<String>, +} + +impl DbSettings { + pub fn db_type(&self) -> DbType { + if self.db_uri.starts_with("postgres://") || self.db_uri.starts_with("postgresql://") { + DbType::Postgres + } else if self.db_uri.starts_with("sqlite://") { + DbType::Sqlite + } else { + DbType::Unknown + } + } +} + +fn redact_db_uri(uri: &str) -> String { + url::Url::parse(uri) + .map(|mut url| { + let _ = url.set_password(Some("****")); + url.to_string() + }) + .unwrap_or_else(|_| uri.to_string()) +} + +// Do our best to redact passwords so they're not logged in the event of an error. +impl Debug for DbSettings { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + if self.db_type() == DbType::Postgres { + let redacted_uri = redact_db_uri(&self.db_uri); + let redacted_read_uri = self.read_db_uri.as_ref().map(|uri| redact_db_uri(uri)); + f.debug_struct("DbSettings") + .field("db_uri", &redacted_uri) + .field("read_db_uri", &redacted_read_uri) + .finish() + } else { + f.debug_struct("DbSettings") + .field("db_uri", &self.db_uri) + .field("read_db_uri", &self.read_db_uri) + .finish() + } + } +} + +#[async_trait] +pub trait Database: Sized + Clone + Send + Sync + 'static { + async fn new(settings: &DbSettings) -> DbResult<Self>; + + async fn get_session(&self, token: &str) -> DbResult<Session>; + async fn get_session_user(&self, token: &str) -> DbResult<User>; + async fn add_session(&self, session: &NewSession) -> DbResult<()>; + + async fn get_user(&self, username: &str) -> DbResult<User>; + async fn get_user_session(&self, u: &User) -> DbResult<Session>; + async fn add_user(&self, user: &NewUser) -> DbResult<i64>; + + async fn update_user_password(&self, u: &User) -> DbResult<()>; + + async fn count_history(&self, user: &User) -> DbResult<i64>; + async fn count_history_cached(&self, user: &User) -> DbResult<i64>; + + async fn delete_user(&self, u: &User) -> DbResult<()>; + async fn delete_history(&self, user: &User, id: String) -> DbResult<()>; + async fn deleted_history(&self, user: &User) -> DbResult<Vec<String>>; + async fn delete_store(&self, user: &User) -> DbResult<()>; + + async fn add_records(&self, user: &User, record: &[Record<EncryptedData>]) -> DbResult<()>; + async fn next_records( + &self, + user: &User, + host: HostId, + tag: String, + start: Option<RecordIdx>, + count: u64, + ) -> DbResult<Vec<Record<EncryptedData>>>; + + // Return the tail record ID for each store, so (HostID, Tag, TailRecordID) + async fn status(&self, user: &User) -> DbResult<RecordStatus>; + + async fn count_history_range(&self, user: &User, range: Range<OffsetDateTime>) + -> DbResult<i64>; + + async fn list_history( + &self, + user: &User, + created_after: OffsetDateTime, + since: OffsetDateTime, + host: &str, + page_size: i64, + ) -> DbResult<Vec<History>>; + + async fn add_history(&self, history: &[NewHistory]) -> DbResult<()>; + + async fn oldest_history(&self, user: &User) -> DbResult<History>; + + #[instrument(skip_all)] + async fn calendar( + &self, + user: &User, + period: TimePeriod, + tz: UtcOffset, + ) -> DbResult<HashMap<u64, TimePeriodInfo>> { + let mut ret = HashMap::new(); + let iter: Box<dyn Iterator<Item = DbResult<(u64, Range<Date>)>> + Send> = match period { + TimePeriod::Year => { + // First we need to work out how far back to calculate. Get the + // oldest history item + let oldest = self + .oldest_history(user) + .await? + .timestamp + .to_offset(tz) + .year(); + let current_year = OffsetDateTime::now_utc().to_offset(tz).year(); + + // All the years we need to get data for + // The upper bound is exclusive, so include current +1 + let years = oldest..current_year + 1; + + Box::new(years.map(|year| { + let start = Date::from_calendar_date(year, time::Month::January, 1)?; + let end = Date::from_calendar_date(year + 1, time::Month::January, 1)?; + + Ok((year as u64, start..end)) + })) + } + + TimePeriod::Month { year } => { + let months = + std::iter::successors(Some(Month::January), |m| Some(m.next())).take(12); + + Box::new(months.map(move |month| { + let start = Date::from_calendar_date(year, month, 1)?; + let days = start.month().length(year); + let end = start + Duration::days(days as i64); + + Ok((month as u64, start..end)) + })) + } + + TimePeriod::Day { year, month } => { + let days = 1..month.length(year); + Box::new(days.map(move |day| { + let start = Date::from_calendar_date(year, month, day)?; + let end = start + .next_day() + .ok_or_else(|| DbError::Other(eyre::eyre!("no next day?")))?; + + Ok((day as u64, start..end)) + })) + } + }; + + for x in iter { + let (index, range) = x?; + + let start = range.start.with_time(Time::MIDNIGHT).assume_offset(tz); + let end = range.end.with_time(Time::MIDNIGHT).assume_offset(tz); + + let count = self.count_history_range(user, start..end).await?; + + ret.insert( + index, + TimePeriodInfo { + count: count as u64, + hash: "".to_string(), + }, + ); + } + + Ok(ret) + } +} diff --git a/crates/atuin-server-database/src/models.rs b/crates/atuin-server-database/src/models.rs new file mode 100644 index 00000000..b71a9bc9 --- /dev/null +++ b/crates/atuin-server-database/src/models.rs @@ -0,0 +1,52 @@ +use time::OffsetDateTime; + +pub struct History { + pub id: i64, + pub client_id: String, // a client generated ID + pub user_id: i64, + pub hostname: String, + pub timestamp: OffsetDateTime, + + /// All the data we have about this command, encrypted. + /// + /// Currently this is an encrypted msgpack object, but this may change in the future. + pub data: String, + + pub created_at: OffsetDateTime, +} + +pub struct NewHistory { + pub client_id: String, + pub user_id: i64, + pub hostname: String, + pub timestamp: OffsetDateTime, + + /// All the data we have about this command, encrypted. + /// + /// Currently this is an encrypted msgpack object, but this may change in the future. + pub data: String, +} + +pub struct User { + pub id: i64, + pub username: String, + pub email: String, + pub password: String, +} + +pub struct Session { + pub id: i64, + pub user_id: i64, + pub token: String, +} + +pub struct NewUser { + pub username: String, + pub email: String, + pub password: String, +} + +pub struct NewSession { + pub user_id: i64, + pub token: String, +} diff --git a/crates/atuin-server-postgres/Cargo.toml b/crates/atuin-server-postgres/Cargo.toml new file mode 100644 index 00000000..35be92ac --- /dev/null +++ b/crates/atuin-server-postgres/Cargo.toml @@ -0,0 +1,25 @@ +[package] +name = "atuin-server-postgres" +edition = "2024" +description = "server postgres database library for atuin" + +version = { workspace = true } +authors = { workspace = true } +license = { workspace = true } +homepage = { workspace = true } +repository = { workspace = true } + +[dependencies] +atuin-common = { path = "../atuin-common", version = "18.13.2" } +atuin-server-database = { path = "../atuin-server-database", version = "18.13.2" } + +eyre = { workspace = true } +tracing = { workspace = true } +time = { workspace = true } +serde = { workspace = true } +sqlx = { workspace = true } +async-trait = { workspace = true } +uuid = { workspace = true } +metrics = "0.24" +futures-util = "0.3" +rand.workspace = true
\ No newline at end of file diff --git a/crates/atuin-server-postgres/build.rs b/crates/atuin-server-postgres/build.rs new file mode 100644 index 00000000..d5068697 --- /dev/null +++ b/crates/atuin-server-postgres/build.rs @@ -0,0 +1,5 @@ +// generated by `sqlx migrate build-script` +fn main() { + // trigger recompilation when a new migration is added + println!("cargo:rerun-if-changed=migrations"); +} diff --git a/crates/atuin-server-postgres/migrations/20210425153745_create_history.sql b/crates/atuin-server-postgres/migrations/20210425153745_create_history.sql new file mode 100644 index 00000000..2c2d17b0 --- /dev/null +++ b/crates/atuin-server-postgres/migrations/20210425153745_create_history.sql @@ -0,0 +1,11 @@ +create table history ( + id bigserial primary key, + client_id text not null unique, -- the client-generated ID + user_id bigserial not null, -- allow multiple users + hostname text not null, -- a unique identifier from the client (can be hashed, random, whatever) + timestamp timestamp not null, -- one of the few non-encrypted metadatas + + data varchar(8192) not null, -- store the actual history data, encrypted. I don't wanna know! + + created_at timestamp not null default current_timestamp +); diff --git a/crates/atuin-server-postgres/migrations/20210425153757_create_users.sql b/crates/atuin-server-postgres/migrations/20210425153757_create_users.sql new file mode 100644 index 00000000..a25dcced --- /dev/null +++ b/crates/atuin-server-postgres/migrations/20210425153757_create_users.sql @@ -0,0 +1,10 @@ +create table users ( + id bigserial primary key, -- also store our own ID + username varchar(32) not null unique, -- being able to contact users is useful + email varchar(128) not null unique, -- being able to contact users is useful + password varchar(128) not null unique +); + +-- the prior index is case sensitive :( +CREATE UNIQUE INDEX email_unique_idx on users (LOWER(email)); +CREATE UNIQUE INDEX username_unique_idx on users (LOWER(username)); diff --git a/crates/atuin-server-postgres/migrations/20210425153800_create_sessions.sql b/crates/atuin-server-postgres/migrations/20210425153800_create_sessions.sql new file mode 100644 index 00000000..c2fb6559 --- /dev/null +++ b/crates/atuin-server-postgres/migrations/20210425153800_create_sessions.sql @@ -0,0 +1,6 @@ +-- Add migration script here +create table sessions ( + id bigserial primary key, + user_id bigserial, + token varchar(128) unique not null +); diff --git a/crates/atuin-server-postgres/migrations/20220419082412_add_count_trigger.sql b/crates/atuin-server-postgres/migrations/20220419082412_add_count_trigger.sql new file mode 100644 index 00000000..dd1afa88 --- /dev/null +++ b/crates/atuin-server-postgres/migrations/20220419082412_add_count_trigger.sql @@ -0,0 +1,51 @@ +-- Prior to this, the count endpoint was super naive and just ran COUNT(1). +-- This is slow asf. Now that we have an amount of actual traffic, +-- stop doing that! +-- This basically maintains a count, so we can read ONE row, instead of ALL the +-- rows. Much better. +-- Future optimisation could use some sort of cache so we don't even need to hit +-- postgres at all. + +create table total_history_count_user( + id bigserial primary key, + user_id bigserial, + total integer -- try and avoid using keywords - hence total, not count +); + +create or replace function user_history_count() +returns trigger as +$func$ +begin + if (TG_OP='INSERT') then + update total_history_count_user set total = total + 1 where user_id = new.user_id; + + if not found then + insert into total_history_count_user(user_id, total) + values ( + new.user_id, + (select count(1) from history where user_id = new.user_id) + ); + end if; + + elsif (TG_OP='DELETE') then + update total_history_count_user set total = total - 1 where user_id = new.user_id; + + if not found then + insert into total_history_count_user(user_id, total) + values ( + new.user_id, + (select count(1) from history where user_id = new.user_id) + ); + end if; + end if; + + return NEW; -- this is actually ignored for an after trigger, but oh well +end; +$func$ +language plpgsql volatile -- pldfplplpflh +cost 100; -- default value + +create trigger tg_user_history_count + after insert or delete on history + for each row + execute procedure user_history_count(); diff --git a/crates/atuin-server-postgres/migrations/20220421073605_fix_count_trigger_delete.sql b/crates/atuin-server-postgres/migrations/20220421073605_fix_count_trigger_delete.sql new file mode 100644 index 00000000..6198f300 --- /dev/null +++ b/crates/atuin-server-postgres/migrations/20220421073605_fix_count_trigger_delete.sql @@ -0,0 +1,35 @@ +-- the old version of this function used NEW in the delete part when it should +-- use OLD + +create or replace function user_history_count() +returns trigger as +$func$ +begin + if (TG_OP='INSERT') then + update total_history_count_user set total = total + 1 where user_id = new.user_id; + + if not found then + insert into total_history_count_user(user_id, total) + values ( + new.user_id, + (select count(1) from history where user_id = new.user_id) + ); + end if; + + elsif (TG_OP='DELETE') then + update total_history_count_user set total = total - 1 where user_id = old.user_id; + + if not found then + insert into total_history_count_user(user_id, total) + values ( + old.user_id, + (select count(1) from history where user_id = old.user_id) + ); + end if; + end if; + + return NEW; -- this is actually ignored for an after trigger, but oh well +end; +$func$ +language plpgsql volatile -- pldfplplpflh +cost 100; -- default value diff --git a/crates/atuin-server-postgres/migrations/20220421174016_larger-commands.sql b/crates/atuin-server-postgres/migrations/20220421174016_larger-commands.sql new file mode 100644 index 00000000..0ac43433 --- /dev/null +++ b/crates/atuin-server-postgres/migrations/20220421174016_larger-commands.sql @@ -0,0 +1,3 @@ +-- Make it 4x larger. Most commands are less than this, but as it's base64 +-- SOME are more than 8192. Should be enough for now. +ALTER TABLE history ALTER COLUMN data TYPE varchar(32768); diff --git a/crates/atuin-server-postgres/migrations/20220426172813_user-created-at.sql b/crates/atuin-server-postgres/migrations/20220426172813_user-created-at.sql new file mode 100644 index 00000000..a9138194 --- /dev/null +++ b/crates/atuin-server-postgres/migrations/20220426172813_user-created-at.sql @@ -0,0 +1 @@ +alter table users add column created_at timestamp not null default now(); diff --git a/crates/atuin-server-postgres/migrations/20220505082442_create-events.sql b/crates/atuin-server-postgres/migrations/20220505082442_create-events.sql new file mode 100644 index 00000000..57e16ec7 --- /dev/null +++ b/crates/atuin-server-postgres/migrations/20220505082442_create-events.sql @@ -0,0 +1,14 @@ +create type event_type as enum ('create', 'delete'); + +create table events ( + id bigserial primary key, + client_id text not null unique, -- the client-generated ID + user_id bigserial not null, -- allow multiple users + hostname text not null, -- a unique identifier from the client (can be hashed, random, whatever) + timestamp timestamp not null, -- one of the few non-encrypted metadatas + + event_type event_type, + data text not null, -- store the actual history data, encrypted. I don't wanna know! + + created_at timestamp not null default current_timestamp +); diff --git a/crates/atuin-server-postgres/migrations/20220610074049_history-length.sql b/crates/atuin-server-postgres/migrations/20220610074049_history-length.sql new file mode 100644 index 00000000..b1c23016 --- /dev/null +++ b/crates/atuin-server-postgres/migrations/20220610074049_history-length.sql @@ -0,0 +1,2 @@ +-- Add migration script here +alter table history alter column data type text; diff --git a/crates/atuin-server-postgres/migrations/20230315220537_drop-events.sql b/crates/atuin-server-postgres/migrations/20230315220537_drop-events.sql new file mode 100644 index 00000000..fe3cae17 --- /dev/null +++ b/crates/atuin-server-postgres/migrations/20230315220537_drop-events.sql @@ -0,0 +1,2 @@ +-- Add migration script here +drop table events; diff --git a/crates/atuin-server-postgres/migrations/20230315224203_create-deleted.sql b/crates/atuin-server-postgres/migrations/20230315224203_create-deleted.sql new file mode 100644 index 00000000..9a9e6263 --- /dev/null +++ b/crates/atuin-server-postgres/migrations/20230315224203_create-deleted.sql @@ -0,0 +1,5 @@ +-- Add migration script here +alter table history add column if not exists deleted_at timestamp; + +-- queries will all be selecting the ids of history for a user, that has been deleted +create index if not exists history_deleted_index on history(client_id, user_id, deleted_at); diff --git a/crates/atuin-server-postgres/migrations/20230515221038_trigger-delete-only.sql b/crates/atuin-server-postgres/migrations/20230515221038_trigger-delete-only.sql new file mode 100644 index 00000000..3d0bba52 --- /dev/null +++ b/crates/atuin-server-postgres/migrations/20230515221038_trigger-delete-only.sql @@ -0,0 +1,30 @@ +-- We do not need to run the trigger on deletes, as the only time we are deleting history is when the user +-- has already been deleted +-- This actually slows down deleting all the history a good bit! + +create or replace function user_history_count() +returns trigger as +$func$ +begin + if (TG_OP='INSERT') then + update total_history_count_user set total = total + 1 where user_id = new.user_id; + + if not found then + insert into total_history_count_user(user_id, total) + values ( + new.user_id, + (select count(1) from history where user_id = new.user_id) + ); + end if; + end if; + + return NEW; -- this is actually ignored for an after trigger, but oh well +end; +$func$ +language plpgsql volatile -- pldfplplpflh +cost 100; -- default value + +create or replace trigger tg_user_history_count + after insert on history + for each row + execute procedure user_history_count(); diff --git a/crates/atuin-server-postgres/migrations/20230623070418_records.sql b/crates/atuin-server-postgres/migrations/20230623070418_records.sql new file mode 100644 index 00000000..22437595 --- /dev/null +++ b/crates/atuin-server-postgres/migrations/20230623070418_records.sql @@ -0,0 +1,15 @@ +-- Add migration script here +create table records ( + id uuid primary key, -- remember to use uuidv7 for happy indices <3 + client_id uuid not null, -- I am too uncomfortable with the idea of a client-generated primary key + host uuid not null, -- a unique identifier for the host + parent uuid default null, -- the ID of the parent record, bearing in mind this is a linked list + timestamp bigint not null, -- not a timestamp type, as those do not have nanosecond precision + version text not null, + tag text not null, -- what is this? history, kv, whatever. Remember clients get a log per tag per host + data text not null, -- store the actual history data, encrypted. I don't wanna know! + cek text not null, + + user_id bigint not null, -- allow multiple users + created_at timestamp not null default current_timestamp +); diff --git a/crates/atuin-server-postgres/migrations/20231202170508_create-store.sql b/crates/atuin-server-postgres/migrations/20231202170508_create-store.sql new file mode 100644 index 00000000..ffb57966 --- /dev/null +++ b/crates/atuin-server-postgres/migrations/20231202170508_create-store.sql @@ -0,0 +1,15 @@ +-- Add migration script here +create table store ( + id uuid primary key, -- remember to use uuidv7 for happy indices <3 + client_id uuid not null, -- I am too uncomfortable with the idea of a client-generated primary key, even though it's fine mathematically + host uuid not null, -- a unique identifier for the host + idx bigint not null, -- the index of the record in this store, identified by (host, tag) + timestamp bigint not null, -- not a timestamp type, as those do not have nanosecond precision + version text not null, + tag text not null, -- what is this? history, kv, whatever. Remember clients get a log per tag per host + data text not null, -- store the actual history data, encrypted. I don't wanna know! + cek text not null, + + user_id bigint not null, -- allow multiple users + created_at timestamp not null default current_timestamp +); diff --git a/crates/atuin-server-postgres/migrations/20231203124112_create-store-idx.sql b/crates/atuin-server-postgres/migrations/20231203124112_create-store-idx.sql new file mode 100644 index 00000000..56d67145 --- /dev/null +++ b/crates/atuin-server-postgres/migrations/20231203124112_create-store-idx.sql @@ -0,0 +1,2 @@ +-- Add migration script here +create unique index record_uniq ON store(user_id, host, tag, idx); diff --git a/crates/atuin-server-postgres/migrations/20240108124837_drop-some-defaults.sql b/crates/atuin-server-postgres/migrations/20240108124837_drop-some-defaults.sql new file mode 100644 index 00000000..ad2af5a1 --- /dev/null +++ b/crates/atuin-server-postgres/migrations/20240108124837_drop-some-defaults.sql @@ -0,0 +1,4 @@ +-- Add migration script here +alter table history alter column user_id drop default; +alter table sessions alter column user_id drop default; +alter table total_history_count_user alter column user_id drop default; diff --git a/crates/atuin-server-postgres/migrations/20240614104159_idx-cache.sql b/crates/atuin-server-postgres/migrations/20240614104159_idx-cache.sql new file mode 100644 index 00000000..76425ed7 --- /dev/null +++ b/crates/atuin-server-postgres/migrations/20240614104159_idx-cache.sql @@ -0,0 +1,8 @@ +create table store_idx_cache( + id bigserial primary key, + user_id bigint, + + host uuid, + tag text, + idx bigint +); diff --git a/crates/atuin-server-postgres/migrations/20240621110731_user-verified.sql b/crates/atuin-server-postgres/migrations/20240621110731_user-verified.sql new file mode 100644 index 00000000..6eba02ec --- /dev/null +++ b/crates/atuin-server-postgres/migrations/20240621110731_user-verified.sql @@ -0,0 +1,8 @@ +alter table users add verified_at timestamp with time zone default null; + +create table user_verification_token( + id bigserial primary key, + user_id bigint unique references users(id), + token text, + valid_until timestamp with time zone +); diff --git a/crates/atuin-server-postgres/migrations/20240702094825_idx_cache_index.sql b/crates/atuin-server-postgres/migrations/20240702094825_idx_cache_index.sql new file mode 100644 index 00000000..d1a7b194 --- /dev/null +++ b/crates/atuin-server-postgres/migrations/20240702094825_idx_cache_index.sql @@ -0,0 +1 @@ +create unique index store_idx_cache_uniq on store_idx_cache(user_id, host, tag); diff --git a/crates/atuin-server-postgres/migrations/20260127000000_remove-email-verification.sql b/crates/atuin-server-postgres/migrations/20260127000000_remove-email-verification.sql new file mode 100644 index 00000000..15309920 --- /dev/null +++ b/crates/atuin-server-postgres/migrations/20260127000000_remove-email-verification.sql @@ -0,0 +1,2 @@ +drop table if exists user_verification_token; +alter table users drop column if exists verified_at; diff --git a/crates/atuin-server-postgres/src/lib.rs b/crates/atuin-server-postgres/src/lib.rs new file mode 100644 index 00000000..ce101d8d --- /dev/null +++ b/crates/atuin-server-postgres/src/lib.rs @@ -0,0 +1,639 @@ +use std::collections::HashMap; +use std::ops::Range; + +use rand::Rng; + +use async_trait::async_trait; +use atuin_common::record::{EncryptedData, HostId, Record, RecordIdx, RecordStatus}; +use atuin_server_database::models::{History, NewHistory, NewSession, NewUser, Session, User}; +use atuin_server_database::{Database, DbError, DbResult, DbSettings}; +use futures_util::TryStreamExt; +use sqlx::Row; +use sqlx::postgres::PgPoolOptions; + +use time::{OffsetDateTime, PrimitiveDateTime, UtcOffset}; +use tracing::instrument; +use uuid::Uuid; +use wrappers::{DbHistory, DbRecord, DbSession, DbUser}; + +mod wrappers; + +const MIN_PG_VERSION: u32 = 14; + +#[derive(Clone)] +pub struct Postgres { + pool: sqlx::Pool<sqlx::postgres::Postgres>, + /// Optional read replica pool for read-only queries + read_pool: Option<sqlx::Pool<sqlx::postgres::Postgres>>, +} + +impl Postgres { + /// Returns the appropriate pool for read operations. + /// Uses read_pool if available, otherwise falls back to the primary pool. + fn read_pool(&self) -> &sqlx::Pool<sqlx::postgres::Postgres> { + self.read_pool.as_ref().unwrap_or(&self.pool) + } +} + +fn fix_error(error: sqlx::Error) -> DbError { + match error { + sqlx::Error::RowNotFound => DbError::NotFound, + error => DbError::Other(error.into()), + } +} + +#[async_trait] +impl Database for Postgres { + async fn new(settings: &DbSettings) -> DbResult<Self> { + let pool = PgPoolOptions::new() + .max_connections(100) + .connect(settings.db_uri.as_str()) + .await + .map_err(fix_error)?; + + // Call server_version_num to get the DB server's major version number + // The call returns None for servers older than 8.x. + let pg_major_version: u32 = pool + .acquire() + .await + .map_err(fix_error)? + .server_version_num() + .ok_or(DbError::Other(eyre::Report::msg( + "could not get PostgreSQL version", + )))? + / 10000; + + if pg_major_version < MIN_PG_VERSION { + return Err(DbError::Other(eyre::Report::msg(format!( + "unsupported PostgreSQL version {pg_major_version}, minimum required is {MIN_PG_VERSION}" + )))); + } + + sqlx::migrate!("./migrations") + .run(&pool) + .await + .map_err(|error| DbError::Other(error.into()))?; + + // Create read replica pool if configured + let read_pool = if let Some(read_db_uri) = &settings.read_db_uri { + tracing::info!("Connecting to read replica database"); + let read_pool = PgPoolOptions::new() + .max_connections(100) + .connect(read_db_uri.as_str()) + .await + .map_err(fix_error)?; + + // Verify the read replica is also a supported PostgreSQL version + let read_pg_major_version: u32 = read_pool + .acquire() + .await + .map_err(fix_error)? + .server_version_num() + .ok_or(DbError::Other(eyre::Report::msg( + "could not get PostgreSQL version from read replica", + )))? + / 10000; + + if read_pg_major_version < MIN_PG_VERSION { + return Err(DbError::Other(eyre::Report::msg(format!( + "unsupported PostgreSQL version {read_pg_major_version} on read replica, minimum required is {MIN_PG_VERSION}" + )))); + } + + Some(read_pool) + } else { + None + }; + + Ok(Self { pool, read_pool }) + } + + #[instrument(skip_all)] + async fn get_session(&self, token: &str) -> DbResult<Session> { + sqlx::query_as("select id, user_id, token from sessions where token = $1") + .bind(token) + .fetch_one(self.read_pool()) + .await + .map_err(fix_error) + .map(|DbSession(session)| session) + } + + #[instrument(skip_all)] + async fn get_user(&self, username: &str) -> DbResult<User> { + sqlx::query_as("select id, username, email, password from users where username = $1") + .bind(username) + .fetch_one(self.read_pool()) + .await + .map_err(fix_error) + .map(|DbUser(user)| user) + } + + #[instrument(skip_all)] + async fn get_session_user(&self, token: &str) -> DbResult<User> { + sqlx::query_as( + "select users.id, users.username, users.email, users.password from users + inner join sessions + on users.id = sessions.user_id + and sessions.token = $1", + ) + .bind(token) + .fetch_one(self.read_pool()) + .await + .map_err(fix_error) + .map(|DbUser(user)| user) + } + + #[instrument(skip_all)] + async fn count_history(&self, user: &User) -> DbResult<i64> { + // The cache is new, and the user might not yet have a cache value. + // They will have one as soon as they post up some new history, but handle that + // edge case. + + let res: (i64,) = sqlx::query_as( + "select count(1) from history + where user_id = $1", + ) + .bind(user.id) + .fetch_one(self.read_pool()) + .await + .map_err(fix_error)?; + + Ok(res.0) + } + + #[instrument(skip_all)] + async fn count_history_cached(&self, user: &User) -> DbResult<i64> { + let res: (i32,) = sqlx::query_as( + "select total from total_history_count_user + where user_id = $1", + ) + .bind(user.id) + .fetch_one(self.read_pool()) + .await + .map_err(fix_error)?; + + Ok(res.0 as i64) + } + + async fn delete_store(&self, user: &User) -> DbResult<()> { + let mut tx = self.pool.begin().await.map_err(fix_error)?; + + sqlx::query( + "delete from store + where user_id = $1", + ) + .bind(user.id) + .execute(&mut *tx) + .await + .map_err(fix_error)?; + + sqlx::query( + "delete from store_idx_cache + where user_id = $1", + ) + .bind(user.id) + .execute(&mut *tx) + .await + .map_err(fix_error)?; + + tx.commit().await.map_err(fix_error)?; + + Ok(()) + } + + async fn delete_history(&self, user: &User, id: String) -> DbResult<()> { + sqlx::query( + "update history + set deleted_at = $3 + where user_id = $1 + and client_id = $2 + and deleted_at is null", // don't just keep setting it + ) + .bind(user.id) + .bind(id) + .bind(OffsetDateTime::now_utc()) + .fetch_all(&self.pool) + .await + .map_err(fix_error)?; + + Ok(()) + } + + #[instrument(skip_all)] + async fn deleted_history(&self, user: &User) -> DbResult<Vec<String>> { + // The cache is new, and the user might not yet have a cache value. + // They will have one as soon as they post up some new history, but handle that + // edge case. + + let res = sqlx::query( + "select client_id from history + where user_id = $1 + and deleted_at is not null", + ) + .bind(user.id) + .fetch_all(self.read_pool()) + .await + .map_err(fix_error)?; + + let res = res + .iter() + .map(|row| row.get::<String, _>("client_id")) + .collect(); + + Ok(res) + } + + #[instrument(skip_all)] + async fn count_history_range( + &self, + user: &User, + range: Range<OffsetDateTime>, + ) -> DbResult<i64> { + let res: (i64,) = sqlx::query_as( + "select count(1) from history + where user_id = $1 + and timestamp >= $2::date + and timestamp < $3::date", + ) + .bind(user.id) + .bind(into_utc(range.start)) + .bind(into_utc(range.end)) + .fetch_one(self.read_pool()) + .await + .map_err(fix_error)?; + + Ok(res.0) + } + + #[instrument(skip_all)] + async fn list_history( + &self, + user: &User, + created_after: OffsetDateTime, + since: OffsetDateTime, + host: &str, + page_size: i64, + ) -> DbResult<Vec<History>> { + let res = sqlx::query_as( + "select id, client_id, user_id, hostname, timestamp, data, created_at from history + where user_id = $1 + and hostname != $2 + and created_at >= $3 + and timestamp >= $4 + order by timestamp asc + limit $5", + ) + .bind(user.id) + .bind(host) + .bind(into_utc(created_after)) + .bind(into_utc(since)) + .bind(page_size) + .fetch(self.read_pool()) + .map_ok(|DbHistory(h)| h) + .try_collect() + .await + .map_err(fix_error)?; + + Ok(res) + } + + #[instrument(skip_all)] + async fn add_history(&self, history: &[NewHistory]) -> DbResult<()> { + let mut tx = self.pool.begin().await.map_err(fix_error)?; + + for i in history { + let client_id: &str = &i.client_id; + let hostname: &str = &i.hostname; + let data: &str = &i.data; + + sqlx::query( + "insert into history + (client_id, user_id, hostname, timestamp, data) + values ($1, $2, $3, $4, $5) + on conflict do nothing + ", + ) + .bind(client_id) + .bind(i.user_id) + .bind(hostname) + .bind(i.timestamp) + .bind(data) + .execute(&mut *tx) + .await + .map_err(fix_error)?; + } + + tx.commit().await.map_err(fix_error)?; + + Ok(()) + } + + #[instrument(skip_all)] + async fn delete_user(&self, u: &User) -> DbResult<()> { + sqlx::query("delete from sessions where user_id = $1") + .bind(u.id) + .execute(&self.pool) + .await + .map_err(fix_error)?; + + sqlx::query("delete from history where user_id = $1") + .bind(u.id) + .execute(&self.pool) + .await + .map_err(fix_error)?; + + sqlx::query("delete from store where user_id = $1") + .bind(u.id) + .execute(&self.pool) + .await + .map_err(fix_error)?; + + sqlx::query("delete from total_history_count_user where user_id = $1") + .bind(u.id) + .execute(&self.pool) + .await + .map_err(fix_error)?; + + sqlx::query("delete from users where id = $1") + .bind(u.id) + .execute(&self.pool) + .await + .map_err(fix_error)?; + + Ok(()) + } + + #[instrument(skip_all)] + async fn update_user_password(&self, user: &User) -> DbResult<()> { + sqlx::query( + "update users + set password = $1 + where id = $2", + ) + .bind(&user.password) + .bind(user.id) + .execute(&self.pool) + .await + .map_err(fix_error)?; + + Ok(()) + } + + #[instrument(skip_all)] + async fn add_user(&self, user: &NewUser) -> DbResult<i64> { + let email: &str = &user.email; + let username: &str = &user.username; + let password: &str = &user.password; + + let res: (i64,) = sqlx::query_as( + "insert into users + (username, email, password) + values($1, $2, $3) + returning id", + ) + .bind(username) + .bind(email) + .bind(password) + .fetch_one(&self.pool) + .await + .map_err(fix_error)?; + + Ok(res.0) + } + + #[instrument(skip_all)] + async fn add_session(&self, session: &NewSession) -> DbResult<()> { + let token: &str = &session.token; + + sqlx::query( + "insert into sessions + (user_id, token) + values($1, $2)", + ) + .bind(session.user_id) + .bind(token) + .execute(&self.pool) + .await + .map_err(fix_error)?; + + Ok(()) + } + + #[instrument(skip_all)] + async fn get_user_session(&self, u: &User) -> DbResult<Session> { + sqlx::query_as("select id, user_id, token from sessions where user_id = $1") + .bind(u.id) + .fetch_one(self.read_pool()) + .await + .map_err(fix_error) + .map(|DbSession(session)| session) + } + + #[instrument(skip_all)] + async fn oldest_history(&self, user: &User) -> DbResult<History> { + sqlx::query_as( + "select id, client_id, user_id, hostname, timestamp, data, created_at from history + where user_id = $1 + order by timestamp asc + limit 1", + ) + .bind(user.id) + .fetch_one(self.read_pool()) + .await + .map_err(fix_error) + .map(|DbHistory(h)| h) + } + + #[instrument(skip_all)] + async fn add_records(&self, user: &User, records: &[Record<EncryptedData>]) -> DbResult<()> { + let mut tx = self.pool.begin().await.map_err(fix_error)?; + + // We won't have uploaded this data if it wasn't the max. Therefore, we can deduce the max + // idx without having to make further database queries. Doing the query on this small + // amount of data should be much, much faster. + // + // Worst case, say we get this wrong. We end up caching data that isn't actually the max + // idx, so clients upload again. The cache logic can be verified with a sql query anyway :) + + let mut heads = HashMap::<(HostId, &str), u64>::new(); + + for i in records { + let id = atuin_common::utils::uuid_v7(); + + let result = sqlx::query( + "insert into store + (id, client_id, host, idx, timestamp, version, tag, data, cek, user_id) + values ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10) + on conflict do nothing + ", + ) + .bind(id) + .bind(i.id) + .bind(i.host.id) + .bind(i.idx as i64) + .bind(i.timestamp as i64) // throwing away some data, but i64 is still big in terms of time + .bind(&i.version) + .bind(&i.tag) + .bind(&i.data.data) + .bind(&i.data.content_encryption_key) + .bind(user.id) + .execute(&mut *tx) + .await + .map_err(fix_error)?; + + // Only update heads if we actually inserted the record + if result.rows_affected() > 0 { + heads + .entry((i.host.id, &i.tag)) + .and_modify(|e| { + if i.idx > *e { + *e = i.idx + } + }) + .or_insert(i.idx); + } + } + + // we've built the map of heads for this push, so commit it to the database + for ((host, tag), idx) in heads { + sqlx::query( + "insert into store_idx_cache + (user_id, host, tag, idx) + values ($1, $2, $3, $4) + on conflict(user_id, host, tag) do update set idx = greatest(store_idx_cache.idx, $4) + ", + ) + .bind(user.id) + .bind(host) + .bind(tag) + .bind(idx as i64) + .execute(&mut *tx) + .await + .map_err(fix_error)?; + } + + tx.commit().await.map_err(fix_error)?; + + Ok(()) + } + + #[instrument(skip_all)] + async fn next_records( + &self, + user: &User, + host: HostId, + tag: String, + start: Option<RecordIdx>, + count: u64, + ) -> DbResult<Vec<Record<EncryptedData>>> { + tracing::debug!("{:?} - {:?} - {:?}", host, tag, start); + let start = start.unwrap_or(0); + + let records: Result<Vec<DbRecord>, DbError> = sqlx::query_as( + "select client_id, host, idx, timestamp, version, tag, data, cek from store + where user_id = $1 + and tag = $2 + and host = $3 + and idx >= $4 + order by idx asc + limit $5", + ) + .bind(user.id) + .bind(tag.clone()) + .bind(host) + .bind(start as i64) + .bind(count as i64) + .fetch_all(self.read_pool()) + .await + .map_err(fix_error); + + let ret = match records { + Ok(records) => { + let records: Vec<Record<EncryptedData>> = records + .into_iter() + .map(|f| { + let record: Record<EncryptedData> = f.into(); + record + }) + .collect(); + + records + } + Err(DbError::NotFound) => { + tracing::debug!("no records found in store: {:?}/{}", host, tag); + return Ok(vec![]); + } + Err(e) => return Err(e), + }; + + Ok(ret) + } + + async fn status(&self, user: &User) -> DbResult<RecordStatus> { + const STATUS_SQL: &str = + "select host, tag, max(idx) from store where user_id = $1 group by host, tag"; + + // If IDX_CACHE_ROLLOUT is set, then we + // 1. Read the value of the var, use it as a % chance of using the cache + // 2. If we use the cache, just read from the cache table + // 3. If we don't use the cache, read from the store table + // IDX_CACHE_ROLLOUT should be between 0 and 100. + + let idx_cache_rollout = std::env::var("IDX_CACHE_ROLLOUT").unwrap_or("0".to_string()); + let idx_cache_rollout = idx_cache_rollout.parse::<f64>().unwrap_or(0.0); + let use_idx_cache = rand::thread_rng().gen_bool(idx_cache_rollout / 100.0); + + let mut res: Vec<(Uuid, String, i64)> = if use_idx_cache { + tracing::debug!("using idx cache for user {}", user.id); + sqlx::query_as("select host, tag, idx from store_idx_cache where user_id = $1") + .bind(user.id) + .fetch_all(self.read_pool()) + .await + .map_err(fix_error)? + } else { + tracing::debug!("using aggregate query for user {}", user.id); + sqlx::query_as(STATUS_SQL) + .bind(user.id) + .fetch_all(self.read_pool()) + .await + .map_err(fix_error)? + }; + + res.sort(); + + let mut status = RecordStatus::new(); + + for i in res.iter() { + status.set_raw(HostId(i.0), i.1.clone(), i.2 as u64); + } + + Ok(status) + } +} + +fn into_utc(x: OffsetDateTime) -> PrimitiveDateTime { + let x = x.to_offset(UtcOffset::UTC); + PrimitiveDateTime::new(x.date(), x.time()) +} + +#[cfg(test)] +mod tests { + use time::macros::datetime; + + use crate::into_utc; + + #[test] + fn utc() { + let dt = datetime!(2023-09-26 15:11:02 +05:30); + assert_eq!(into_utc(dt), datetime!(2023-09-26 09:41:02)); + assert_eq!(into_utc(dt).assume_utc(), dt); + + let dt = datetime!(2023-09-26 15:11:02 -07:00); + assert_eq!(into_utc(dt), datetime!(2023-09-26 22:11:02)); + assert_eq!(into_utc(dt).assume_utc(), dt); + + let dt = datetime!(2023-09-26 15:11:02 +00:00); + assert_eq!(into_utc(dt), datetime!(2023-09-26 15:11:02)); + assert_eq!(into_utc(dt).assume_utc(), dt); + } +} diff --git a/crates/atuin-server-postgres/src/wrappers.rs b/crates/atuin-server-postgres/src/wrappers.rs new file mode 100644 index 00000000..cde4134c --- /dev/null +++ b/crates/atuin-server-postgres/src/wrappers.rs @@ -0,0 +1,77 @@ +use ::sqlx::{FromRow, Result}; +use atuin_common::record::{EncryptedData, Host, Record}; +use atuin_server_database::models::{History, Session, User}; +use sqlx::{Row, postgres::PgRow}; +use time::PrimitiveDateTime; + +pub struct DbUser(pub User); +pub struct DbSession(pub Session); +pub struct DbHistory(pub History); +pub struct DbRecord(pub Record<EncryptedData>); + +impl<'a> FromRow<'a, PgRow> for DbUser { + fn from_row(row: &'a PgRow) -> Result<Self> { + Ok(Self(User { + id: row.try_get("id")?, + username: row.try_get("username")?, + email: row.try_get("email")?, + password: row.try_get("password")?, + })) + } +} + +impl<'a> ::sqlx::FromRow<'a, PgRow> for DbSession { + fn from_row(row: &'a PgRow) -> ::sqlx::Result<Self> { + Ok(Self(Session { + id: row.try_get("id")?, + user_id: row.try_get("user_id")?, + token: row.try_get("token")?, + })) + } +} + +impl<'a> ::sqlx::FromRow<'a, PgRow> for DbHistory { + fn from_row(row: &'a PgRow) -> ::sqlx::Result<Self> { + Ok(Self(History { + id: row.try_get("id")?, + client_id: row.try_get("client_id")?, + user_id: row.try_get("user_id")?, + hostname: row.try_get("hostname")?, + timestamp: row + .try_get::<PrimitiveDateTime, _>("timestamp")? + .assume_utc(), + data: row.try_get("data")?, + created_at: row + .try_get::<PrimitiveDateTime, _>("created_at")? + .assume_utc(), + })) + } +} + +impl<'a> ::sqlx::FromRow<'a, PgRow> for DbRecord { + fn from_row(row: &'a PgRow) -> ::sqlx::Result<Self> { + let timestamp: i64 = row.try_get("timestamp")?; + let idx: i64 = row.try_get("idx")?; + + let data = EncryptedData { + data: row.try_get("data")?, + content_encryption_key: row.try_get("cek")?, + }; + + Ok(Self(Record { + id: row.try_get("client_id")?, + host: Host::new(row.try_get("host")?), + idx: idx as u64, + timestamp: timestamp as u64, + version: row.try_get("version")?, + tag: row.try_get("tag")?, + data, + })) + } +} + +impl From<DbRecord> for Record<EncryptedData> { + fn from(other: DbRecord) -> Record<EncryptedData> { + Record { ..other.0 } + } +} diff --git a/crates/atuin-server-sqlite/Cargo.toml b/crates/atuin-server-sqlite/Cargo.toml new file mode 100644 index 00000000..a0776746 --- /dev/null +++ b/crates/atuin-server-sqlite/Cargo.toml @@ -0,0 +1,24 @@ +[package] +name = "atuin-server-sqlite" +edition = "2024" +description = "server sqlite database library for atuin" + +version = { workspace = true } +authors = { workspace = true } +license = { workspace = true } +homepage = { workspace = true } +repository = { workspace = true } + +[dependencies] +atuin-common = { path = "../atuin-common", version = "18.13.2" } +atuin-server-database = { path = "../atuin-server-database", version = "18.13.2" } + +eyre = { workspace = true } +tracing = { workspace = true } +time = { workspace = true } +serde = { workspace = true } +sqlx = { workspace = true, features = ["sqlite", "regexp"] } +async-trait = { workspace = true } +uuid = { workspace = true } +metrics = "0.24" +futures-util = "0.3" diff --git a/crates/atuin-server-sqlite/build.rs b/crates/atuin-server-sqlite/build.rs new file mode 100644 index 00000000..d5068697 --- /dev/null +++ b/crates/atuin-server-sqlite/build.rs @@ -0,0 +1,5 @@ +// generated by `sqlx migrate build-script` +fn main() { + // trigger recompilation when a new migration is added + println!("cargo:rerun-if-changed=migrations"); +} diff --git a/crates/atuin-server-sqlite/migrations/20231203124112_create-store.sql b/crates/atuin-server-sqlite/migrations/20231203124112_create-store.sql new file mode 100644 index 00000000..ca19ed62 --- /dev/null +++ b/crates/atuin-server-sqlite/migrations/20231203124112_create-store.sql @@ -0,0 +1,17 @@ +create table store ( + id text primary key, -- remember to use uuidv7 for happy indices <3 + client_id text not null, -- I am too uncomfortable with the idea of a client-generated primary key, even though it's fine mathematically + host text not null, -- a unique identifier for the host + idx bigint not null, -- the index of the record in this store, identified by (host, tag) + timestamp bigint not null, -- not a timestamp type, as those do not have nanosecond precision + version text not null, + tag text not null, -- what is this? history, kv, whatever. Remember clients get a log per tag per host + data text not null, -- store the actual history data, encrypted. I don't wanna know! + cek text not null, + + user_id bigint not null, -- allow multiple users + created_at timestamp not null default current_timestamp +); + +create unique index record_uniq ON store(user_id, host, tag, idx); + diff --git a/crates/atuin-server-sqlite/migrations/20240108124830_create-history.sql b/crates/atuin-server-sqlite/migrations/20240108124830_create-history.sql new file mode 100644 index 00000000..7bd653ba --- /dev/null +++ b/crates/atuin-server-sqlite/migrations/20240108124830_create-history.sql @@ -0,0 +1,15 @@ +create table history ( + id integer primary key autoincrement, + client_id text not null unique, -- the client-generated ID + user_id bigserial not null, -- allow multiple users + hostname text not null, -- a unique identifier from the client (can be hashed, random, whatever) + timestamp timestamp not null, -- one of the few non-encrypted metadatas + + data text not null, -- store the actual history data, encrypted. I don't wanna know! + + created_at timestamp not null default current_timestamp, + deleted_at timestamp +); + +create unique index history_deleted_index on history(client_id, user_id, deleted_at); + diff --git a/crates/atuin-server-sqlite/migrations/20240108124831_create-sessions.sql b/crates/atuin-server-sqlite/migrations/20240108124831_create-sessions.sql new file mode 100644 index 00000000..3120c35d --- /dev/null +++ b/crates/atuin-server-sqlite/migrations/20240108124831_create-sessions.sql @@ -0,0 +1,6 @@ +create table sessions ( + id integer primary key autoincrement, + user_id integer, + token text unique not null +); + diff --git a/crates/atuin-server-sqlite/migrations/20240621110730_create-users.sql b/crates/atuin-server-sqlite/migrations/20240621110730_create-users.sql new file mode 100644 index 00000000..852c159d --- /dev/null +++ b/crates/atuin-server-sqlite/migrations/20240621110730_create-users.sql @@ -0,0 +1,12 @@ +create table users ( + id integer primary key autoincrement, -- also store our own ID + username text not null unique, -- being able to contact users is useful + email text not null unique, -- being able to contact users is useful + password text not null unique, + created_at timestamp not null default (datetime('now','localtime')), + verified_at timestamp with time zone default null +); + +-- the prior index is case sensitive :( +CREATE UNIQUE INDEX email_unique_idx on users (LOWER(email)); +CREATE UNIQUE INDEX username_unique_idx on users (LOWER(username)); diff --git a/crates/atuin-server-sqlite/migrations/20240621110731_create-user-verification-token.sql b/crates/atuin-server-sqlite/migrations/20240621110731_create-user-verification-token.sql new file mode 100644 index 00000000..36eb14de --- /dev/null +++ b/crates/atuin-server-sqlite/migrations/20240621110731_create-user-verification-token.sql @@ -0,0 +1,6 @@ +create table user_verification_token( + id integer primary key autoincrement, + user_id bigint unique references users(id), + token text, + valid_until timestamp with time zone +); diff --git a/crates/atuin-server-sqlite/migrations/20240702094825_create-store-idx-cache.sql b/crates/atuin-server-sqlite/migrations/20240702094825_create-store-idx-cache.sql new file mode 100644 index 00000000..cd54cb18 --- /dev/null +++ b/crates/atuin-server-sqlite/migrations/20240702094825_create-store-idx-cache.sql @@ -0,0 +1,10 @@ +create table store_idx_cache( + id integer primary key autoincrement, + user_id bigint, + + host uuid, + tag text, + idx bigint +); + +create unique index store_idx_cache_uniq on store_idx_cache(user_id, host, tag); diff --git a/crates/atuin-server-sqlite/migrations/20260127000000_remove-email-verification.sql b/crates/atuin-server-sqlite/migrations/20260127000000_remove-email-verification.sql new file mode 100644 index 00000000..0bde89d7 --- /dev/null +++ b/crates/atuin-server-sqlite/migrations/20260127000000_remove-email-verification.sql @@ -0,0 +1,2 @@ +drop table if exists user_verification_token; +alter table users drop column verified_at; diff --git a/crates/atuin-server-sqlite/src/lib.rs b/crates/atuin-server-sqlite/src/lib.rs new file mode 100644 index 00000000..d69258c4 --- /dev/null +++ b/crates/atuin-server-sqlite/src/lib.rs @@ -0,0 +1,462 @@ +use std::str::FromStr; + +use async_trait::async_trait; +use atuin_common::record::{EncryptedData, HostId, Record, RecordIdx, RecordStatus}; +use atuin_server_database::{ + Database, DbError, DbResult, DbSettings, + models::{History, NewHistory, NewSession, NewUser, Session, User}, +}; +use futures_util::TryStreamExt; +use sqlx::{ + Row, + sqlite::{SqliteConnectOptions, SqliteJournalMode, SqlitePoolOptions}, + types::Uuid, +}; +use time::{OffsetDateTime, PrimitiveDateTime, UtcOffset}; +use tracing::instrument; +use wrappers::{DbHistory, DbRecord, DbSession, DbUser}; + +mod wrappers; + +#[derive(Clone)] +pub struct Sqlite { + pool: sqlx::Pool<sqlx::sqlite::Sqlite>, +} + +fn fix_error(error: sqlx::Error) -> DbError { + match error { + sqlx::Error::RowNotFound => DbError::NotFound, + error => DbError::Other(error.into()), + } +} + +#[async_trait] +impl Database for Sqlite { + async fn new(settings: &DbSettings) -> DbResult<Self> { + let opts = SqliteConnectOptions::from_str(&settings.db_uri) + .map_err(fix_error)? + .journal_mode(SqliteJournalMode::Wal) + .create_if_missing(true); + + let pool = SqlitePoolOptions::new() + .connect_with(opts) + .await + .map_err(fix_error)?; + + sqlx::migrate!("./migrations") + .run(&pool) + .await + .map_err(|error| DbError::Other(error.into()))?; + + Ok(Self { pool }) + } + + #[instrument(skip_all)] + async fn get_session(&self, token: &str) -> DbResult<Session> { + sqlx::query_as("select id, user_id, token from sessions where token = $1") + .bind(token) + .fetch_one(&self.pool) + .await + .map_err(fix_error) + .map(|DbSession(session)| session) + } + + #[instrument(skip_all)] + async fn get_session_user(&self, token: &str) -> DbResult<User> { + sqlx::query_as( + "select users.id, users.username, users.email, users.password from users + inner join sessions + on users.id = sessions.user_id + and sessions.token = $1", + ) + .bind(token) + .fetch_one(&self.pool) + .await + .map_err(fix_error) + .map(|DbUser(user)| user) + } + + #[instrument(skip_all)] + async fn add_session(&self, session: &NewSession) -> DbResult<()> { + let token: &str = &session.token; + + sqlx::query( + "insert into sessions + (user_id, token) + values($1, $2)", + ) + .bind(session.user_id) + .bind(token) + .execute(&self.pool) + .await + .map_err(fix_error)?; + + Ok(()) + } + + #[instrument(skip_all)] + async fn get_user(&self, username: &str) -> DbResult<User> { + sqlx::query_as("select id, username, email, password from users where username = $1") + .bind(username) + .fetch_one(&self.pool) + .await + .map_err(fix_error) + .map(|DbUser(user)| user) + } + + #[instrument(skip_all)] + async fn get_user_session(&self, u: &User) -> DbResult<Session> { + sqlx::query_as("select id, user_id, token from sessions where user_id = $1") + .bind(u.id) + .fetch_one(&self.pool) + .await + .map_err(fix_error) + .map(|DbSession(session)| session) + } + + #[instrument(skip_all)] + async fn add_user(&self, user: &NewUser) -> DbResult<i64> { + let email: &str = &user.email; + let username: &str = &user.username; + let password: &str = &user.password; + + let res: (i64,) = sqlx::query_as( + "insert into users + (username, email, password) + values($1, $2, $3) + returning id", + ) + .bind(username) + .bind(email) + .bind(password) + .fetch_one(&self.pool) + .await + .map_err(fix_error)?; + + Ok(res.0) + } + + #[instrument(skip_all)] + async fn update_user_password(&self, user: &User) -> DbResult<()> { + sqlx::query( + "update users + set password = $1 + where id = $2", + ) + .bind(&user.password) + .bind(user.id) + .execute(&self.pool) + .await + .map_err(fix_error)?; + + Ok(()) + } + + #[instrument(skip_all)] + async fn count_history(&self, user: &User) -> DbResult<i64> { + // The cache is new, and the user might not yet have a cache value. + // They will have one as soon as they post up some new history, but handle that + // edge case. + + let res: (i64,) = sqlx::query_as( + "select count(1) from history + where user_id = $1", + ) + .bind(user.id) + .fetch_one(&self.pool) + .await + .map_err(fix_error)?; + + Ok(res.0) + } + + #[instrument(skip_all)] + async fn count_history_cached(&self, _user: &User) -> DbResult<i64> { + Err(DbError::NotFound) + } + + #[instrument(skip_all)] + async fn delete_user(&self, u: &User) -> DbResult<()> { + sqlx::query("delete from sessions where user_id = $1") + .bind(u.id) + .execute(&self.pool) + .await + .map_err(fix_error)?; + + sqlx::query("delete from users where id = $1") + .bind(u.id) + .execute(&self.pool) + .await + .map_err(fix_error)?; + + sqlx::query("delete from history where user_id = $1") + .bind(u.id) + .execute(&self.pool) + .await + .map_err(fix_error)?; + + Ok(()) + } + + async fn delete_history(&self, user: &User, id: String) -> DbResult<()> { + sqlx::query( + "update history + set deleted_at = $3 + where user_id = $1 + and client_id = $2 + and deleted_at is null", // don't just keep setting it + ) + .bind(user.id) + .bind(id) + .bind(time::OffsetDateTime::now_utc()) + .fetch_all(&self.pool) + .await + .map_err(fix_error)?; + + Ok(()) + } + + #[instrument(skip_all)] + async fn deleted_history(&self, user: &User) -> DbResult<Vec<String>> { + // The cache is new, and the user might not yet have a cache value. + // They will have one as soon as they post up some new history, but handle that + // edge case. + + let res = sqlx::query( + "select client_id from history + where user_id = $1 + and deleted_at is not null", + ) + .bind(user.id) + .fetch_all(&self.pool) + .await + .map_err(fix_error)?; + + let res = res.iter().map(|row| row.get("client_id")).collect(); + + Ok(res) + } + + async fn delete_store(&self, user: &User) -> DbResult<()> { + sqlx::query( + "delete from store + where user_id = $1", + ) + .bind(user.id) + .execute(&self.pool) + .await + .map_err(fix_error)?; + + Ok(()) + } + + #[instrument(skip_all)] + async fn add_records(&self, user: &User, records: &[Record<EncryptedData>]) -> DbResult<()> { + let mut tx = self.pool.begin().await.map_err(fix_error)?; + + for i in records { + let id = atuin_common::utils::uuid_v7(); + + sqlx::query( + "insert into store + (id, client_id, host, idx, timestamp, version, tag, data, cek, user_id) + values ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10) + on conflict do nothing + ", + ) + .bind(id) + .bind(i.id) + .bind(i.host.id) + .bind(i.idx as i64) + .bind(i.timestamp as i64) // throwing away some data, but i64 is still big in terms of time + .bind(&i.version) + .bind(&i.tag) + .bind(&i.data.data) + .bind(&i.data.content_encryption_key) + .bind(user.id) + .execute(&mut *tx) + .await + .map_err(fix_error)?; + } + + tx.commit().await.map_err(fix_error)?; + + Ok(()) + } + + #[instrument(skip_all)] + async fn next_records( + &self, + user: &User, + host: HostId, + tag: String, + start: Option<RecordIdx>, + count: u64, + ) -> DbResult<Vec<Record<EncryptedData>>> { + tracing::debug!("{:?} - {:?} - {:?}", host, tag, start); + let start = start.unwrap_or(0); + + let records: Result<Vec<DbRecord>, DbError> = sqlx::query_as( + "select client_id, host, idx, timestamp, version, tag, data, cek from store + where user_id = $1 + and tag = $2 + and host = $3 + and idx >= $4 + order by idx asc + limit $5", + ) + .bind(user.id) + .bind(tag.clone()) + .bind(host) + .bind(start as i64) + .bind(count as i64) + .fetch_all(&self.pool) + .await + .map_err(fix_error); + + let ret = match records { + Ok(records) => { + let records: Vec<Record<EncryptedData>> = records + .into_iter() + .map(|f| { + let record: Record<EncryptedData> = f.into(); + record + }) + .collect(); + + records + } + Err(DbError::NotFound) => { + tracing::debug!("no records found in store: {:?}/{}", host, tag); + return Ok(vec![]); + } + Err(e) => return Err(e), + }; + + Ok(ret) + } + + async fn status(&self, user: &User) -> DbResult<RecordStatus> { + const STATUS_SQL: &str = + "select host, tag, max(idx) from store where user_id = $1 group by host, tag"; + + let res: Vec<(Uuid, String, i64)> = sqlx::query_as(STATUS_SQL) + .bind(user.id) + .fetch_all(&self.pool) + .await + .map_err(fix_error)?; + + let mut status = RecordStatus::new(); + + for i in res { + status.set_raw(HostId(i.0), i.1, i.2 as u64); + } + + Ok(status) + } + + #[instrument(skip_all)] + async fn count_history_range( + &self, + user: &User, + range: std::ops::Range<time::OffsetDateTime>, + ) -> DbResult<i64> { + let res: (i64,) = sqlx::query_as( + "select count(1) from history + where user_id = $1 + and timestamp >= $2::date + and timestamp < $3::date", + ) + .bind(user.id) + .bind(into_utc(range.start)) + .bind(into_utc(range.end)) + .fetch_one(&self.pool) + .await + .map_err(fix_error)?; + + Ok(res.0) + } + + #[instrument(skip_all)] + async fn list_history( + &self, + user: &User, + created_after: time::OffsetDateTime, + since: time::OffsetDateTime, + host: &str, + page_size: i64, + ) -> DbResult<Vec<History>> { + let res = sqlx::query_as( + "select id, client_id, user_id, hostname, timestamp, data, created_at from history + where user_id = $1 + and hostname != $2 + and created_at >= $3 + and timestamp >= $4 + order by timestamp asc + limit $5", + ) + .bind(user.id) + .bind(host) + .bind(into_utc(created_after)) + .bind(into_utc(since)) + .bind(page_size) + .fetch(&self.pool) + .map_ok(|DbHistory(h)| h) + .try_collect() + .await + .map_err(fix_error)?; + + Ok(res) + } + + #[instrument(skip_all)] + async fn add_history(&self, history: &[NewHistory]) -> DbResult<()> { + let mut tx = self.pool.begin().await.map_err(fix_error)?; + + for i in history { + let client_id: &str = &i.client_id; + let hostname: &str = &i.hostname; + let data: &str = &i.data; + + sqlx::query( + "insert into history + (client_id, user_id, hostname, timestamp, data) + values ($1, $2, $3, $4, $5) + on conflict do nothing + ", + ) + .bind(client_id) + .bind(i.user_id) + .bind(hostname) + .bind(i.timestamp) + .bind(data) + .execute(&mut *tx) + .await + .map_err(fix_error)?; + } + + tx.commit().await.map_err(fix_error)?; + + Ok(()) + } + + #[instrument(skip_all)] + async fn oldest_history(&self, user: &User) -> DbResult<History> { + sqlx::query_as( + "select id, client_id, user_id, hostname, timestamp, data, created_at from history + where user_id = $1 + order by timestamp asc + limit 1", + ) + .bind(user.id) + .fetch_one(&self.pool) + .await + .map_err(fix_error) + .map(|DbHistory(h)| h) + } +} + +fn into_utc(x: OffsetDateTime) -> PrimitiveDateTime { + let x = x.to_offset(UtcOffset::UTC); + PrimitiveDateTime::new(x.date(), x.time()) +} diff --git a/crates/atuin-server-sqlite/src/wrappers.rs b/crates/atuin-server-sqlite/src/wrappers.rs new file mode 100644 index 00000000..2f1230c2 --- /dev/null +++ b/crates/atuin-server-sqlite/src/wrappers.rs @@ -0,0 +1,72 @@ +use ::sqlx::{FromRow, Result}; +use atuin_common::record::{EncryptedData, Host, Record}; +use atuin_server_database::models::{History, Session, User}; +use sqlx::{Row, sqlite::SqliteRow}; + +pub struct DbUser(pub User); +pub struct DbSession(pub Session); +pub struct DbHistory(pub History); +pub struct DbRecord(pub Record<EncryptedData>); + +impl<'a> FromRow<'a, SqliteRow> for DbUser { + fn from_row(row: &'a SqliteRow) -> Result<Self> { + Ok(Self(User { + id: row.try_get("id")?, + username: row.try_get("username")?, + email: row.try_get("email")?, + password: row.try_get("password")?, + })) + } +} + +impl<'a> ::sqlx::FromRow<'a, SqliteRow> for DbSession { + fn from_row(row: &'a SqliteRow) -> ::sqlx::Result<Self> { + Ok(Self(Session { + id: row.try_get("id")?, + user_id: row.try_get("user_id")?, + token: row.try_get("token")?, + })) + } +} + +impl<'a> ::sqlx::FromRow<'a, SqliteRow> for DbHistory { + fn from_row(row: &'a SqliteRow) -> ::sqlx::Result<Self> { + Ok(Self(History { + id: row.try_get("id")?, + client_id: row.try_get("client_id")?, + user_id: row.try_get("user_id")?, + hostname: row.try_get("hostname")?, + timestamp: row.try_get("timestamp")?, + data: row.try_get("data")?, + created_at: row.try_get("created_at")?, + })) + } +} + +impl<'a> ::sqlx::FromRow<'a, SqliteRow> for DbRecord { + fn from_row(row: &'a SqliteRow) -> ::sqlx::Result<Self> { + let idx: i64 = row.try_get("idx")?; + let timestamp: i64 = row.try_get("timestamp")?; + + let data = EncryptedData { + data: row.try_get("data")?, + content_encryption_key: row.try_get("cek")?, + }; + + Ok(Self(Record { + id: row.try_get("client_id")?, + host: Host::new(row.try_get("host")?), + idx: idx as u64, + timestamp: timestamp as u64, + version: row.try_get("version")?, + tag: row.try_get("tag")?, + data, + })) + } +} + +impl From<DbRecord> for Record<EncryptedData> { + fn from(other: DbRecord) -> Record<EncryptedData> { + Record { ..other.0 } + } +} diff --git a/crates/atuin-server/Cargo.toml b/crates/atuin-server/Cargo.toml new file mode 100644 index 00000000..b7779899 --- /dev/null +++ b/crates/atuin-server/Cargo.toml @@ -0,0 +1,45 @@ +[package] +name = "atuin-server" +edition = "2024" +description = "server library for atuin" + +rust-version = { workspace = true } +version = { workspace = true } +authors = { workspace = true } +license = { workspace = true } +homepage = { workspace = true } +repository = { workspace = true } + +[lib] +name = "atuin_server" +path = "src/lib.rs" + +[[bin]] +name = "atuin-server" +path = "src/bin/main.rs" + +[dependencies] +atuin-common = { workspace = true } +atuin-server-database = { workspace = true } +atuin-server-postgres = { workspace = true } +atuin-server-sqlite = { workspace = true } + +tracing = { workspace = true } +time = { workspace = true } +eyre = { workspace = true } +config = { workspace = true } +serde = { workspace = true } +serde_json = { workspace = true } +rand = { workspace = true } +tokio = { workspace = true } +axum = "0.8" +fs-err = { workspace = true } +tower = { workspace = true } +tower-http = { version = "0.6", features = ["trace"] } +reqwest = { workspace = true } +argon2 = "0.5" +semver = { workspace = true } +metrics-exporter-prometheus = { version = "0.18", default-features = false } +metrics = "0.24" +clap = { workspace = true } +tracing-subscriber = { workspace = true } diff --git a/crates/atuin-server/server.toml b/crates/atuin-server/server.toml new file mode 100644 index 00000000..9ff95890 --- /dev/null +++ b/crates/atuin-server/server.toml @@ -0,0 +1,38 @@ +## host to bind, can also be passed via CLI args +# host = "127.0.0.1" + +## port to bind, can also be passed via CLI args +# port = 8888 + +## whether to allow anyone to register an account +# open_registration = false + +## URI for postgres (using development creds here) +# db_uri="postgres://username:password@localhost/atuin" +# db_uri="sqlite:///config/atuin-server.db" + +## Optional: URI for read replica database +## If set, read-only queries will be routed to this database +# read_db_uri="postgres://username:password@localhost-replica/atuin" + +## Maximum size for one history entry +# max_history_length = 8192 + +## Maximum size for one record entry +## 1024 * 1024 * 1024 +# max_record_size = 1073741824 + +## Webhook to be called when user registers on the servers +# register_webhook_username = "" + +## Default page size for requests +# page_size = 1100 + +# [metrics] +# enable = false +# host = 127.0.0.1 +# port = 9001 + +## Enable legacy sync v1 routes (history-based sync) +## Set to false to disable and use only the newer record-based sync +# sync_v1_enabled = true diff --git a/crates/atuin-server/src/bin/main.rs b/crates/atuin-server/src/bin/main.rs new file mode 100644 index 00000000..960bdf6e --- /dev/null +++ b/crates/atuin-server/src/bin/main.rs @@ -0,0 +1,73 @@ +#![forbid(unsafe_code)] + +use std::net::SocketAddr; + +use atuin_server::{Settings, example_config, launch, launch_metrics_server}; +use atuin_server_database::DbType; +use atuin_server_postgres::Postgres; +use atuin_server_sqlite::Sqlite; + +use clap::Parser; +use eyre::{Context, Result, eyre}; +use tracing_subscriber::{EnvFilter, fmt, prelude::*}; + +#[derive(Parser, Debug)] +#[clap( + name = "atuin-server", + about = "Atuin sync server", + version, + infer_subcommands = true +)] +enum Cmd { + /// Start the server + Start { + /// The host address to bind + #[clap(long)] + host: Option<String>, + + /// The port to bind + #[clap(long, short)] + port: Option<u16>, + }, + + /// Print server example configuration + DefaultConfig, +} + +#[tokio::main] +async fn main() -> Result<()> { + let cmd = Cmd::parse(); + + tracing_subscriber::registry() + .with(fmt::layer()) + .with(EnvFilter::from_default_env()) + .init(); + + tracing::trace!(command = ?cmd, "server command"); + + match cmd { + Cmd::Start { host, port } => { + let settings = Settings::new().wrap_err("could not load server settings")?; + let host = host.as_ref().unwrap_or(&settings.host).clone(); + let port = port.unwrap_or(settings.port); + let addr = SocketAddr::new(host.parse()?, port); + + if settings.metrics.enable { + tokio::spawn(launch_metrics_server( + settings.metrics.host.clone(), + settings.metrics.port, + )); + } + + match settings.db_settings.db_type() { + DbType::Postgres => launch::<Postgres>(settings, addr).await, + DbType::Sqlite => launch::<Sqlite>(settings, addr).await, + DbType::Unknown => Err(eyre!("db_uri must start with postgres:// or sqlite://")), + } + } + Cmd::DefaultConfig => { + println!("{}", example_config()); + Ok(()) + } + } +} diff --git a/crates/atuin-server/src/handlers/health.rs b/crates/atuin-server/src/handlers/health.rs new file mode 100644 index 00000000..aebd1e8f --- /dev/null +++ b/crates/atuin-server/src/handlers/health.rs @@ -0,0 +1,15 @@ +use axum::{Json, http, response::IntoResponse}; + +use serde::Serialize; + +#[derive(Serialize)] +pub struct HealthResponse { + pub status: &'static str, +} + +pub async fn health_check() -> impl IntoResponse { + ( + http::StatusCode::OK, + Json(HealthResponse { status: "healthy" }), + ) +} diff --git a/crates/atuin-server/src/handlers/history.rs b/crates/atuin-server/src/handlers/history.rs new file mode 100644 index 00000000..bdafcc60 --- /dev/null +++ b/crates/atuin-server/src/handlers/history.rs @@ -0,0 +1,237 @@ +use std::{collections::HashMap, convert::TryFrom}; + +use axum::{ + Json, + extract::{Path, Query, State}, + http::{HeaderMap, StatusCode}, +}; +use metrics::counter; +use time::{Month, UtcOffset}; +use tracing::{debug, error, instrument}; + +use super::{ErrorResponse, ErrorResponseStatus, RespExt}; +use crate::{ + router::{AppState, UserAuth}, + utils::client_version_min, +}; +use atuin_server_database::{ + Database, + calendar::{TimePeriod, TimePeriodInfo}, + models::NewHistory, +}; + +use atuin_common::api::*; + +#[instrument(skip_all, fields(user.id = user.id))] +pub async fn count<DB: Database>( + UserAuth(user): UserAuth, + state: State<AppState<DB>>, +) -> Result<Json<CountResponse>, ErrorResponseStatus<'static>> { + let db = &state.0.database; + match db.count_history_cached(&user).await { + // By default read out the cached value + Ok(count) => Ok(Json(CountResponse { count })), + + // If that fails, fallback on a full COUNT. Cache is built on a POST + // only + Err(_) => match db.count_history(&user).await { + Ok(count) => Ok(Json(CountResponse { count })), + Err(_) => Err(ErrorResponse::reply("failed to query history count") + .with_status(StatusCode::INTERNAL_SERVER_ERROR)), + }, + } +} + +#[instrument(skip_all, fields(user.id = user.id))] +pub async fn list<DB: Database>( + req: Query<SyncHistoryRequest>, + UserAuth(user): UserAuth, + headers: HeaderMap, + state: State<AppState<DB>>, +) -> Result<Json<SyncHistoryResponse>, ErrorResponseStatus<'static>> { + let db = &state.0.database; + + let agent = headers + .get("user-agent") + .map_or("", |v| v.to_str().unwrap_or("")); + + let variable_page_size = client_version_min(agent, ">=15.0.0").unwrap_or(false); + + let page_size = if variable_page_size { + state.settings.page_size + } else { + 100 + }; + + if req.sync_ts.unix_timestamp_nanos() < 0 || req.history_ts.unix_timestamp_nanos() < 0 { + error!("client asked for history from < epoch 0"); + counter!("atuin_history_epoch_before_zero").increment(1); + + return Err( + ErrorResponse::reply("asked for history from before epoch 0") + .with_status(StatusCode::BAD_REQUEST), + ); + } + + let history = db + .list_history(&user, req.sync_ts, req.history_ts, &req.host, page_size) + .await; + + if let Err(e) = history { + error!("failed to load history: {}", e); + return Err(ErrorResponse::reply("failed to load history") + .with_status(StatusCode::INTERNAL_SERVER_ERROR)); + } + + let history: Vec<String> = history + .unwrap() + .iter() + .map(|i| i.data.to_string()) + .collect(); + + debug!( + "loaded {} items of history for user {}", + history.len(), + user.id + ); + + counter!("atuin_history_returned").increment(history.len() as u64); + + Ok(Json(SyncHistoryResponse { history })) +} + +#[instrument(skip_all, fields(user.id = user.id))] +pub async fn delete<DB: Database>( + UserAuth(user): UserAuth, + state: State<AppState<DB>>, + Json(req): Json<DeleteHistoryRequest>, +) -> Result<Json<MessageResponse>, ErrorResponseStatus<'static>> { + let db = &state.0.database; + + // user_id is the ID of the history, as set by the user (the server has its own ID) + let deleted = db.delete_history(&user, req.client_id).await; + + if let Err(e) = deleted { + error!("failed to delete history: {}", e); + return Err(ErrorResponse::reply("failed to delete history") + .with_status(StatusCode::INTERNAL_SERVER_ERROR)); + } + + Ok(Json(MessageResponse { + message: String::from("deleted OK"), + })) +} + +#[instrument(skip_all, fields(user.id = user.id))] +pub async fn add<DB: Database>( + UserAuth(user): UserAuth, + state: State<AppState<DB>>, + Json(req): Json<Vec<AddHistoryRequest>>, +) -> Result<(), ErrorResponseStatus<'static>> { + let State(AppState { database, settings }) = state; + + debug!("request to add {} history items", req.len()); + counter!("atuin_history_uploaded").increment(req.len() as u64); + + let mut history: Vec<NewHistory> = req + .into_iter() + .map(|h| NewHistory { + client_id: h.id, + user_id: user.id, + hostname: h.hostname, + timestamp: h.timestamp, + data: h.data, + }) + .collect(); + + history.retain(|h| { + // keep if within limit, or limit is 0 (unlimited) + let keep = h.data.len() <= settings.max_history_length || settings.max_history_length == 0; + + // Don't return an error here. We want to insert as much of the + // history list as we can, so log the error and continue going. + if !keep { + counter!("atuin_history_too_long").increment(1); + + tracing::warn!( + "history too long, got length {}, max {}", + h.data.len(), + settings.max_history_length + ); + } + + keep + }); + + if let Err(e) = database.add_history(&history).await { + error!("failed to add history: {}", e); + + return Err(ErrorResponse::reply("failed to add history") + .with_status(StatusCode::INTERNAL_SERVER_ERROR)); + }; + + Ok(()) +} + +#[derive(serde::Deserialize, Debug)] +pub struct CalendarQuery { + #[serde(default = "serde_calendar::zero")] + year: i32, + #[serde(default = "serde_calendar::one")] + month: u8, + + #[serde(default = "serde_calendar::utc")] + tz: UtcOffset, +} + +mod serde_calendar { + use time::UtcOffset; + + pub fn zero() -> i32 { + 0 + } + + pub fn one() -> u8 { + 1 + } + + pub fn utc() -> UtcOffset { + UtcOffset::UTC + } +} + +#[instrument(skip_all, fields(user.id = user.id))] +pub async fn calendar<DB: Database>( + Path(focus): Path<String>, + Query(params): Query<CalendarQuery>, + UserAuth(user): UserAuth, + state: State<AppState<DB>>, +) -> Result<Json<HashMap<u64, TimePeriodInfo>>, ErrorResponseStatus<'static>> { + let focus = focus.as_str(); + + let year = params.year; + let month = Month::try_from(params.month).map_err(|e| ErrorResponseStatus { + error: ErrorResponse { + reason: e.to_string().into(), + }, + status: StatusCode::BAD_REQUEST, + })?; + + let period = match focus { + "year" => TimePeriod::Year, + "month" => TimePeriod::Month { year }, + "day" => TimePeriod::Day { year, month }, + _ => { + return Err(ErrorResponse::reply("invalid focus: use year/month/day") + .with_status(StatusCode::BAD_REQUEST)); + } + }; + + let db = &state.0.database; + let focus = db.calendar(&user, period, params.tz).await.map_err(|_| { + ErrorResponse::reply("failed to query calendar") + .with_status(StatusCode::INTERNAL_SERVER_ERROR) + })?; + + Ok(Json(focus)) +} diff --git a/crates/atuin-server/src/handlers/mod.rs b/crates/atuin-server/src/handlers/mod.rs new file mode 100644 index 00000000..2176ac5e --- /dev/null +++ b/crates/atuin-server/src/handlers/mod.rs @@ -0,0 +1,60 @@ +use atuin_common::api::{ErrorResponse, IndexResponse}; +use atuin_server_database::Database; +use axum::{Json, extract::State, http, response::IntoResponse}; + +use crate::router::AppState; + +pub mod health; +pub mod history; +pub mod record; +pub mod status; +pub mod user; +pub mod v0; + +const VERSION: &str = env!("CARGO_PKG_VERSION"); + +pub async fn index<DB: Database>(state: State<AppState<DB>>) -> Json<IndexResponse> { + let homage = r#""Through the fathomless deeps of space swims the star turtle Great A'Tuin, bearing on its back the four giant elephants who carry on their shoulders the mass of the Discworld." -- Sir Terry Pratchett"#; + + let version = state + .settings + .fake_version + .clone() + .unwrap_or(VERSION.to_string()); + + Json(IndexResponse { + homage: homage.to_string(), + version, + }) +} + +impl IntoResponse for ErrorResponseStatus<'_> { + fn into_response(self) -> axum::response::Response { + (self.status, Json(self.error)).into_response() + } +} + +pub struct ErrorResponseStatus<'a> { + pub error: ErrorResponse<'a>, + pub status: http::StatusCode, +} + +pub trait RespExt<'a> { + fn with_status(self, status: http::StatusCode) -> ErrorResponseStatus<'a>; + fn reply(reason: &'a str) -> Self; +} + +impl<'a> RespExt<'a> for ErrorResponse<'a> { + fn with_status(self, status: http::StatusCode) -> ErrorResponseStatus<'a> { + ErrorResponseStatus { + error: self, + status, + } + } + + fn reply(reason: &'a str) -> ErrorResponse<'a> { + Self { + reason: reason.into(), + } + } +} diff --git a/crates/atuin-server/src/handlers/record.rs b/crates/atuin-server/src/handlers/record.rs new file mode 100644 index 00000000..410c54bd --- /dev/null +++ b/crates/atuin-server/src/handlers/record.rs @@ -0,0 +1,42 @@ +use axum::{Json, http::StatusCode, response::IntoResponse}; +use serde_json::json; +use tracing::instrument; + +use super::{ErrorResponse, ErrorResponseStatus, RespExt}; +use crate::router::UserAuth; + +use atuin_common::record::{EncryptedData, Record}; + +#[instrument(skip_all, fields(user.id = user.id))] +pub async fn post(UserAuth(user): UserAuth) -> Result<(), ErrorResponseStatus<'static>> { + // anyone who has actually used the old record store (a very small number) will see this error + // upon trying to sync. + // 1. The status endpoint will say that the server has nothing + // 2. The client will try to upload local records + // 3. Sync will fail with this error + + // If the client has no local records, they will see the empty index and do nothing. For the + // vast majority of users, this is the case. + return Err( + ErrorResponse::reply("record store deprecated; please upgrade") + .with_status(StatusCode::BAD_REQUEST), + ); +} + +#[instrument(skip_all, fields(user.id = user.id))] +pub async fn index(UserAuth(user): UserAuth) -> axum::response::Response { + let ret = json!({ + "hosts": {} + }); + + ret.to_string().into_response() +} + +#[instrument(skip_all, fields(user.id = user.id))] +pub async fn next( + UserAuth(user): UserAuth, +) -> Result<Json<Vec<Record<EncryptedData>>>, ErrorResponseStatus<'static>> { + let records = Vec::new(); + + Ok(Json(records)) +} diff --git a/crates/atuin-server/src/handlers/status.rs b/crates/atuin-server/src/handlers/status.rs new file mode 100644 index 00000000..9c152d51 --- /dev/null +++ b/crates/atuin-server/src/handlers/status.rs @@ -0,0 +1,45 @@ +use axum::{Json, extract::State, http::StatusCode}; +use tracing::instrument; + +use super::{ErrorResponse, ErrorResponseStatus, RespExt}; +use crate::router::{AppState, UserAuth}; +use atuin_server_database::Database; + +use atuin_common::api::*; + +const VERSION: &str = env!("CARGO_PKG_VERSION"); + +#[instrument(skip_all, fields(user.id = user.id))] +pub async fn status<DB: Database>( + UserAuth(user): UserAuth, + state: State<AppState<DB>>, +) -> Result<Json<StatusResponse>, ErrorResponseStatus<'static>> { + let db = &state.0.database; + + let deleted = db.deleted_history(&user).await.unwrap_or(vec![]); + + let count = match db.count_history_cached(&user).await { + // By default read out the cached value + Ok(count) => count, + + // If that fails, fallback on a full COUNT. Cache is built on a POST + // only + Err(_) => match db.count_history(&user).await { + Ok(count) => count, + Err(_) => { + return Err(ErrorResponse::reply("failed to query history count") + .with_status(StatusCode::INTERNAL_SERVER_ERROR)); + } + }, + }; + + tracing::debug!(user = user.username, "requested sync status"); + + Ok(Json(StatusResponse { + count, + deleted, + username: user.username, + version: VERSION.to_string(), + page_size: state.settings.page_size, + })) +} diff --git a/crates/atuin-server/src/handlers/user.rs b/crates/atuin-server/src/handlers/user.rs new file mode 100644 index 00000000..6436e327 --- /dev/null +++ b/crates/atuin-server/src/handlers/user.rs @@ -0,0 +1,265 @@ +use std::borrow::Borrow; +use std::collections::HashMap; +use std::time::Duration; + +use argon2::{ + Algorithm, Argon2, Params, PasswordHash, PasswordHasher, PasswordVerifier, Version, + password_hash::SaltString, +}; +use axum::{ + Json, + extract::{Path, State}, + http::StatusCode, +}; +use metrics::counter; + +use rand::rngs::OsRng; +use tracing::{debug, error, info, instrument}; + +use atuin_common::tls::ensure_crypto_provider; + +use super::{ErrorResponse, ErrorResponseStatus, RespExt}; +use crate::router::{AppState, UserAuth}; +use atuin_server_database::{ + Database, DbError, + models::{NewSession, NewUser}, +}; + +use reqwest::header::CONTENT_TYPE; + +use atuin_common::{api::*, utils::crypto_random_string}; + +pub fn verify_str(hash: &str, password: &str) -> bool { + let arg2 = Argon2::new(Algorithm::Argon2id, Version::V0x13, Params::default()); + let Ok(hash) = PasswordHash::new(hash) else { + return false; + }; + arg2.verify_password(password.as_bytes(), &hash).is_ok() +} + +// Try to send a Discord webhook once - if it fails, we don't retry. "At most once", and best effort. +// Don't return the status because if this fails, we don't really care. +async fn send_register_hook(url: &str, username: String, registered: String) { + ensure_crypto_provider(); + let hook = HashMap::from([ + ("username", username), + ("content", format!("{registered} has just signed up!")), + ]); + + let client = reqwest::Client::new(); + + let resp = client + .post(url) + .timeout(Duration::new(5, 0)) + .header(CONTENT_TYPE, "application/json") + .json(&hook) + .send() + .await; + + match resp { + Ok(_) => info!("register webhook sent ok!"), + Err(e) => error!("failed to send register webhook: {}", e), + } +} + +#[instrument(skip_all, fields(user.username = username.as_str()))] +pub async fn get<DB: Database>( + Path(username): Path<String>, + state: State<AppState<DB>>, +) -> Result<Json<UserResponse>, ErrorResponseStatus<'static>> { + let db = &state.0.database; + let user = match db.get_user(username.as_ref()).await { + Ok(user) => user, + Err(DbError::NotFound) => { + debug!("user not found: {}", username); + return Err(ErrorResponse::reply("user not found").with_status(StatusCode::NOT_FOUND)); + } + Err(DbError::Other(err)) => { + error!("database error: {}", err); + return Err(ErrorResponse::reply("database error") + .with_status(StatusCode::INTERNAL_SERVER_ERROR)); + } + }; + + Ok(Json(UserResponse { + username: user.username, + })) +} + +#[instrument(skip_all)] +pub async fn register<DB: Database>( + state: State<AppState<DB>>, + Json(register): Json<RegisterRequest>, +) -> Result<Json<RegisterResponse>, ErrorResponseStatus<'static>> { + if !state.settings.open_registration { + return Err( + ErrorResponse::reply("this server is not open for registrations") + .with_status(StatusCode::BAD_REQUEST), + ); + } + + for c in register.username.chars() { + match c { + 'a'..='z' | 'A'..='Z' | '0'..='9' | '-' => {} + _ => { + return Err(ErrorResponse::reply( + "Only alphanumeric and hyphens (-) are allowed in usernames", + ) + .with_status(StatusCode::BAD_REQUEST)); + } + } + } + + let hashed = hash_secret(®ister.password); + + let new_user = NewUser { + email: register.email.clone(), + username: register.username.clone(), + password: hashed, + }; + + let db = &state.0.database; + let user_id = match db.add_user(&new_user).await { + Ok(id) => id, + Err(e) => { + error!("failed to add user: {}", e); + return Err( + ErrorResponse::reply("failed to add user").with_status(StatusCode::BAD_REQUEST) + ); + } + }; + + // 24 bytes encoded as base64 + let token = crypto_random_string::<24>(); + + let new_session = NewSession { + user_id, + token: (&token).into(), + }; + + if let Some(url) = &state.settings.register_webhook_url { + // Could probs be run on another thread, but it's ok atm + send_register_hook( + url, + state.settings.register_webhook_username.clone(), + register.username, + ) + .await; + } + + counter!("atuin_users_registered").increment(1); + + match db.add_session(&new_session).await { + Ok(_) => Ok(Json(RegisterResponse { session: token })), + Err(e) => { + error!("failed to add session: {}", e); + Err(ErrorResponse::reply("failed to register user") + .with_status(StatusCode::BAD_REQUEST)) + } + } +} + +#[instrument(skip_all, fields(user.id = user.id))] +pub async fn delete<DB: Database>( + UserAuth(user): UserAuth, + state: State<AppState<DB>>, +) -> Result<Json<DeleteUserResponse>, ErrorResponseStatus<'static>> { + debug!("request to delete user {}", user.id); + + let db = &state.0.database; + if let Err(e) = db.delete_user(&user).await { + error!("failed to delete user: {}", e); + + return Err(ErrorResponse::reply("failed to delete user") + .with_status(StatusCode::INTERNAL_SERVER_ERROR)); + }; + + counter!("atuin_users_deleted").increment(1); + + Ok(Json(DeleteUserResponse {})) +} + +#[instrument(skip_all, fields(user.id = user.id, change_password))] +pub async fn change_password<DB: Database>( + UserAuth(mut user): UserAuth, + state: State<AppState<DB>>, + Json(change_password): Json<ChangePasswordRequest>, +) -> Result<Json<ChangePasswordResponse>, ErrorResponseStatus<'static>> { + let db = &state.0.database; + + let verified = verify_str( + user.password.as_str(), + change_password.current_password.borrow(), + ); + if !verified { + return Err( + ErrorResponse::reply("password is not correct").with_status(StatusCode::UNAUTHORIZED) + ); + } + + let hashed = hash_secret(&change_password.new_password); + user.password = hashed; + + if let Err(e) = db.update_user_password(&user).await { + error!("failed to change user password: {}", e); + + return Err(ErrorResponse::reply("failed to change user password") + .with_status(StatusCode::INTERNAL_SERVER_ERROR)); + }; + Ok(Json(ChangePasswordResponse {})) +} + +#[instrument(skip_all, fields(user.username = login.username.as_str()))] +pub async fn login<DB: Database>( + state: State<AppState<DB>>, + login: Json<LoginRequest>, +) -> Result<Json<LoginResponse>, ErrorResponseStatus<'static>> { + let db = &state.0.database; + let user = match db.get_user(login.username.borrow()).await { + Ok(u) => u, + Err(DbError::NotFound) => { + return Err(ErrorResponse::reply("user not found").with_status(StatusCode::NOT_FOUND)); + } + Err(DbError::Other(e)) => { + error!("failed to get user {}: {}", login.username.clone(), e); + + return Err(ErrorResponse::reply("database error") + .with_status(StatusCode::INTERNAL_SERVER_ERROR)); + } + }; + + let session = match db.get_user_session(&user).await { + Ok(u) => u, + Err(DbError::NotFound) => { + debug!("user session not found for user id={}", user.id); + return Err(ErrorResponse::reply("user not found").with_status(StatusCode::NOT_FOUND)); + } + Err(DbError::Other(err)) => { + error!("database error for user {}: {}", login.username, err); + return Err(ErrorResponse::reply("database error") + .with_status(StatusCode::INTERNAL_SERVER_ERROR)); + } + }; + + let verified = verify_str(user.password.as_str(), login.password.borrow()); + + if !verified { + debug!(user = user.username, "login failed"); + return Err( + ErrorResponse::reply("password is not correct").with_status(StatusCode::UNAUTHORIZED) + ); + } + + debug!(user = user.username, "login success"); + + Ok(Json(LoginResponse { + session: session.token, + })) +} + +fn hash_secret(password: &str) -> String { + let arg2 = Argon2::new(Algorithm::Argon2id, Version::V0x13, Params::default()); + let salt = SaltString::generate(&mut OsRng); + let hash = arg2.hash_password(password.as_bytes(), &salt).unwrap(); + hash.to_string() +} diff --git a/crates/atuin-server/src/handlers/v0/me.rs b/crates/atuin-server/src/handlers/v0/me.rs new file mode 100644 index 00000000..7960b479 --- /dev/null +++ b/crates/atuin-server/src/handlers/v0/me.rs @@ -0,0 +1,16 @@ +use axum::Json; +use tracing::instrument; + +use crate::handlers::ErrorResponseStatus; +use crate::router::UserAuth; + +use atuin_common::api::*; + +#[instrument(skip_all, fields(user.id = user.id))] +pub async fn get( + UserAuth(user): UserAuth, +) -> Result<Json<MeResponse>, ErrorResponseStatus<'static>> { + Ok(Json(MeResponse { + username: user.username, + })) +} diff --git a/crates/atuin-server/src/handlers/v0/mod.rs b/crates/atuin-server/src/handlers/v0/mod.rs new file mode 100644 index 00000000..d6f880f2 --- /dev/null +++ b/crates/atuin-server/src/handlers/v0/mod.rs @@ -0,0 +1,3 @@ +pub(crate) mod me; +pub(crate) mod record; +pub(crate) mod store; diff --git a/crates/atuin-server/src/handlers/v0/record.rs b/crates/atuin-server/src/handlers/v0/record.rs new file mode 100644 index 00000000..5c57910b --- /dev/null +++ b/crates/atuin-server/src/handlers/v0/record.rs @@ -0,0 +1,114 @@ +use axum::{Json, extract::Query, extract::State, http::StatusCode}; +use metrics::counter; +use serde::Deserialize; +use tracing::{error, instrument}; + +use crate::{ + handlers::{ErrorResponse, ErrorResponseStatus, RespExt}, + router::{AppState, UserAuth}, +}; +use atuin_server_database::Database; + +use atuin_common::record::{EncryptedData, HostId, Record, RecordIdx, RecordStatus}; + +#[instrument(skip_all, fields(user.id = user.id))] +pub async fn post<DB: Database>( + UserAuth(user): UserAuth, + state: State<AppState<DB>>, + Json(records): Json<Vec<Record<EncryptedData>>>, +) -> Result<(), ErrorResponseStatus<'static>> { + let State(AppState { database, settings }) = state; + + tracing::debug!( + count = records.len(), + user = user.username, + "request to add records" + ); + + counter!("atuin_record_uploaded").increment(records.len() as u64); + + let keep = records + .iter() + .all(|r| r.data.data.len() <= settings.max_record_size || settings.max_record_size == 0); + + if !keep { + counter!("atuin_record_too_large").increment(1); + + return Err( + ErrorResponse::reply("could not add records; record too large") + .with_status(StatusCode::BAD_REQUEST), + ); + } + + if let Err(e) = database.add_records(&user, &records).await { + error!("failed to add record: {}", e); + + return Err(ErrorResponse::reply("failed to add record") + .with_status(StatusCode::INTERNAL_SERVER_ERROR)); + }; + + Ok(()) +} + +#[instrument(skip_all, fields(user.id = user.id))] +pub async fn index<DB: Database>( + UserAuth(user): UserAuth, + state: State<AppState<DB>>, +) -> Result<Json<RecordStatus>, ErrorResponseStatus<'static>> { + let State(AppState { + database, + settings: _, + }) = state; + + let record_index = match database.status(&user).await { + Ok(index) => index, + Err(e) => { + error!("failed to get record index: {}", e); + + return Err(ErrorResponse::reply("failed to calculate record index") + .with_status(StatusCode::INTERNAL_SERVER_ERROR)); + } + }; + + tracing::debug!(user = user.username, "record index request"); + + Ok(Json(record_index)) +} + +#[derive(Deserialize)] +pub struct NextParams { + host: HostId, + tag: String, + start: Option<RecordIdx>, + count: u64, +} + +#[instrument(skip_all, fields(user.id = user.id))] +pub async fn next<DB: Database>( + params: Query<NextParams>, + UserAuth(user): UserAuth, + state: State<AppState<DB>>, +) -> Result<Json<Vec<Record<EncryptedData>>>, ErrorResponseStatus<'static>> { + let State(AppState { + database, + settings: _, + }) = state; + let params = params.0; + + let records = match database + .next_records(&user, params.host, params.tag, params.start, params.count) + .await + { + Ok(records) => records, + Err(e) => { + error!("failed to get record index: {}", e); + + return Err(ErrorResponse::reply("failed to calculate record index") + .with_status(StatusCode::INTERNAL_SERVER_ERROR)); + } + }; + + counter!("atuin_record_downloaded").increment(records.len() as u64); + + Ok(Json(records)) +} diff --git a/crates/atuin-server/src/handlers/v0/store.rs b/crates/atuin-server/src/handlers/v0/store.rs new file mode 100644 index 00000000..6ca455d7 --- /dev/null +++ b/crates/atuin-server/src/handlers/v0/store.rs @@ -0,0 +1,37 @@ +use axum::{extract::Query, extract::State, http::StatusCode}; +use metrics::counter; +use serde::Deserialize; +use tracing::{error, instrument}; + +use crate::{ + handlers::{ErrorResponse, ErrorResponseStatus, RespExt}, + router::{AppState, UserAuth}, +}; +use atuin_server_database::Database; + +#[derive(Deserialize)] +pub struct DeleteParams {} + +#[instrument(skip_all, fields(user.id = user.id))] +pub async fn delete<DB: Database>( + _params: Query<DeleteParams>, + UserAuth(user): UserAuth, + state: State<AppState<DB>>, +) -> Result<(), ErrorResponseStatus<'static>> { + let State(AppState { + database, + settings: _, + }) = state; + + if let Err(e) = database.delete_store(&user).await { + counter!("atuin_store_delete_failed").increment(1); + error!("failed to delete store {e:?}"); + + return Err(ErrorResponse::reply("failed to delete store") + .with_status(StatusCode::INTERNAL_SERVER_ERROR)); + } + + counter!("atuin_store_deleted").increment(1); + + Ok(()) +} diff --git a/crates/atuin-server/src/lib.rs b/crates/atuin-server/src/lib.rs new file mode 100644 index 00000000..fcf5dde6 --- /dev/null +++ b/crates/atuin-server/src/lib.rs @@ -0,0 +1,98 @@ +#![forbid(unsafe_code)] + +use std::future::Future; +use std::net::SocketAddr; + +use atuin_server_database::Database; +use axum::{Router, serve}; +use eyre::{Context, Result}; + +mod handlers; +mod metrics; +mod router; +mod utils; + +pub use settings::Settings; +pub use settings::example_config; + +pub mod settings; + +use tokio::net::TcpListener; +use tokio::signal; + +#[cfg(target_family = "unix")] +async fn shutdown_signal() { + let mut term = signal::unix::signal(signal::unix::SignalKind::terminate()) + .expect("failed to register signal handler"); + let mut interrupt = signal::unix::signal(signal::unix::SignalKind::interrupt()) + .expect("failed to register signal handler"); + + tokio::select! { + _ = term.recv() => {}, + _ = interrupt.recv() => {}, + }; + eprintln!("Shutting down gracefully..."); +} + +#[cfg(target_family = "windows")] +async fn shutdown_signal() { + signal::windows::ctrl_c() + .expect("failed to register signal handler") + .recv() + .await; + eprintln!("Shutting down gracefully..."); +} + +pub async fn launch<Db: Database>(settings: Settings, addr: SocketAddr) -> Result<()> { + launch_with_tcp_listener::<Db>( + settings, + TcpListener::bind(addr) + .await + .context("could not connect to socket")?, + shutdown_signal(), + ) + .await +} + +pub async fn launch_with_tcp_listener<Db: Database>( + settings: Settings, + listener: TcpListener, + shutdown: impl Future<Output = ()> + Send + 'static, +) -> Result<()> { + let r = make_router::<Db>(settings).await?; + + serve(listener, r.into_make_service()) + .with_graceful_shutdown(shutdown) + .await?; + + Ok(()) +} + +// The separate listener means it's much easier to ensure metrics are not accidentally exposed to +// the public. +pub async fn launch_metrics_server(host: String, port: u16) -> Result<()> { + let listener = TcpListener::bind((host, port)) + .await + .context("failed to bind metrics tcp")?; + + let recorder_handle = metrics::setup_metrics_recorder(); + + let router = Router::new().route( + "/metrics", + axum::routing::get(move || std::future::ready(recorder_handle.render())), + ); + + serve(listener, router.into_make_service()) + .with_graceful_shutdown(shutdown_signal()) + .await?; + + Ok(()) +} + +async fn make_router<Db: Database>(settings: Settings) -> Result<Router, eyre::Error> { + let db = Db::new(&settings.db_settings) + .await + .wrap_err_with(|| format!("failed to connect to db: {:?}", settings.db_settings))?; + let r = router::router(db, settings); + Ok(r) +} diff --git a/crates/atuin-server/src/metrics.rs b/crates/atuin-server/src/metrics.rs new file mode 100644 index 00000000..ebd0dd2d --- /dev/null +++ b/crates/atuin-server/src/metrics.rs @@ -0,0 +1,55 @@ +use std::time::Instant; + +use axum::{ + extract::{MatchedPath, Request}, + middleware::Next, + response::IntoResponse, +}; +use metrics_exporter_prometheus::{Matcher, PrometheusBuilder, PrometheusHandle}; + +pub fn setup_metrics_recorder() -> PrometheusHandle { + const EXPONENTIAL_SECONDS: &[f64] = &[ + 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0, + ]; + + PrometheusBuilder::new() + .set_buckets_for_metric( + Matcher::Full("http_requests_duration_seconds".to_string()), + EXPONENTIAL_SECONDS, + ) + .unwrap() + .install_recorder() + .unwrap() +} + +/// Middleware to record some common HTTP metrics +/// Generic over B to allow for arbitrary body types (eg Vec<u8>, Streams, a deserialized thing, etc) +/// Someday tower-http might provide a metrics middleware: https://github.com/tower-rs/tower-http/issues/57 +pub async fn track_metrics(req: Request, next: Next) -> impl IntoResponse { + let start = Instant::now(); + + let path = match req.extensions().get::<MatchedPath>() { + Some(matched_path) => matched_path.as_str().to_owned(), + _ => req.uri().path().to_owned(), + }; + + let method = req.method().clone(); + + // Run the rest of the request handling first, so we can measure it and get response + // codes. + let response = next.run(req).await; + + let latency = start.elapsed().as_secs_f64(); + let status = response.status().as_u16().to_string(); + + let labels = [ + ("method", method.to_string()), + ("path", path), + ("status", status), + ]; + + metrics::counter!("http_requests_total", &labels).increment(1); + metrics::histogram!("http_requests_duration_seconds", &labels).record(latency); + + response +} diff --git a/crates/atuin-server/src/router.rs b/crates/atuin-server/src/router.rs new file mode 100644 index 00000000..2d679759 --- /dev/null +++ b/crates/atuin-server/src/router.rs @@ -0,0 +1,155 @@ +use atuin_common::api::{ATUIN_CARGO_VERSION, ATUIN_HEADER_VERSION, ErrorResponse}; +use axum::{ + Router, + extract::{FromRequestParts, Request}, + http::{self, request::Parts}, + middleware::Next, + response::{IntoResponse, Response}, + routing::{delete, get, patch, post}, +}; +use eyre::Result; +use tower::ServiceBuilder; +use tower_http::trace::TraceLayer; + +use super::handlers; +use crate::{ + handlers::{ErrorResponseStatus, RespExt}, + metrics, + settings::Settings, +}; +use atuin_server_database::{Database, DbError, models::User}; + +pub struct UserAuth(pub User); + +impl<DB: Send + Sync> FromRequestParts<AppState<DB>> for UserAuth +where + DB: Database, +{ + type Rejection = ErrorResponseStatus<'static>; + + async fn from_request_parts( + req: &mut Parts, + state: &AppState<DB>, + ) -> Result<Self, Self::Rejection> { + let auth_header = req + .headers + .get(http::header::AUTHORIZATION) + .ok_or_else(|| { + ErrorResponse::reply("missing authorization header") + .with_status(http::StatusCode::BAD_REQUEST) + })?; + let auth_header = auth_header.to_str().map_err(|_| { + ErrorResponse::reply("invalid authorization header encoding") + .with_status(http::StatusCode::BAD_REQUEST) + })?; + let (typ, token) = auth_header.split_once(' ').ok_or_else(|| { + ErrorResponse::reply("invalid authorization header encoding") + .with_status(http::StatusCode::BAD_REQUEST) + })?; + + if typ != "Token" { + return Err( + ErrorResponse::reply("invalid authorization header encoding") + .with_status(http::StatusCode::BAD_REQUEST), + ); + } + + let user = state + .database + .get_session_user(token) + .await + .map_err(|e| match e { + DbError::NotFound => ErrorResponse::reply("session not found") + .with_status(http::StatusCode::FORBIDDEN), + DbError::Other(e) => { + tracing::error!(error = ?e, "could not query user session"); + ErrorResponse::reply("could not query user session") + .with_status(http::StatusCode::INTERNAL_SERVER_ERROR) + } + })?; + + Ok(UserAuth(user)) + } +} + +async fn teapot() -> impl IntoResponse { + // This used to return 418: 🫖 + // Much as it was fun, it wasn't as useful or informative as it should be + (http::StatusCode::NOT_FOUND, "404 not found") +} + +async fn clacks_overhead(request: Request, next: Next) -> Response { + let mut response = next.run(request).await; + + let gnu_terry_value = "GNU Terry Pratchett, Kris Nova"; + let gnu_terry_header = "X-Clacks-Overhead"; + + response + .headers_mut() + .insert(gnu_terry_header, gnu_terry_value.parse().unwrap()); + response +} + +/// Ensure that we only try and sync with clients on the same major version +async fn semver(request: Request, next: Next) -> Response { + let mut response = next.run(request).await; + response + .headers_mut() + .insert(ATUIN_HEADER_VERSION, ATUIN_CARGO_VERSION.parse().unwrap()); + + response +} + +#[derive(Clone)] +pub struct AppState<DB: Database> { + pub database: DB, + pub settings: Settings, +} + +pub fn router<DB: Database>(database: DB, settings: Settings) -> Router { + let mut routes = Router::new() + .route("/", get(handlers::index)) + .route("/healthz", get(handlers::health::health_check)); + + // Sync v1 routes - can be disabled in favor of record-based sync + if settings.sync_v1_enabled { + routes = routes + .route("/sync/count", get(handlers::history::count)) + .route("/sync/history", get(handlers::history::list)) + .route("/sync/calendar/{focus}", get(handlers::history::calendar)) + .route("/sync/status", get(handlers::status::status)) + .route("/history", post(handlers::history::add)) + .route("/history", delete(handlers::history::delete)); + } + + let routes = routes + .route("/user/{username}", get(handlers::user::get)) + .route("/account", delete(handlers::user::delete)) + .route("/account/password", patch(handlers::user::change_password)) + .route("/register", post(handlers::user::register)) + .route("/login", post(handlers::user::login)) + .route("/record", post(handlers::record::post)) + .route("/record", get(handlers::record::index)) + .route("/record/next", get(handlers::record::next)) + .route("/api/v0/me", get(handlers::v0::me::get)) + .route("/api/v0/record", post(handlers::v0::record::post)) + .route("/api/v0/record", get(handlers::v0::record::index)) + .route("/api/v0/record/next", get(handlers::v0::record::next)) + .route("/api/v0/store", delete(handlers::v0::store::delete)); + + let path = settings.path.as_str(); + if path.is_empty() { + routes + } else { + Router::new().nest(path, routes) + } + .fallback(teapot) + .with_state(AppState { database, settings }) + .layer( + ServiceBuilder::new() + .layer(axum::middleware::from_fn(clacks_overhead)) + .layer(TraceLayer::new_for_http()) + .layer(axum::middleware::from_fn(metrics::track_metrics)) + .layer(axum::middleware::from_fn(semver)), + ) +} diff --git a/crates/atuin-server/src/settings.rs b/crates/atuin-server/src/settings.rs new file mode 100644 index 00000000..3a612be9 --- /dev/null +++ b/crates/atuin-server/src/settings.rs @@ -0,0 +1,113 @@ +use std::{io::prelude::*, path::PathBuf}; + +use atuin_server_database::DbSettings; +use config::{Config, Environment, File as ConfigFile, FileFormat}; +use eyre::{Result, eyre}; +use fs_err::{File, create_dir_all}; +use serde::{Deserialize, Serialize}; + +static EXAMPLE_CONFIG: &str = include_str!("../server.toml"); + +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct Metrics { + #[serde(alias = "enabled")] + pub enable: bool, + pub host: String, + pub port: u16, +} + +impl Default for Metrics { + fn default() -> Self { + Self { + enable: false, + host: String::from("127.0.0.1"), + port: 9001, + } + } +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct Settings { + pub host: String, + pub port: u16, + pub path: String, + pub open_registration: bool, + pub max_history_length: usize, + pub max_record_size: usize, + pub page_size: i64, + pub register_webhook_url: Option<String>, + pub register_webhook_username: String, + pub metrics: Metrics, + + /// Enable legacy sync v1 routes (history-based sync) + /// Set to false to use only the newer record-based sync + pub sync_v1_enabled: bool, + + /// Advertise a version that is not what we are _actually_ running + /// Many clients compare their version with api.atuin.sh, and if they differ, notify the user + /// that an update is available. + /// Now that we take beta releases, we should be able to advertise a different version to avoid + /// notifying users when the server runs something that is not a stable release. + pub fake_version: Option<String>, + + #[serde(flatten)] + pub db_settings: DbSettings, +} + +impl Settings { + pub fn new() -> Result<Self> { + let mut config_file = if let Ok(p) = std::env::var("ATUIN_CONFIG_DIR") { + PathBuf::from(p) + } else { + let mut config_file = PathBuf::new(); + let config_dir = atuin_common::utils::config_dir(); + config_file.push(config_dir); + config_file + }; + + config_file.push("server.toml"); + + // create the config file if it does not exist + let mut config_builder = Config::builder() + .set_default("host", "127.0.0.1")? + .set_default("port", 8888)? + .set_default("open_registration", false)? + .set_default("max_history_length", 8192)? + .set_default("max_record_size", 1024 * 1024 * 1024)? // pretty chonky + .set_default("path", "")? + .set_default("register_webhook_username", "")? + .set_default("page_size", 1100)? + .set_default("metrics.enable", false)? + .set_default("metrics.host", "127.0.0.1")? + .set_default("metrics.port", 9001)? + .set_default("sync_v1_enabled", true)? + .add_source( + Environment::with_prefix("atuin") + .prefix_separator("_") + .separator("__"), + ); + + config_builder = if config_file.exists() { + config_builder.add_source(ConfigFile::new( + config_file.to_str().unwrap(), + FileFormat::Toml, + )) + } else { + create_dir_all(config_file.parent().unwrap())?; + let mut file = File::create(config_file)?; + file.write_all(EXAMPLE_CONFIG.as_bytes())?; + + config_builder + }; + + let config = config_builder.build()?; + + config + .try_deserialize() + .map_err(|e| eyre!("failed to deserialize: {}", e)) + } +} + +pub fn example_config() -> &'static str { + EXAMPLE_CONFIG +} diff --git a/crates/atuin-server/src/utils.rs b/crates/atuin-server/src/utils.rs new file mode 100644 index 00000000..12e9ac1b --- /dev/null +++ b/crates/atuin-server/src/utils.rs @@ -0,0 +1,15 @@ +use eyre::Result; +use semver::{Version, VersionReq}; + +pub fn client_version_min(user_agent: &str, req: &str) -> Result<bool> { + if user_agent.is_empty() { + return Ok(false); + } + + let version = user_agent.replace("atuin/", ""); + + let req = VersionReq::parse(req)?; + let version = Version::parse(version.as_str())?; + + Ok(req.matches(&version)) +} diff --git a/crates/atuin/Cargo.toml b/crates/atuin/Cargo.toml new file mode 100644 index 00000000..5eed945d --- /dev/null +++ b/crates/atuin/Cargo.toml @@ -0,0 +1,115 @@ +[package] +name = "atuin" +edition = "2024" +description = "atuin - magical shell history" +readme = "./README.md" + +rust-version = { workspace = true } +version = { workspace = true } +authors = { workspace = true } +license = { workspace = true } +homepage = { workspace = true } +repository = { workspace = true } + +[package.metadata.binstall] +pkg-url = "{ repo }/releases/download/v{ version }/{ name }-{ target }.tar.gz" +bin-dir = "{ name }-{ target }/{ bin }{ binary-ext }" +pkg-fmt = "tgz" + +[package.metadata.deb] +maintainer = "Ellie Huxtable <ellie@elliehuxtable.com>" +copyright = "2021, Ellie Huxtable <ellie@elliehuxtable.com>" +license-file = ["LICENSE"] +depends = "$auto" +section = "utility" + +[package.metadata.rpm] +package = "atuin" + +[package.metadata.rpm.cargo] +buildflags = ["--release"] + +[package.metadata.rpm.targets] +atuin = { path = "/usr/bin/atuin" } + +[features] +default = ["client", "sync", "clipboard", "check-update", "daemon", "ai", "hex"] +client = ["atuin-client"] +sync = ["atuin-client/sync"] +daemon = ["atuin-client/daemon", "atuin-daemon"] +ai = ["atuin-ai"] +hex = ["atuin-hex"] +clipboard = ["arboard"] +check-update = ["atuin-client/check-update"] + +[dependencies] +atuin-ai = { path = "../atuin-ai", version = "18.13.2", optional = true, default-features = false } +atuin-client = { path = "../atuin-client", version = "18.13.2", optional = true, default-features = false } +atuin-common = { workspace = true } +atuin-dotfiles = { workspace = true } +atuin-history = { workspace = true } +atuin-daemon = { path = "../atuin-daemon", version = "18.13.2", optional = true, default-features = false } +atuin-hex = { path = "../atuin-hex", version = "18.13.2", optional = true, default-features = false } +atuin-scripts = { workspace = true } +atuin-kv = { workspace = true } + +log = { workspace = true } +time = { workspace = true } +eyre = { workspace = true } +indicatif = "0.18.0" +serde = { workspace = true } +serde_json = { workspace = true } +crossterm = { workspace = true, features = ["use-dev-tty"] } +unicode-width = "0.2" +itertools = { workspace = true } +tokio = { workspace = true } +async-trait = { workspace = true } +interim = { workspace = true } +clap = { workspace = true } +clap_complete = "4.5.8" +clap_complete_nushell = "4.5.4" +fs-err = { workspace = true } +fs4 = "0.13.1" +rpassword = "7.0" +semver = { workspace = true } +rustix = { workspace = true } +runtime-format = "0.1.3" +tiny-bip39 = "2" +futures-util = "0.3" +fuzzy-matcher = "0.3.7" +colored = "2.0.4" +open = "5" +ratatui = { workspace = true } +tracing = "0.1" +tracing-subscriber = { workspace = true } +tracing-appender = "0.2" +uuid = { workspace = true } +sysinfo = "0.30.7" +regex = "1.10.5" +norm = { version = "0.1.1", features = ["fzf-v2"] } +nucleo-matcher = { git = "https://github.com/atuinsh/nucleo-ext.git", rev = "74bd786" } +tempfile = { workspace = true } +shlex = "1.3.0" + +# settings editor with comment and relative ordering preservation +toml_edit = "0.25.4" + +[target.'cfg(any(target_os = "windows", target_os = "macos"))'.dependencies] +arboard = { version = "3.4", optional = true } + +[target.'cfg(target_os = "linux")'.dependencies] +arboard = { version = "3.4", optional = true, features = [ + "wayland-data-control", +] } + +[target.'cfg(unix)'.dependencies] +daemonize = "0.5.0" + +[dev-dependencies] +tracing-tree = "0.4" + +# Integration tests in tests/ spin up a test server to verify sync functionality. +# TODO: Consider moving these tests to atuin-server crate instead (client would become a dev dep there) +atuin-server = { workspace = true } +atuin-server-database = { workspace = true } +atuin-server-postgres = { workspace = true } diff --git a/crates/atuin/LICENSE b/crates/atuin/LICENSE new file mode 100644 index 00000000..7dfc9b58 --- /dev/null +++ b/crates/atuin/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2021 Ellie Huxtable + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/crates/atuin/README.md b/crates/atuin/README.md new file mode 120000 index 00000000..fe840054 --- /dev/null +++ b/crates/atuin/README.md @@ -0,0 +1 @@ +../../README.md
\ No newline at end of file diff --git a/crates/atuin/build.rs b/crates/atuin/build.rs new file mode 100644 index 00000000..75d53ee0 --- /dev/null +++ b/crates/atuin/build.rs @@ -0,0 +1,11 @@ +use std::process::Command; +fn main() { + let output = Command::new("git").args(["rev-parse", "HEAD"]).output(); + + let sha = match output { + Ok(sha) => String::from_utf8(sha.stdout).unwrap(), + Err(_) => String::from("NO_GIT"), + }; + + println!("cargo:rustc-env=GIT_HASH={sha}"); +} diff --git a/crates/atuin/src/command/CONTRIBUTORS b/crates/atuin/src/command/CONTRIBUTORS new file mode 120000 index 00000000..1ca4115a --- /dev/null +++ b/crates/atuin/src/command/CONTRIBUTORS @@ -0,0 +1 @@ +../../../../CONTRIBUTORS
\ No newline at end of file diff --git a/crates/atuin/src/command/client.rs b/crates/atuin/src/command/client.rs new file mode 100644 index 00000000..02d64205 --- /dev/null +++ b/crates/atuin/src/command/client.rs @@ -0,0 +1,381 @@ +use std::fs::{self, OpenOptions}; +use std::path::{Path, PathBuf}; + +use clap::Subcommand; +use eyre::{Result, WrapErr}; + +use atuin_client::{ + database::Sqlite, record::sqlite_store::SqliteStore, settings::Settings, theme, +}; +use tracing_appender::rolling::{RollingFileAppender, Rotation}; +use tracing_subscriber::{ + Layer, filter::EnvFilter, filter::LevelFilter, fmt, fmt::format::FmtSpan, prelude::*, +}; + +fn cleanup_old_logs(log_dir: &Path, prefix: &str, retention_days: u64) { + let cutoff = std::time::SystemTime::now() + - std::time::Duration::from_secs(retention_days * 24 * 60 * 60); + + let Ok(entries) = fs::read_dir(log_dir) else { + return; + }; + + for entry in entries.flatten() { + let path = entry.path(); + let Some(name) = path.file_name().and_then(|n| n.to_str()) else { + continue; + }; + + // Match files like "search.log.2024-02-23" or "daemon.log.2024-02-23" + if !name.starts_with(prefix) || name == prefix { + continue; + } + + if let Ok(metadata) = entry.metadata() + && let Ok(modified) = metadata.modified() + && modified < cutoff + { + let _ = fs::remove_file(&path); + } + } +} + +#[cfg(feature = "sync")] +mod sync; + +#[cfg(feature = "sync")] +mod account; + +#[cfg(feature = "daemon")] +mod daemon; + +mod default_config; +mod doctor; +mod dotfiles; +mod history; +mod import; +mod info; +mod init; +mod kv; +mod scripts; +mod search; +mod setup; +mod stats; +mod store; +mod wrapped; + +#[derive(Subcommand, Debug)] +#[command(infer_subcommands = true)] +pub enum Cmd { + /// Setup Atuin features + #[command()] + Setup, + + /// Manipulate shell history + #[command(subcommand)] + History(history::Cmd), + + /// Import shell history from file + #[command(subcommand)] + Import(import::Cmd), + + /// Calculate statistics for your history + Stats(stats::Cmd), + + /// Interactive history search + Search(search::Cmd), + + #[cfg(feature = "sync")] + #[command(flatten)] + Sync(sync::Cmd), + + /// Manage your sync account + #[cfg(feature = "sync")] + Account(account::Cmd), + + /// Get or set small key-value pairs + #[command(subcommand)] + Kv(kv::Cmd), + + /// Manage the atuin data store + #[command(subcommand)] + Store(store::Cmd), + + /// Manage your dotfiles with Atuin + #[command(subcommand)] + Dotfiles(dotfiles::Cmd), + + /// Manage your scripts with Atuin + #[command(subcommand)] + Scripts(scripts::Cmd), + + /// Print Atuin's shell init script + #[command()] + Init(init::Cmd), + + /// Information about dotfiles locations and ENV vars + #[command()] + Info, + + /// Run the doctor to check for common issues + #[command()] + Doctor, + + #[command()] + Wrapped { year: Option<i32> }, + + /// *Experimental* Manage the background daemon + #[cfg(feature = "daemon")] + #[command()] + Daemon(daemon::Cmd), + + /// Print the default atuin configuration (config.toml) + #[command()] + DefaultConfig, + + /// Run the AI assistant + #[cfg(feature = "ai")] + #[command(subcommand)] + Ai(atuin_ai::commands::Commands), +} + +impl Cmd { + pub fn run(self) -> Result<()> { + // Daemonize before creating the async runtime – fork() inside a live + // tokio runtime corrupts its internal state. + #[cfg(all(unix, feature = "daemon"))] + if let Self::Daemon(ref cmd) = self + && cmd.should_daemonize() + { + daemon::daemonize_current_process()?; + } + + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + + let settings = Settings::new().wrap_err("could not load client settings")?; + let theme_manager = theme::ThemeManager::new(settings.theme.debug, None); + let res = runtime.block_on(self.run_inner(settings, theme_manager)); + + runtime.shutdown_timeout(std::time::Duration::from_millis(50)); + + res + } + + #[allow(clippy::too_many_lines)] + async fn run_inner( + self, + mut settings: Settings, + mut theme_manager: theme::ThemeManager, + ) -> Result<()> { + // ATUIN_LOG env var overrides config file level settings + let env_log_set = std::env::var("ATUIN_LOG").is_ok(); + + // Base filter from env var (or empty if not set) + let base_filter = + EnvFilter::from_env("ATUIN_LOG").add_directive("sqlx_sqlite::regexp=off".parse()?); + + let is_interactive_search = matches!(&self, Self::Search(cmd) if cmd.is_interactive()); + // Use file-based logging for interactive search (TUI mode) + let use_search_logging = is_interactive_search && settings.logs.search_enabled(); + + // Use file-based logging for daemon + #[cfg(feature = "daemon")] + let use_daemon_logging = matches!(&self, Self::Daemon(_)) && settings.logs.daemon_enabled(); + + #[cfg(not(feature = "daemon"))] + let use_daemon_logging = false; + + // Check if daemon should also log to console + #[cfg(feature = "daemon")] + let daemon_show_logs = matches!(&self, Self::Daemon(cmd) if cmd.show_logs()); + + #[cfg(not(feature = "daemon"))] + let daemon_show_logs = false; + + // Set up span timing JSON logs if ATUIN_SPAN is set + let span_path = std::env::var("ATUIN_SPAN").ok().map(|p| { + if p.is_empty() { + "atuin-spans.json".to_string() + } else { + p + } + }); + + // Helper to create span timing layer + macro_rules! make_span_layer { + ($path:expr) => {{ + let span_file = OpenOptions::new() + .create(true) + .truncate(true) + .write(true) + .open($path)?; + Some( + fmt::layer() + .json() + .with_writer(span_file) + .with_span_events(FmtSpan::NEW | FmtSpan::CLOSE) + .with_filter(LevelFilter::TRACE), + ) + }}; + } + + // Build the subscriber with all configured layers + if use_search_logging { + let search_filename = settings.logs.search.file.clone(); + let log_dir = PathBuf::from(&settings.logs.dir); + fs::create_dir_all(&log_dir)?; + + // Clean up old log files + cleanup_old_logs(&log_dir, &search_filename, settings.logs.search_retention()); + + let file_appender = + RollingFileAppender::new(Rotation::DAILY, &log_dir, &search_filename); + + // Use config level unless ATUIN_LOG is set + let filter = if env_log_set { + base_filter + } else { + EnvFilter::default() + .add_directive(settings.logs.search_level().as_directive().parse()?) + .add_directive("sqlx_sqlite::regexp=off".parse()?) + }; + + let base = tracing_subscriber::registry().with( + fmt::layer() + .with_writer(file_appender) + .with_ansi(false) + .with_filter(filter), + ); + + match &span_path { + Some(sp) => { + base.with(make_span_layer!(sp)).init(); + } + None => { + base.init(); + } + } + } else if use_daemon_logging { + let daemon_filename = settings.logs.daemon.file.clone(); + let log_dir = PathBuf::from(&settings.logs.dir); + fs::create_dir_all(&log_dir)?; + + // Clean up old log files + cleanup_old_logs(&log_dir, &daemon_filename, settings.logs.daemon_retention()); + + let file_appender = + RollingFileAppender::new(Rotation::DAILY, &log_dir, &daemon_filename); + + // Use config level unless ATUIN_LOG is set + let file_filter = if env_log_set { + base_filter + } else { + EnvFilter::default() + .add_directive(settings.logs.daemon_level().as_directive().parse()?) + .add_directive("sqlx_sqlite::regexp=off".parse()?) + }; + + let file_layer = fmt::layer() + .with_writer(file_appender) + .with_ansi(false) + .with_filter(file_filter); + + // Optionally add console layer for --show-logs + if daemon_show_logs { + let console_filter = EnvFilter::from_env("ATUIN_LOG") + .add_directive("sqlx_sqlite::regexp=off".parse()?); + + let console_layer = fmt::layer().with_filter(console_filter); + + let base = tracing_subscriber::registry() + .with(file_layer) + .with(console_layer); + + match &span_path { + Some(sp) => { + base.with(make_span_layer!(sp)).init(); + } + None => { + base.init(); + } + } + } else { + let base = tracing_subscriber::registry().with(file_layer); + + match &span_path { + Some(sp) => { + base.with(make_span_layer!(sp)).init(); + } + None => { + base.init(); + } + } + } + } + + tracing::trace!(command = ?self, "client command"); + + // Skip initializing any databases for history + // This is a pretty hot path, as it runs before and after every single command the user + // runs + match self { + Self::History(history) => return history.run(&settings).await, + Self::Init(init) => return init.run(&settings).await, + Self::Doctor => return doctor::run(&settings).await, + _ => {} + } + + let db_path = PathBuf::from(settings.db_path.as_str()); + let record_store_path = PathBuf::from(settings.record_store_path.as_str()); + + let db = Sqlite::new(db_path, settings.local_timeout).await?; + let sqlite_store = SqliteStore::new(record_store_path, settings.local_timeout).await?; + + let theme_name = settings.theme.name.clone(); + let theme = theme_manager.load_theme(theme_name.as_str(), settings.theme.max_depth); + + match self { + Self::Setup => setup::run(&settings).await, + Self::Import(import) => import.run(&db).await, + Self::Stats(stats) => stats.run(&db, &settings, theme).await, + Self::Search(search) => search.run(db, &mut settings, sqlite_store, theme).await, + + #[cfg(feature = "sync")] + Self::Sync(sync) => sync.run(settings, &db, sqlite_store).await, + + #[cfg(feature = "sync")] + Self::Account(account) => account.run(settings, sqlite_store).await, + + Self::Kv(kv) => kv.run(&settings, &sqlite_store).await, + + Self::Store(store) => store.run(&settings, &db, sqlite_store).await, + + Self::Dotfiles(dotfiles) => dotfiles.run(&settings, sqlite_store).await, + + Self::Scripts(scripts) => scripts.run(&settings, sqlite_store, &db).await, + + Self::Info => { + info::run(&settings); + Ok(()) + } + + Self::DefaultConfig => { + default_config::run(); + Ok(()) + } + + Self::Wrapped { year } => wrapped::run(year, &db, &settings, sqlite_store, theme).await, + + #[cfg(feature = "daemon")] + Self::Daemon(cmd) => cmd.run(settings, sqlite_store, db).await, + + Self::History(_) | Self::Init(_) | Self::Doctor => unreachable!(), + + #[cfg(feature = "ai")] + Self::Ai(cli) => atuin_ai::commands::run(cli, &settings).await, + } + } +} diff --git a/crates/atuin/src/command/client/account.rs b/crates/atuin/src/command/client/account.rs new file mode 100644 index 00000000..e9861c0a --- /dev/null +++ b/crates/atuin/src/command/client/account.rs @@ -0,0 +1,54 @@ +use clap::{Args, Subcommand}; +use eyre::Result; + +use atuin_client::record::sqlite_store::SqliteStore; +use atuin_client::settings::Settings; + +pub mod change_password; +pub mod delete; +pub mod link; +pub mod login; +pub mod logout; +pub mod register; + +const DEFAULT_HUB_ENDPOINT: &str = "https://hub.atuin.sh"; + +#[derive(Args, Debug)] +pub struct Cmd { + #[command(subcommand)] + command: Commands, +} + +#[derive(Subcommand, Debug)] +pub enum Commands { + /// Login to the configured server + Login(login::Cmd), + + /// Register a new account + Register(register::Cmd), + + /// Log out + Logout, + + /// Delete your account, and all synced data + Delete, + + /// Change your password + ChangePassword(change_password::Cmd), + + /// Link your CLI sync account to your Hub account + Link, +} + +impl Cmd { + pub async fn run(self, settings: Settings, store: SqliteStore) -> Result<()> { + match self.command { + Commands::Login(l) => l.run(&settings, &store).await, + Commands::Register(r) => r.run(&settings, &store).await, + Commands::Logout => logout::run().await, + Commands::Delete => delete::run(&settings).await, + Commands::ChangePassword(c) => c.run(&settings).await, + Commands::Link => link::run(&settings).await, + } + } +} diff --git a/crates/atuin/src/command/client/account/change_password.rs b/crates/atuin/src/command/client/account/change_password.rs new file mode 100644 index 00000000..0f4a31cd --- /dev/null +++ b/crates/atuin/src/command/client/account/change_password.rs @@ -0,0 +1,77 @@ +use clap::Parser; +use eyre::{Result, bail}; + +use atuin_client::{api_client, settings::Settings}; +use rpassword::prompt_password; + +use crate::command::client::account::DEFAULT_HUB_ENDPOINT; + +#[derive(Parser, Debug)] +pub struct Cmd { + #[clap(long, short)] + pub current_password: Option<String>, + + #[clap(long, short)] + pub new_password: Option<String>, +} + +impl Cmd { + pub async fn run(self, settings: &Settings) -> Result<()> { + run(settings, self.current_password, self.new_password).await + } +} + +pub async fn run( + settings: &Settings, + current_password: Option<String>, + new_password: Option<String>, +) -> Result<()> { + let using_hub_sync = settings.is_hub_sync(); + let has_sync_session = settings.session_token().await.is_ok(); + let has_hub_session = settings.hub_session_token().await.is_ok(); + + if using_hub_sync && has_hub_session { + let endpoint = settings + .active_hub_endpoint() + .unwrap_or_else(|| DEFAULT_HUB_ENDPOINT.to_string()); + + println!("You are authenticated with Atuin Hub."); + println!("Manage your account on the site: {endpoint}/settings/account"); + return Ok(()); + } + + if !has_sync_session { + bail!("You are not logged in"); + } + + let client = api_client::Client::new( + &settings.sync_address, + settings.sync_auth_token().await?, + settings.network_connect_timeout, + settings.network_timeout, + )?; + + let current_password = current_password.clone().unwrap_or_else(|| { + prompt_password("Please enter the current password: ").expect("Failed to read from input") + }); + + if current_password.is_empty() { + bail!("please provide the current password"); + } + + let new_password = new_password.clone().unwrap_or_else(|| { + prompt_password("Please enter the new password: ").expect("Failed to read from input") + }); + + if new_password.is_empty() { + bail!("please provide a new password"); + } + + client + .change_password(current_password, new_password) + .await?; + + println!("Account password successfully changed!"); + + Ok(()) +} diff --git a/crates/atuin/src/command/client/account/delete.rs b/crates/atuin/src/command/client/account/delete.rs new file mode 100644 index 00000000..5c0439a3 --- /dev/null +++ b/crates/atuin/src/command/client/account/delete.rs @@ -0,0 +1,40 @@ +use atuin_client::{api_client, settings::Settings}; +use eyre::{Result, bail}; + +use crate::command::client::account::DEFAULT_HUB_ENDPOINT; + +pub async fn run(settings: &Settings) -> Result<()> { + let using_hub_sync = settings.is_hub_sync(); + let has_sync_session = settings.session_token().await.is_ok(); + let has_hub_session = settings.hub_session_token().await.is_ok(); + + if using_hub_sync && has_hub_session { + let endpoint = settings + .active_hub_endpoint() + .unwrap_or_else(|| DEFAULT_HUB_ENDPOINT.to_string()); + println!("You are authenticated with Atuin Hub."); + println!("Manage your account on the site: {endpoint}/settings/account"); + return Ok(()); + } + + if !has_sync_session { + bail!("You are not logged in"); + } + + let client = api_client::Client::new( + &settings.sync_address, + settings.sync_auth_token().await?, + settings.network_connect_timeout, + settings.network_timeout, + )?; + + client.delete().await?; + + // Clean up session from meta store + Settings::meta_store().await?.delete_session().await?; + Settings::meta_store().await?.delete_hub_session().await?; + + println!("Your account is deleted"); + + Ok(()) +} diff --git a/crates/atuin/src/command/client/account/link.rs b/crates/atuin/src/command/client/account/link.rs new file mode 100644 index 00000000..5a2e4044 --- /dev/null +++ b/crates/atuin/src/command/client/account/link.rs @@ -0,0 +1,45 @@ +use eyre::{Result, bail}; + +use atuin_client::settings::Settings; + +use super::DEFAULT_HUB_ENDPOINT; + +pub async fn run(settings: &Settings) -> Result<()> { + let meta = Settings::meta_store().await?; + + let cli_token = meta.session_token().await?; + let hub_token = meta.hub_session_token().await?; + + let Some(cli_token) = cli_token else { + bail!("No CLI session found. Please log in first with 'atuin login'."); + }; + + let hub_address = settings + .active_hub_endpoint() + .unwrap_or_else(|| DEFAULT_HUB_ENDPOINT.to_string()); + + if hub_token.is_some() { + println!("Found both Hub and CLI sessions. Linking accounts..."); + } else { + println!("Found CLI session but no Hub session. Logging in to Hub first..."); + + let session = atuin_client::hub::HubAuthSession::start(&hub_address).await?; + println!("Open this URL to authenticate with Atuin Hub:"); + println!("{}", session.auth_url); + + let token = session + .wait_for_completion( + atuin_client::hub::DEFAULT_AUTH_TIMEOUT, + atuin_client::hub::DEFAULT_POLL_INTERVAL, + ) + .await?; + + atuin_client::hub::save_session(&token).await?; + println!("Hub authentication complete."); + } + + atuin_client::hub::link_account(&hub_address, &cli_token).await?; + println!("Successfully linked CLI account to Hub."); + + Ok(()) +} diff --git a/crates/atuin/src/command/client/account/login.rs b/crates/atuin/src/command/client/account/login.rs new file mode 100644 index 00000000..b8aad5a9 --- /dev/null +++ b/crates/atuin/src/command/client/account/login.rs @@ -0,0 +1,257 @@ +use std::{io, path::PathBuf}; + +use clap::Parser; +use eyre::{Context, Result, bail}; +use tokio::{fs::File, io::AsyncWriteExt}; + +use atuin_client::{ + api_client, + encryption::{Key, decode_key, encode_key, load_key}, + record::sqlite_store::SqliteStore, + record::store::Store, + settings::Settings, +}; +use atuin_common::api::LoginRequest; +use rpassword::prompt_password; + +#[derive(Parser, Debug)] +pub struct Cmd { + #[clap(long, short)] + pub username: Option<String>, + + #[clap(long, short)] + pub password: Option<String>, + + /// The encryption key for your account + #[clap(long, short)] + pub key: Option<String>, + + #[clap(long, hide = true)] + pub from_registration: bool, +} + +fn get_input() -> Result<String> { + let mut input = String::new(); + io::stdin().read_line(&mut input)?; + Ok(input.trim_end_matches(&['\r', '\n'][..]).to_string()) +} + +impl Cmd { + pub async fn run(&self, settings: &Settings, store: &SqliteStore) -> Result<()> { + if let Some(endpoint) = settings.active_hub_endpoint() { + if settings.hub_session_token().await.is_ok() { + println!("You are authenticated with Atuin Hub."); + println!("Run 'atuin logout' to log out."); + return Ok(()); + } + + // The only difference between login and registration is that registration doesn't prompt for a key + if self.from_registration { + load_key(settings)?; + } else { + self.prompt_and_store_key(settings, store).await?; + } + + self.ensure_hub_session(settings, endpoint.as_str()).await?; + println!("Successfully authenticated with Atuin Hub."); + return Ok(()); + } + + if settings.logged_in().await? { + println!("You are already logged in."); + println!("Run 'atuin logout' to log out."); + return Ok(()); + } + + self.run_sync_login(settings, store).await + } + + async fn ensure_hub_session(&self, settings: &Settings, hub_address: &str) -> Result<()> { + tracing::info!("Authenticating with Atuin Hub..."); + + let session = atuin_client::hub::HubAuthSession::start(hub_address).await?; + println!("Open this URL to continue authenticating with Atuin Hub:"); + println!("{}", session.auth_url); + + let token = session + .wait_for_completion( + atuin_client::hub::DEFAULT_AUTH_TIMEOUT, + atuin_client::hub::DEFAULT_POLL_INTERVAL, + ) + .await?; + + tracing::info!("Authentication complete, saving session token"); + + atuin_client::hub::save_session(&token).await?; + + // Silently attempt to link CLI account to Hub if one exists + // This enables unified auth - users can use their Hub token for sync + if let Ok(cli_token) = settings.session_token().await { + tracing::debug!("CLI session found, attempting to link accounts"); + if let Err(e) = atuin_client::hub::link_account(hub_address, &cli_token).await { + tracing::debug!("Could not link CLI account to Hub: {}", e); + } else { + tracing::info!("Successfully linked CLI account to Hub"); + } + } + + Ok(()) + } + + async fn run_sync_login(&self, settings: &Settings, store: &SqliteStore) -> Result<()> { + // TODO(ellie): Replace this with a call to atuin_client::login::login + // The reason I haven't done this yet is that this implementation allows for + // an empty key. This will use an existing key file. + // + // I'd quite like to ditch that behaviour, so have not brought it into the library + // function. + let username = or_user_input(self.username.clone(), "username"); + let password = self.password.clone().unwrap_or_else(read_user_password); + + self.prompt_and_store_key(settings, store).await?; + + let session = api_client::login( + settings.sync_address.as_str(), + LoginRequest { username, password }, + ) + .await?; + + Settings::meta_store() + .await? + .save_session(&session.session) + .await?; + + println!("Logged in!"); + + Ok(()) + } + + async fn prompt_and_store_key(&self, settings: &Settings, store: &SqliteStore) -> Result<()> { + let key_path = settings.key_path.as_str(); + let key_path = PathBuf::from(key_path); + + println!("IMPORTANT"); + println!( + "If you are already logged in on another machine, you must ensure that the key you use here is the same as the key you used there." + ); + println!("You can find your key by running 'atuin key' on the other machine."); + println!("Do not share this key with anyone."); + println!("\nRead more here: https://docs.atuin.sh/guide/sync/#login \n"); + + let key = or_user_input( + self.key.clone(), + "encryption key [blank to use existing key file]", + ); + + // if provided, the key may be EITHER base64, or a bip mnemonic + // try to normalize on base64 + let key = if key.is_empty() { + key + } else { + // try parse the key as a mnemonic... + match bip39::Mnemonic::from_phrase(&key, bip39::Language::English) { + Ok(mnemonic) => encode_key(Key::from_slice(mnemonic.entropy()))?, + Err(err) => { + match err { + // assume they copied in the base64 key + bip39::ErrorKind::InvalidWord(_) => key, + bip39::ErrorKind::InvalidChecksum => { + bail!("Key mnemonic is not valid") + } + bip39::ErrorKind::InvalidKeysize(_) + | bip39::ErrorKind::InvalidWordLength(_) + | bip39::ErrorKind::InvalidEntropyLength(_, _) => { + bail!("Key is not the correct length") + } + } + } + } + }; + + if key.is_empty() { + if key_path.exists() { + let bytes = fs_err::read_to_string(&key_path).context(format!( + "Existing key file at '{}' could not be read", + key_path.to_string_lossy() + ))?; + if decode_key(bytes).is_err() { + bail!(format!( + "The key in existing key file at '{}' is invalid", + key_path.to_string_lossy() + )); + } + } else { + panic!( + "No key provided and no existing key file found. Please use 'atuin key' on your other machine, or recover your key from a backup" + ) + } + } else if !key_path.exists() { + if decode_key(key.clone()).is_err() { + bail!("The specified key is invalid"); + } + + let mut file = File::create(&key_path).await?; + file.write_all(key.as_bytes()).await?; + } else { + // we now know that the user has logged in specifying a key, AND that the key path + // exists + + // 1. check if the saved key and the provided key match. if so, nothing to do. + // 2. if not, re-encrypt the local history and overwrite the key + let current_key: [u8; 32] = load_key(settings)?.into(); + + let encoded = key.clone(); // gonna want to save it in a bit + let new_key: [u8; 32] = decode_key(key) + .context("Could not decode provided key; is not valid base64-encoded key")? + .into(); + + if new_key != current_key { + println!("\nRe-encrypting local store with new key"); + + store.re_encrypt(¤t_key, &new_key).await?; + + println!("Writing new key"); + let mut file = File::create(&key_path).await?; + file.write_all(encoded.as_bytes()).await?; + } + } + + Ok(()) + } +} + +pub(super) fn or_user_input(value: Option<String>, name: &'static str) -> String { + value.unwrap_or_else(|| read_user_input(name)) +} + +pub(super) fn read_user_password() -> String { + let password = prompt_password("Please enter password: "); + password.expect("Failed to read from input") +} + +fn read_user_input(name: &'static str) -> String { + eprint!("Please enter {name}: "); + get_input().expect("Failed to read from input") +} + +#[cfg(test)] +mod tests { + use atuin_client::encryption::Key; + + #[test] + fn mnemonic_round_trip() { + let key = Key::from([ + 3, 1, 4, 1, 5, 9, 2, 6, 5, 3, 5, 8, 9, 7, 9, 3, 2, 3, 8, 4, 6, 2, 6, 4, 3, 3, 8, 3, 2, + 7, 9, 5, + ]); + let phrase = bip39::Mnemonic::from_entropy(&key, bip39::Language::English) + .unwrap() + .into_phrase(); + let mnemonic = bip39::Mnemonic::from_phrase(&phrase, bip39::Language::English).unwrap(); + assert_eq!(mnemonic.entropy(), key.as_slice()); + assert_eq!( + phrase, + "adapt amused able anxiety mother adapt beef gaze amount else seat alcohol cage lottery avoid scare alcohol cactus school avoid coral adjust catch pink" + ); + } +} diff --git a/crates/atuin/src/command/client/account/logout.rs b/crates/atuin/src/command/client/account/logout.rs new file mode 100644 index 00000000..b958e65a --- /dev/null +++ b/crates/atuin/src/command/client/account/logout.rs @@ -0,0 +1,5 @@ +use eyre::Result; + +pub async fn run() -> Result<()> { + atuin_client::logout::logout().await +} diff --git a/crates/atuin/src/command/client/account/register.rs b/crates/atuin/src/command/client/account/register.rs new file mode 100644 index 00000000..a2f4edfd --- /dev/null +++ b/crates/atuin/src/command/client/account/register.rs @@ -0,0 +1,87 @@ +use clap::Parser; +use eyre::{Result, bail}; + +use super::login::or_user_input; +use atuin_client::{api_client, record::sqlite_store::SqliteStore, settings::Settings}; + +#[derive(Parser, Debug)] +pub struct Cmd { + #[clap(long, short)] + pub username: Option<String>, + + #[clap(long, short)] + pub password: Option<String>, + + #[clap(long, short)] + pub email: Option<String>, +} + +impl Cmd { + pub async fn run(self, settings: &Settings, store: &SqliteStore) -> Result<()> { + run(settings, store, self.username, self.email, self.password).await + } +} + +pub async fn run( + settings: &Settings, + store: &SqliteStore, + username: Option<String>, + email: Option<String>, + password: Option<String>, +) -> Result<()> { + if let Some(_endpoint) = settings.active_hub_endpoint() { + if settings.hub_session_token().await.is_ok() { + println!("You are already authenticated with Atuin Hub."); + println!("Run 'atuin logout' to log out."); + return Ok(()); + } + + // Login can also handle registration, as the registration piece for Hub auth lives on the server + // (e.g. create a new Hub account, then log in as normal) + super::login::Cmd { + username: None, + password: None, + key: None, + from_registration: true, + } + .run(settings, store) + .await?; + return Ok(()); + } + + if settings.session_token().await.is_ok() { + println!("You are already logged in."); + println!("Run 'atuin logout' to log out."); + return Ok(()); + } + + println!("Registering for an Atuin Sync account"); + + let username = or_user_input(username, "username"); + let email = or_user_input(email, "email"); + + let password = password + .clone() + .unwrap_or_else(super::login::read_user_password); + + if password.is_empty() { + bail!("please provide a password"); + } + + let session = + api_client::register(settings.sync_address.as_str(), &username, &email, &password).await?; + + let meta = Settings::meta_store().await?; + meta.save_session(&session.session).await?; + + let _key = atuin_client::encryption::load_key(settings)?; + + println!( + "Registration successful! Please make a note of your key (run 'atuin key') and keep it safe." + ); + println!( + "You will need it to log in on other devices, and we cannot help recover it if you lose it." + ); + + Ok(()) +} diff --git a/crates/atuin/src/command/client/daemon.rs b/crates/atuin/src/command/client/daemon.rs new file mode 100644 index 00000000..64547505 --- /dev/null +++ b/crates/atuin/src/command/client/daemon.rs @@ -0,0 +1,714 @@ +use std::fs::{self, File, OpenOptions}; +use std::io::{ErrorKind, Write}; +#[cfg(unix)] +use std::os::unix::net::UnixStream as StdUnixStream; +use std::path::{Path, PathBuf}; +use std::process::{Command, Stdio}; +use std::time::{Duration, Instant}; + +use atuin_client::{ + database::Sqlite, history::History, record::sqlite_store::SqliteStore, settings::Settings, +}; +use atuin_daemon::client::{DaemonClientErrorKind, HistoryClient, classify_error}; +use clap::Subcommand; +#[cfg(unix)] +use daemonize::Daemonize; +use eyre::{Result, WrapErr, bail, eyre}; +use fs4::fs_std::FileExt; +use tokio::time::sleep; + +#[derive(clap::Args, Debug)] +pub struct Cmd { + /// Internal flag for daemonization + #[arg(long, hide = true)] + daemonize: bool, + + /// Also write daemon logs to the console (useful for debugging) + #[arg(long)] + show_logs: bool, + + #[command(subcommand)] + subcmd: Option<SubCmd>, +} + +#[derive(Subcommand, Debug)] +#[command(infer_subcommands = true)] +pub enum SubCmd { + /// Start the daemon server + Start { + #[arg(long, hide = true)] + daemonize: bool, + + /// Also write daemon logs to the console (useful for debugging) + #[arg(long)] + show_logs: bool, + + /// Force start: kill existing daemon process and reset the socket + #[arg(long)] + force: bool, + }, + + /// Show the daemon's current status + Status, + + /// Stop the daemon gracefully + Stop, + + /// Restart the daemon (stop, then start in background) + Restart, +} + +impl Cmd { + /// Returns `true` when the process should daemonize before creating the + /// async runtime or opening any database connections. + #[cfg(unix)] + pub fn should_daemonize(&self) -> bool { + match &self.subcmd { + Some(SubCmd::Start { daemonize, .. }) => *daemonize, + None => self.daemonize, + _ => false, + } + } + + /// Returns `true` when logs should also be written to the console. + pub fn show_logs(&self) -> bool { + match &self.subcmd { + Some(SubCmd::Start { show_logs, .. }) => *show_logs, + None => self.show_logs, + _ => false, + } + } + + pub async fn run( + self, + settings: Settings, + store: SqliteStore, + history_db: Sqlite, + ) -> Result<()> { + match self.subcmd { + None => { + eprintln!("Warning: `atuin daemon` is deprecated, use `atuin daemon start`"); + run(settings, store, history_db, false).await + } + Some(SubCmd::Start { force, .. }) => run(settings, store, history_db, force).await, + Some(SubCmd::Status) => status_cmd(&settings).await, + Some(SubCmd::Stop) => stop_cmd(&settings).await, + Some(SubCmd::Restart) => restart_cmd(&settings).await, + } + } +} + +const DAEMON_VERSION: &str = env!("CARGO_PKG_VERSION"); +const DAEMON_PROTOCOL_VERSION: u32 = 1; +const STARTUP_POLL: Duration = Duration::from_millis(40); +const LOCK_POLL: Duration = Duration::from_millis(20); +const LEGACY_DAEMON_RESTART_MESSAGE: &str = "legacy daemon detected; restart daemon manually"; + +struct PidfileGuard { + file: File, +} + +impl PidfileGuard { + fn acquire(path: &Path) -> Result<Self> { + let mut file = open_lock_file(path)?; + + if !file.try_lock_exclusive()? { + bail!( + "daemon already running (pidfile lock busy at {})", + path.display() + ); + } + + file.set_len(0) + .wrap_err_with(|| format!("could not truncate daemon pidfile {}", path.display()))?; + writeln!(file, "{}", std::process::id()) + .and_then(|()| writeln!(file, "{DAEMON_VERSION}")) + .wrap_err_with(|| format!("could not write daemon pidfile {}", path.display()))?; + + Ok(Self { file }) + } +} + +impl Drop for PidfileGuard { + fn drop(&mut self) { + let _ = self.file.unlock(); + } +} + +enum Probe { + Ready(HistoryClient), + NeedsRestart(String), + Unreachable(eyre::Report), +} + +fn daemon_matches_expected(version: &str, protocol: u32) -> bool { + version == DAEMON_VERSION && protocol == DAEMON_PROTOCOL_VERSION +} + +fn daemon_mismatch_message(version: &str, protocol: u32) -> String { + if protocol == DAEMON_PROTOCOL_VERSION { + format!("daemon is out of date: expected {DAEMON_VERSION}, got {version}") + } else { + format!("daemon protocol mismatch: expected {DAEMON_PROTOCOL_VERSION}, got {protocol}") + } +} + +fn is_legacy_daemon_error(err: &eyre::Report) -> bool { + matches!(classify_error(err), DaemonClientErrorKind::Unimplemented) +} + +fn should_retry_after_error(err: &eyre::Report) -> bool { + matches!( + classify_error(err), + DaemonClientErrorKind::Connect + | DaemonClientErrorKind::Unavailable + | DaemonClientErrorKind::Unimplemented + ) +} + +fn daemon_startup_lock_path(pidfile_path: &Path) -> PathBuf { + let mut os = pidfile_path.as_os_str().to_os_string(); + os.push(".startup.lock"); + PathBuf::from(os) +} + +fn open_lock_file(path: &Path) -> Result<File> { + if let Some(parent) = path.parent() { + fs::create_dir_all(parent) + .wrap_err_with(|| format!("could not create lock directory {}", parent.display()))?; + } + + OpenOptions::new() + .read(true) + .write(true) + .create(true) + .truncate(false) + .open(path) + .wrap_err_with(|| format!("could not open lock file {}", path.display())) +} + +async fn wait_for_lock(path: &Path, timeout: Duration) -> Result<File> { + let file = open_lock_file(path)?; + let start = Instant::now(); + + loop { + match file.try_lock_exclusive() { + Ok(true) => return Ok(file), + Ok(false) => { + if start.elapsed() >= timeout { + bail!("timed out waiting for lock at {}", path.display()); + } + + sleep(LOCK_POLL).await; + } + Err(err) => { + return Err(eyre!("could not lock {}: {err}", path.display())); + } + } + } +} + +async fn wait_for_pidfile_available(path: &Path, timeout: Duration) -> Result<()> { + let file = wait_for_lock(path, timeout).await?; + file.unlock() + .wrap_err_with(|| format!("failed to unlock {}", path.display()))?; + Ok(()) +} + +async fn connect_client(settings: &Settings) -> Result<HistoryClient> { + HistoryClient::new( + #[cfg(not(unix))] + settings.daemon.tcp_port, + #[cfg(unix)] + settings.daemon.socket_path.clone(), + ) + .await +} + +async fn probe(settings: &Settings) -> Probe { + let mut client = match connect_client(settings).await { + Ok(client) => client, + Err(err) => return Probe::Unreachable(err), + }; + + match client.status().await { + Ok(status) => { + if daemon_matches_expected(&status.version, status.protocol) { + Probe::Ready(client) + } else { + Probe::NeedsRestart(daemon_mismatch_message(&status.version, status.protocol)) + } + } + Err(err) => Probe::Unreachable(err), + } +} + +async fn request_shutdown(settings: &Settings) { + if let Ok(mut client) = connect_client(settings).await { + let _ = client.shutdown().await; + } +} + +fn spawn_daemon_process() -> Result<()> { + let exe = std::env::current_exe().wrap_err("could not locate atuin executable")?; + + let mut cmd = Command::new(exe); + cmd.arg("daemon") + .arg("start") + .stdin(Stdio::null()) + .stdout(Stdio::null()) + .stderr(Stdio::null()); + + #[cfg(unix)] + cmd.arg("--daemonize"); + + cmd.spawn().wrap_err("failed to spawn daemon process")?; + + Ok(()) +} + +fn startup_timeout(settings: &Settings) -> Duration { + Duration::from_secs_f64(settings.local_timeout.max(0.5) + 2.0) +} + +#[cfg(unix)] +fn remove_stale_socket_if_present(settings: &Settings) -> Result<()> { + if settings.daemon.systemd_socket { + return Ok(()); + } + + let socket_path = Path::new(&settings.daemon.socket_path); + if !socket_path.exists() { + return Ok(()); + } + + match StdUnixStream::connect(socket_path) { + Ok(stream) => { + drop(stream); + Ok(()) + } + Err(err) if err.kind() == ErrorKind::ConnectionRefused => { + fs::remove_file(socket_path).wrap_err_with(|| { + format!( + "failed to remove stale daemon socket {}", + socket_path.display() + ) + })?; + Ok(()) + } + Err(err) if err.kind() == ErrorKind::NotFound => Ok(()), + Err(_) => Ok(()), + } +} + +async fn wait_until_ready(settings: &Settings, timeout: Duration) -> Result<HistoryClient> { + let start = Instant::now(); + let mut last_error = eyre!("daemon did not become ready"); + + loop { + match probe(settings).await { + Probe::Ready(client) => return Ok(client), + Probe::NeedsRestart(reason) => { + last_error = eyre!(reason); + } + Probe::Unreachable(err) => { + if is_legacy_daemon_error(&err) { + return Err(err.wrap_err(LEGACY_DAEMON_RESTART_MESSAGE)); + } + last_error = err; + } + } + + if start.elapsed() >= timeout { + return Err(last_error.wrap_err(format!( + "timed out waiting for daemon startup after {}ms", + timeout.as_millis() + ))); + } + + sleep(STARTUP_POLL).await; + } +} + +fn ensure_autostart_supported(settings: &Settings) -> Result<()> { + #[cfg(unix)] + if settings.daemon.systemd_socket { + bail!( + "daemon autostart is incompatible with `daemon.systemd_socket = true`; use systemd to manage the daemon" + ); + } + #[cfg(not(unix))] + let _ = settings; + + Ok(()) +} + +async fn restart_daemon(settings: &Settings) -> Result<HistoryClient> { + ensure_autostart_supported(settings)?; + + let timeout = startup_timeout(settings); + let pidfile_path = PathBuf::from(&settings.daemon.pidfile_path); + let startup_lock_path = daemon_startup_lock_path(&pidfile_path); + let startup_lock = wait_for_lock(&startup_lock_path, timeout).await?; + + match probe(settings).await { + Probe::Ready(client) => { + drop(startup_lock); + return Ok(client); + } + Probe::NeedsRestart(_) => { + request_shutdown(settings).await; + } + Probe::Unreachable(err) => { + if is_legacy_daemon_error(&err) { + return Err(err.wrap_err(LEGACY_DAEMON_RESTART_MESSAGE)); + } + } + } + + // This prevents rapid-fire hook invocations from racing daemon restart. + wait_for_pidfile_available(&pidfile_path, timeout).await?; + + #[cfg(unix)] + remove_stale_socket_if_present(settings)?; + + spawn_daemon_process()?; + let client = wait_until_ready(settings, timeout).await?; + + drop(startup_lock); + Ok(client) +} + +fn ensure_reply_compatible(settings: &Settings, version: &str, protocol: u32) -> Result<()> { + if daemon_matches_expected(version, protocol) { + return Ok(()); + } + + let message = daemon_mismatch_message(version, protocol); + if settings.daemon.autostart { + bail!("{message}"); + } + + bail!("{message}. Enable `daemon.autostart = true` or restart the daemon manually"); +} + +pub async fn start_history(settings: &Settings, history: History) -> Result<String> { + match async { + connect_client(settings) + .await? + .start_history(history.clone()) + .await + } + .await + { + Ok(resp) => { + if daemon_matches_expected(&resp.version, resp.protocol) { + return Ok(resp.id); + } + + if !settings.daemon.autostart { + return Err(eyre!( + "{}. Enable `daemon.autostart = true` or restart the daemon manually", + daemon_mismatch_message(&resp.version, resp.protocol) + )); + } + } + Err(err) if !settings.daemon.autostart => return Err(err), + Err(err) if !should_retry_after_error(&err) => return Err(err), + Err(_) => {} + } + + let resp = restart_daemon(settings) + .await? + .start_history(history) + .await?; + ensure_reply_compatible(settings, &resp.version, resp.protocol)?; + Ok(resp.id) +} + +pub async fn end_history(settings: &Settings, id: String, duration: u64, exit: i64) -> Result<()> { + match async { + connect_client(settings) + .await? + .end_history(id.clone(), duration, exit) + .await + } + .await + { + Ok(resp) => { + if daemon_matches_expected(&resp.version, resp.protocol) { + return Ok(()); + } + + if !settings.daemon.autostart { + return Err(eyre!( + "{}. Enable `daemon.autostart = true` or restart the daemon manually", + daemon_mismatch_message(&resp.version, resp.protocol) + )); + } + + // End succeeded on the running daemon, so avoid replaying it. + // We only restart to make subsequent hook calls target the expected version. + let _ = restart_daemon(settings).await; + return Ok(()); + } + Err(err) if !settings.daemon.autostart => return Err(err), + Err(err) if !should_retry_after_error(&err) => return Err(err), + Err(_) => {} + } + + let resp = restart_daemon(settings) + .await? + .end_history(id, duration, exit) + .await?; + ensure_reply_compatible(settings, &resp.version, resp.protocol)?; + Ok(()) +} + +async fn status_cmd(settings: &Settings) -> Result<()> { + match probe(settings).await { + Probe::Ready(mut client) => { + let status = client.status().await?; + println!("Daemon running"); + println!(" PID: {}", status.pid); + println!(" Version: {}", status.version); + println!(" Protocol: {}", status.protocol); + println!(" Healthy: {}", status.healthy); + #[cfg(unix)] + println!(" Socket: {}", settings.daemon.socket_path); + #[cfg(not(unix))] + println!(" Port: {}", settings.daemon.tcp_port); + } + Probe::NeedsRestart(reason) => { + println!("Daemon running (needs restart)"); + println!(" Reason: {reason}"); + } + Probe::Unreachable(_) => { + println!("Daemon is not running"); + } + } + + Ok(()) +} + +async fn stop_cmd(settings: &Settings) -> Result<()> { + let Ok(mut client) = connect_client(settings).await else { + println!("Daemon is not running"); + return Ok(()); + }; + + match client.shutdown().await { + Ok(true) => { + println!("Shutdown requested"); + + let pidfile_path = PathBuf::from(&settings.daemon.pidfile_path); + let timeout = Duration::from_secs(5); + match wait_for_pidfile_available(&pidfile_path, timeout).await { + Ok(()) => println!("Daemon stopped"), + Err(_) => println!("Daemon may still be shutting down"), + } + + Ok(()) + } + Ok(false) => bail!("Daemon rejected shutdown request"), + Err(err) => Err(err.wrap_err("Failed to send shutdown request")), + } +} + +async fn restart_cmd(settings: &Settings) -> Result<()> { + // Stop if running + match probe(settings).await { + Probe::Ready(_) | Probe::NeedsRestart(_) => { + request_shutdown(settings).await; + println!("Stopping daemon..."); + + let pidfile_path = PathBuf::from(&settings.daemon.pidfile_path); + let timeout = Duration::from_secs(5); + wait_for_pidfile_available(&pidfile_path, timeout) + .await + .wrap_err("Timed out waiting for old daemon to stop")?; + } + Probe::Unreachable(_) => { + println!("No daemon running"); + } + } + + #[cfg(unix)] + remove_stale_socket_if_present(settings)?; + + spawn_daemon_process()?; + println!("Starting daemon..."); + + let timeout = startup_timeout(settings); + let status = wait_until_ready(settings, timeout).await?.status().await?; + + println!("Daemon restarted"); + println!(" PID: {}", status.pid); + println!(" Version: {}", status.version); + + Ok(()) +} + +/// Daemonize the current process. Must be called before creating the tokio +/// runtime or opening database connections, since `fork()` inside an async +/// runtime corrupts its internal state. +#[cfg(unix)] +pub fn daemonize_current_process() -> Result<()> { + let cwd = + std::env::current_dir().wrap_err("could not determine current directory for daemon")?; + + Daemonize::new() + .working_directory(cwd) + .start() + .wrap_err("failed to daemonize process")?; + + Ok(()) +} + +async fn run( + settings: Settings, + store: SqliteStore, + history_db: Sqlite, + force: bool, +) -> Result<()> { + if force { + force_cleanup(&settings); + } + + let pidfile_path = PathBuf::from(&settings.daemon.pidfile_path); + let _pidfile_guard = PidfileGuard::acquire(&pidfile_path)?; + + atuin_daemon::boot(settings, store, history_db).await?; + + Ok(()) +} + +/// Force cleanup: kill existing daemon process and remove socket. +fn force_cleanup(settings: &Settings) { + let pidfile_path = Path::new(&settings.daemon.pidfile_path); + + // Read and kill the existing process if pidfile exists + if pidfile_path.exists() { + if let Ok(contents) = fs::read_to_string(pidfile_path) + && let Some(pid_str) = contents.lines().next() + && let Ok(pid) = pid_str.parse::<u32>() + { + kill_process(pid); + // Give it a moment to release resources + std::thread::sleep(Duration::from_millis(100)); + } + + // Remove the pidfile + if let Err(e) = fs::remove_file(pidfile_path) + && e.kind() != ErrorKind::NotFound + { + tracing::warn!("failed to remove pidfile: {e}"); + } + } + + // Remove the socket file + #[cfg(unix)] + { + let socket_path = Path::new(&settings.daemon.socket_path); + if socket_path.exists() + && let Err(e) = fs::remove_file(socket_path) + && e.kind() != ErrorKind::NotFound + { + tracing::warn!("failed to remove socket: {e}"); + } + } +} + +/// Kill a process by PID. +#[cfg(unix)] +fn kill_process(pid: u32) { + // Use kill command to send SIGTERM for graceful shutdown + let _ = Command::new("kill") + .args(["-TERM", &pid.to_string()]) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .status(); +} + +/// Kill a process by PID. +#[cfg(not(unix))] +fn kill_process(pid: u32) { + // On Windows, use taskkill + let _ = Command::new("taskkill") + .args(["/PID", &pid.to_string(), "/F"]) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .status(); +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_version_matches() { + assert!(daemon_matches_expected( + DAEMON_VERSION, + DAEMON_PROTOCOL_VERSION + )); + } + + #[test] + fn test_version_mismatch() { + assert!(!daemon_matches_expected("0.0.0", DAEMON_PROTOCOL_VERSION)); + assert!(!daemon_matches_expected(DAEMON_VERSION, 999)); + assert!(!daemon_matches_expected("0.0.0", 999)); + } + + #[test] + fn test_mismatch_message_version() { + let msg = daemon_mismatch_message("0.0.0", DAEMON_PROTOCOL_VERSION); + assert!(msg.contains("out of date"), "got: {msg}"); + assert!(msg.contains("0.0.0")); + assert!(msg.contains(DAEMON_VERSION)); + } + + #[test] + fn test_mismatch_message_protocol() { + let msg = daemon_mismatch_message(DAEMON_VERSION, 999); + assert!(msg.contains("protocol mismatch"), "got: {msg}"); + } + + #[test] + fn test_startup_lock_path() { + let pidfile = Path::new("/tmp/atuin-daemon.pid"); + let lock = daemon_startup_lock_path(pidfile); + assert_eq!(lock, PathBuf::from("/tmp/atuin-daemon.pid.startup.lock")); + } + + #[test] + fn test_pidfile_guard_acquire_and_drop() { + let tmp = tempfile::tempdir().unwrap(); + let pidfile = tmp.path().join("daemon.pid"); + + { + let _guard = PidfileGuard::acquire(&pidfile).unwrap(); + // Guard holds an exclusive lock — on Windows other handles cannot + // read the file, so we verify contents after the guard is dropped. + } + + let contents = std::fs::read_to_string(&pidfile).unwrap(); + let lines: Vec<&str> = contents.lines().collect(); + assert_eq!(lines.len(), 2); + assert_eq!(lines[0], std::process::id().to_string()); + assert_eq!(lines[1], DAEMON_VERSION); + + // After guard is dropped, lock should be released — acquiring again must succeed. + let _guard2 = PidfileGuard::acquire(&pidfile).unwrap(); + } + + #[test] + fn test_pidfile_guard_prevents_double_acquire() { + let tmp = tempfile::tempdir().unwrap(); + let pidfile = tmp.path().join("daemon.pid"); + + let _guard = PidfileGuard::acquire(&pidfile).unwrap(); + let result = PidfileGuard::acquire(&pidfile); + assert!(result.is_err()); + } +} diff --git a/crates/atuin/src/command/client/default_config.rs b/crates/atuin/src/command/client/default_config.rs new file mode 100644 index 00000000..f51e45c2 --- /dev/null +++ b/crates/atuin/src/command/client/default_config.rs @@ -0,0 +1,5 @@ +use atuin_client::settings::Settings; + +pub fn run() { + println!("{}", Settings::example_config()); +} diff --git a/crates/atuin/src/command/client/doctor.rs b/crates/atuin/src/command/client/doctor.rs new file mode 100644 index 00000000..c2c47a58 --- /dev/null +++ b/crates/atuin/src/command/client/doctor.rs @@ -0,0 +1,401 @@ +use std::process::Command; +use std::{env, str::FromStr}; + +use atuin_client::database::Sqlite; +use atuin_client::settings::Settings; +use atuin_common::shell::{Shell, shell_name}; +use atuin_common::utils; +use colored::Colorize; +use eyre::Result; +use serde::Serialize; + +use sysinfo::{Disks, System, get_current_pid}; + +#[derive(Debug, Serialize)] +struct ShellInfo { + pub name: String, + + // best-effort, not supported on all OSes + pub default: String, + + // Detect some shell plugins that the user has installed. + // I'm just going to start with preexec/blesh + pub plugins: Vec<String>, + + // The preexec framework used in the current session, if Atuin is loaded. + pub preexec: Option<String>, +} + +impl ShellInfo { + // HACK ALERT! + // Many of the shell vars we need to detect are not exported :( + // So, we're going to run a interactive session and directly check the + // variable. There's a chance this won't work, so it should not be fatal. + // + // Every shell we support handles `shell -ic 'command'` + fn shellvar_exists(shell: &str, var: &str) -> bool { + let cmd = Command::new(shell) + .args([ + "-ic", + format!("[ -z ${var} ] || echo ATUIN_DOCTOR_ENV_FOUND").as_str(), + ]) + .output() + .map_or(String::new(), |v| { + let out = v.stdout; + String::from_utf8(out).unwrap_or_default() + }); + + cmd.contains("ATUIN_DOCTOR_ENV_FOUND") + } + + fn detect_preexec_framework(shell: &str) -> Option<String> { + if env::var("ATUIN_SESSION").ok().is_none() { + None + } else if shell.starts_with("bash") || shell == "sh" { + env::var("ATUIN_PREEXEC_BACKEND") + .ok() + .filter(|value| !value.is_empty()) + .and_then(|atuin_preexec_backend| { + atuin_preexec_backend.rfind(':').and_then(|pos_colon| { + u32::from_str(&atuin_preexec_backend[..pos_colon]) + .ok() + .is_some_and(|preexec_shlvl| { + env::var("SHLVL") + .ok() + .and_then(|shlvl| u32::from_str(&shlvl).ok()) + .is_some_and(|shlvl| shlvl == preexec_shlvl) + }) + .then(|| atuin_preexec_backend[pos_colon + 1..].to_string()) + }) + }) + } else { + Some("built-in".to_string()) + } + } + + fn validate_plugin_blesh( + _shell: &str, + shell_process: &sysinfo::Process, + ble_session_id: &str, + ) -> Option<String> { + ble_session_id + .split('/') + .nth(1) + .and_then(|field| u32::from_str(field).ok()) + .filter(|&blesh_pid| blesh_pid == shell_process.pid().as_u32()) + .map(|_| "blesh".to_string()) + } + + pub fn plugins(shell: &str, shell_process: &sysinfo::Process) -> Vec<String> { + // consider a different detection approach if there are plugins + // that don't set shell vars + + enum PluginShellType { + Any, + Bash, + + // Note: these are currently unused + #[allow(dead_code)] + Zsh, + #[allow(dead_code)] + Fish, + #[allow(dead_code)] + Nushell, + #[allow(dead_code)] + Xonsh, + } + + enum PluginProbeType { + EnvironmentVariable(&'static str), + InteractiveShellVariable(&'static str), + } + + type PluginValidator = fn(&str, &sysinfo::Process, &str) -> Option<String>; + + let plugin_list: [( + &str, + PluginShellType, + PluginProbeType, + Option<PluginValidator>, + ); 3] = [ + ( + "atuin", + PluginShellType::Any, + PluginProbeType::EnvironmentVariable("ATUIN_SESSION"), + None, + ), + ( + "blesh", + PluginShellType::Bash, + PluginProbeType::EnvironmentVariable("BLE_SESSION_ID"), + Some(Self::validate_plugin_blesh), + ), + ( + "bash-preexec", + PluginShellType::Bash, + PluginProbeType::InteractiveShellVariable("bash_preexec_imported"), + None, + ), + ]; + + plugin_list + .into_iter() + .filter(|(_, shell_type, _, _)| match shell_type { + PluginShellType::Any => true, + PluginShellType::Bash => shell.starts_with("bash") || shell == "sh", + PluginShellType::Zsh => shell.starts_with("zsh"), + PluginShellType::Fish => shell.starts_with("fish"), + PluginShellType::Nushell => shell.starts_with("nu"), + PluginShellType::Xonsh => shell.starts_with("xonsh"), + }) + .filter_map(|(plugin, _, probe_type, validator)| -> Option<String> { + match probe_type { + PluginProbeType::EnvironmentVariable(env) => { + env::var(env).ok().filter(|value| !value.is_empty()) + } + PluginProbeType::InteractiveShellVariable(shellvar) => { + ShellInfo::shellvar_exists(shell, shellvar).then_some(String::default()) + } + } + .and_then(|value| { + validator.map_or_else( + || Some(plugin.to_string()), + |validator| validator(shell, shell_process, &value), + ) + }) + }) + .collect() + } + + pub fn new() -> Self { + // TODO: rework to use atuin_common::Shell + + let sys = System::new_all(); + + let process = sys + .process(get_current_pid().expect("Failed to get current PID")) + .expect("Process with current pid does not exist"); + + let parent = sys + .process(process.parent().expect("Atuin running with no parent!")) + .expect("Process with parent pid does not exist"); + + let name = shell_name(Some(parent)); + + let plugins = ShellInfo::plugins(name.as_str(), parent); + + let default = Shell::default_shell().unwrap_or(Shell::Unknown).to_string(); + + let preexec = Self::detect_preexec_framework(name.as_str()); + + Self { + name, + default, + plugins, + preexec, + } + } +} + +#[derive(Debug, Serialize)] +struct DiskInfo { + pub name: String, + pub filesystem: String, +} + +#[derive(Debug, Serialize)] +struct SystemInfo { + pub os: String, + + pub arch: String, + + pub version: String, + pub disks: Vec<DiskInfo>, +} + +impl SystemInfo { + pub fn new() -> Self { + let disks = Disks::new_with_refreshed_list(); + let disks = disks + .list() + .iter() + .map(|d| DiskInfo { + name: d.name().to_os_string().into_string().unwrap(), + filesystem: d.file_system().to_os_string().into_string().unwrap(), + }) + .collect(); + + Self { + os: System::name().unwrap_or_else(|| "unknown".to_string()), + arch: System::cpu_arch().unwrap_or_else(|| "unknown".to_string()), + version: System::os_version().unwrap_or_else(|| "unknown".to_string()), + disks, + } + } +} + +#[derive(Debug, Serialize)] +struct SyncInfo { + /// Whether the main Atuin sync server is in use + /// I'm just calling it Atuin Cloud for lack of a better name atm + pub cloud: bool, + pub records: bool, + pub auto_sync: bool, + + pub last_sync: String, +} + +impl SyncInfo { + pub async fn new(settings: &Settings) -> Self { + Self { + cloud: settings.is_hub_sync(), + auto_sync: settings.auto_sync, + records: settings.sync.records, + last_sync: Settings::last_sync() + .await + .map_or_else(|_| "no last sync".to_string(), |v| v.to_string()), + } + } +} + +#[derive(Debug)] +struct SettingPaths { + db: String, + record_store: String, + key: String, +} + +impl SettingPaths { + pub fn new(settings: &Settings) -> Self { + Self { + db: settings.db_path.clone(), + record_store: settings.record_store_path.clone(), + key: settings.key_path.clone(), + } + } + + pub fn verify(&self) { + let paths = vec![ + ("ATUIN_DB_PATH", &self.db), + ("ATUIN_RECORD_STORE", &self.record_store), + ("ATUIN_KEY", &self.key), + ]; + + for (path_env_var, path) in paths { + if utils::broken_symlink(path) { + eprintln!( + "{path} (${path_env_var}) is a broken symlink. This may cause issues with Atuin." + ); + } + } + } +} + +#[derive(Debug, Serialize)] +struct AtuinInfo { + pub version: String, + pub commit: String, + + /// Whether the main Atuin sync server is in use + /// I'm just calling it Atuin Cloud for lack of a better name atm + pub sync: Option<SyncInfo>, + + pub sqlite_version: String, + + #[serde(skip)] // probably unnecessary to expose this + pub setting_paths: SettingPaths, +} + +impl AtuinInfo { + pub async fn new(settings: &Settings) -> Self { + let logged_in = settings.logged_in().await.unwrap_or(false); + + let sync = if logged_in { + Some(SyncInfo::new(settings).await) + } else { + None + }; + + let sqlite_version = match Sqlite::new("sqlite::memory:", 0.1).await { + Ok(db) => db + .sqlite_version() + .await + .unwrap_or_else(|_| "unknown".to_string()), + Err(_) => "error".to_string(), + }; + + Self { + version: crate::VERSION.to_string(), + commit: crate::SHA.to_string(), + sync, + sqlite_version, + setting_paths: SettingPaths::new(settings), + } + } +} + +#[derive(Debug, Serialize)] +struct DoctorDump { + pub atuin: AtuinInfo, + pub shell: ShellInfo, + pub system: SystemInfo, +} + +impl DoctorDump { + pub async fn new(settings: &Settings) -> Self { + Self { + atuin: AtuinInfo::new(settings).await, + shell: ShellInfo::new(), + system: SystemInfo::new(), + } + } +} + +fn checks(info: &DoctorDump) { + println!(); // spacing + // + let zfs_error = "[Filesystem] ZFS is known to have some issues with SQLite. Atuin uses SQLite heavily. If you are having poor performance, there are some workarounds here: https://github.com/atuinsh/atuin/issues/952".bold().red(); + let bash_plugin_error = "[Shell] If you are using Bash, Atuin requires that either bash-preexec or ble.sh (>= 0.4) be installed. An older ble.sh may not be detected. so ignore this if you have ble.sh >= 0.4 set up! Read more here: https://docs.atuin.sh/guide/installation/#bash".bold().red(); + let blesh_integration_error = "[Shell] Atuin and ble.sh seem to be loaded in the session, but the integration does not seem to be working. Please check the setup in .bashrc.".bold().red(); + + // ZFS: https://github.com/atuinsh/atuin/issues/952 + if info.system.disks.iter().any(|d| d.filesystem == "zfs") { + println!("{zfs_error}"); + } + + info.atuin.setting_paths.verify(); + + // Shell + if info.shell.name == "bash" { + if !info + .shell + .plugins + .iter() + .any(|p| p == "blesh" || p == "bash-preexec") + { + println!("{bash_plugin_error}"); + } + + if info.shell.plugins.iter().any(|plugin| plugin == "atuin") + && info.shell.plugins.iter().any(|plugin| plugin == "blesh") + && info.shell.preexec.as_ref().is_some_and(|val| val == "none") + { + println!("{blesh_integration_error}"); + } + } +} + +pub async fn run(settings: &Settings) -> Result<()> { + println!("{}", "Atuin Doctor".bold()); + println!("Checking for diagnostics"); + let dump = DoctorDump::new(settings).await; + + checks(&dump); + + let dump = serde_json::to_string_pretty(&dump)?; + + println!("\nPlease include the output below with any bug reports or issues\n"); + println!("{dump}"); + + Ok(()) +} diff --git a/crates/atuin/src/command/client/dotfiles.rs b/crates/atuin/src/command/client/dotfiles.rs new file mode 100644 index 00000000..f42b18f2 --- /dev/null +++ b/crates/atuin/src/command/client/dotfiles.rs @@ -0,0 +1,28 @@ +use clap::Subcommand; +use eyre::Result; + +use atuin_client::{record::sqlite_store::SqliteStore, settings::Settings}; + +mod alias; +mod var; + +#[derive(Subcommand, Debug)] +#[command(infer_subcommands = true)] +pub enum Cmd { + /// Manage shell aliases with Atuin + #[command(subcommand)] + Alias(alias::Cmd), + + /// Manage shell and environment variables with Atuin + #[command(subcommand)] + Var(var::Cmd), +} + +impl Cmd { + pub async fn run(self, settings: &Settings, store: SqliteStore) -> Result<()> { + match self { + Self::Alias(cmd) => cmd.run(settings, store).await, + Self::Var(cmd) => cmd.run(settings, store).await, + } + } +} diff --git a/crates/atuin/src/command/client/dotfiles/alias.rs b/crates/atuin/src/command/client/dotfiles/alias.rs new file mode 100644 index 00000000..983c67f1 --- /dev/null +++ b/crates/atuin/src/command/client/dotfiles/alias.rs @@ -0,0 +1,187 @@ +use clap::{Subcommand, ValueEnum}; +use eyre::{Context, Result, eyre}; + +use atuin_client::{encryption, record::sqlite_store::SqliteStore, settings::Settings}; + +use atuin_dotfiles::{shell::Alias, store::AliasStore}; + +#[derive(Clone, Copy, Debug, Default, ValueEnum)] +pub enum SortBy { + /// Sort by alias name + #[default] + Name, + /// Sort by alias value + Value, +} + +#[derive(Subcommand, Debug)] +#[command(infer_subcommands = true)] +pub enum Cmd { + /// Set an alias + Set { name: String, value: String }, + + /// Delete an alias + Delete { name: String }, + + /// List all aliases + List { + /// Sort results by field + #[arg(long, value_enum, default_value_t = SortBy::Name)] + sort_by: SortBy, + + /// Sort in reverse (descending) order + #[arg(long, short)] + reverse: bool, + + /// Filter aliases by name (substring match) + #[arg(long, short)] + name: Option<String>, + + /// Filter aliases by value (substring match) + #[arg(long, short)] + value: Option<String>, + }, + + /// Delete all aliases + Clear, + // There are too many edge cases to parse at the moment. Disable for now. + // Import, +} + +impl Cmd { + async fn set(&self, store: &AliasStore, name: String, value: String) -> Result<()> { + let illegal_char = regex::Regex::new("[ \t\n&();<>|\\\"'`$/]").unwrap(); + if illegal_char.is_match(name.as_str()) { + return Err(eyre!("Illegal character in alias name")); + } + + let aliases = store.aliases().await?; + let found: Vec<Alias> = aliases.into_iter().filter(|a| a.name == name).collect(); + + if found.is_empty() { + println!("Aliasing '{name}={value}'."); + } else { + println!( + "Overwriting alias '{name}={}' with '{name}={value}'.", + found[0].value + ); + } + + store.set(&name, &value).await?; + + Ok(()) + } + + async fn list( + &self, + store: &AliasStore, + sort_by: SortBy, + reverse: bool, + name_filter: Option<String>, + value_filter: Option<String>, + ) -> Result<()> { + let mut aliases = store.aliases().await?; + + // Apply filters + if let Some(ref name_pattern) = name_filter { + let pattern = name_pattern.to_lowercase(); + aliases.retain(|a| a.name.to_lowercase().contains(&pattern)); + } + if let Some(ref value_pattern) = value_filter { + let pattern = value_pattern.to_lowercase(); + aliases.retain(|a| a.value.to_lowercase().contains(&pattern)); + } + + // Apply sorting + match sort_by { + SortBy::Name => { + aliases.sort_by(|a, b| a.name.to_lowercase().cmp(&b.name.to_lowercase())); + } + SortBy::Value => { + aliases.sort_by(|a, b| a.value.to_lowercase().cmp(&b.value.to_lowercase())); + } + } + + // Apply reverse if requested + if reverse { + aliases.reverse(); + } + + for i in aliases { + println!("{}={}", i.name, i.value); + } + + Ok(()) + } + + async fn clear(&self, store: &AliasStore) -> Result<()> { + let aliases = store.aliases().await?; + + for i in aliases { + self.delete(store, i.name).await?; + } + + Ok(()) + } + + async fn delete(&self, store: &AliasStore, name: String) -> Result<()> { + let mut aliases = store.aliases().await?.into_iter(); + if let Some(alias) = aliases.find(|alias| alias.name == name) { + println!("Deleting '{name}={}'.", alias.value); + store.delete(&name).await?; + } else { + eprintln!("Cannot delete '{name}': Alias not set."); + } + Ok(()) + } + + /* + async fn import(&self, store: &AliasStore) -> Result<()> { + let aliases = atuin_dotfiles::shell::import_aliases(store).await?; + + for i in aliases { + println!("Importing {}={}", i.name, i.value); + } + + Ok(()) + } + */ + + pub async fn run(&self, settings: &Settings, store: SqliteStore) -> Result<()> { + if !settings.dotfiles.enabled { + eprintln!( + "Dotfiles are not enabled. Add\n\n[dotfiles]\nenabled = true\n\nto your configuration file to enable them.\n" + ); + eprintln!("The default configuration file is located at ~/.config/atuin/config.toml."); + return Ok(()); + } + + let encryption_key: [u8; 32] = encryption::load_key(settings) + .context("could not load encryption key")? + .into(); + let host_id = Settings::host_id().await?; + + let alias_store = AliasStore::new(store, host_id, encryption_key); + + match self { + Self::Set { name, value } => self.set(&alias_store, name.clone(), value.clone()).await, + Self::Delete { name } => self.delete(&alias_store, name.clone()).await, + Self::List { + sort_by, + reverse, + name, + value, + } => { + self.list( + &alias_store, + *sort_by, + *reverse, + name.clone(), + value.clone(), + ) + .await + } + Self::Clear => self.clear(&alias_store).await, + } + } +} diff --git a/crates/atuin/src/command/client/dotfiles/var.rs b/crates/atuin/src/command/client/dotfiles/var.rs new file mode 100644 index 00000000..a63231ec --- /dev/null +++ b/crates/atuin/src/command/client/dotfiles/var.rs @@ -0,0 +1,197 @@ +use clap::{Subcommand, ValueEnum}; +use eyre::{Context, Result}; + +use atuin_client::{encryption, record::sqlite_store::SqliteStore, settings::Settings}; + +use atuin_dotfiles::{shell::Var, store::var::VarStore}; + +#[derive(Clone, Copy, Debug, Default, ValueEnum)] +pub enum SortBy { + /// Sort by variable name + #[default] + Name, + /// Sort by variable value + Value, +} + +#[derive(Subcommand, Debug)] +#[command(infer_subcommands = true)] +pub enum Cmd { + /// Set a variable + Set { + name: String, + value: String, + + #[clap(long, short, action)] + no_export: bool, + }, + + /// Delete a variable + Delete { name: String }, + + /// List all variables + List { + /// Sort results by field + #[arg(long, value_enum, default_value_t = SortBy::Name)] + sort_by: SortBy, + + /// Sort in reverse (descending) order + #[arg(long, short)] + reverse: bool, + + /// Filter variables by name (substring match) + #[arg(long, short)] + name: Option<String>, + + /// Filter variables by value (substring match) + #[arg(long, short)] + value: Option<String>, + + /// Show only exported variables + #[arg(long, conflicts_with = "shell_only")] + exports_only: bool, + + /// Show only non-exported (shell) variables + #[arg(long, conflicts_with = "exports_only")] + shell_only: bool, + }, +} + +impl Cmd { + async fn set(&self, store: VarStore, name: String, value: String, export: bool) -> Result<()> { + let vars = store.vars().await?; + let found: Vec<Var> = vars.into_iter().filter(|a| a.name == name).collect(); + let show_export = if export { "export " } else { "" }; + + if found.is_empty() { + println!("Setting '{show_export}{name}={value}'."); + } else { + println!( + "Overwriting var '{show_export}{name}={}' with '{name}={value}'.", + found[0].value + ); + } + + store.set(&name, &value, export).await?; + + Ok(()) + } + + #[allow(clippy::too_many_arguments)] + async fn list( + &self, + store: VarStore, + sort_by: SortBy, + reverse: bool, + name_filter: Option<String>, + value_filter: Option<String>, + exports_only: bool, + shell_only: bool, + ) -> Result<()> { + let mut vars = store.vars().await?; + + // Apply export/shell filters + if exports_only { + vars.retain(|v| v.export); + } + if shell_only { + vars.retain(|v| !v.export); + } + + // Apply name/value filters + if let Some(ref name_pattern) = name_filter { + let pattern = name_pattern.to_lowercase(); + vars.retain(|v| v.name.to_lowercase().contains(&pattern)); + } + if let Some(ref value_pattern) = value_filter { + let pattern = value_pattern.to_lowercase(); + vars.retain(|v| v.value.to_lowercase().contains(&pattern)); + } + + // Apply sorting + match sort_by { + SortBy::Name => { + vars.sort_by(|a, b| a.name.to_lowercase().cmp(&b.name.to_lowercase())); + } + SortBy::Value => { + vars.sort_by(|a, b| a.value.to_lowercase().cmp(&b.value.to_lowercase())); + } + } + + // Apply reverse if requested + if reverse { + vars.reverse(); + } + + for i in vars { + if i.export { + println!("export {}={}", i.name, i.value); + } else { + println!("{}={}", i.name, i.value); + } + } + + Ok(()) + } + + async fn delete(&self, store: VarStore, name: String) -> Result<()> { + let mut vars = store.vars().await?.into_iter(); + + if let Some(var) = vars.find(|var| var.name == name) { + println!("Deleting '{name}={}'.", var.value); + store.delete(&name).await?; + } else { + eprintln!("Cannot delete '{name}': Var not set."); + } + + Ok(()) + } + + pub async fn run(&self, settings: &Settings, store: SqliteStore) -> Result<()> { + if !settings.dotfiles.enabled { + eprintln!( + "Dotfiles are not enabled. Add\n\n[dotfiles]\nenabled = true\n\nto your configuration file to enable them.\n" + ); + eprintln!("The default configuration file is located at ~/.config/atuin/config.toml."); + return Ok(()); + } + + let encryption_key: [u8; 32] = encryption::load_key(settings) + .context("could not load encryption key")? + .into(); + let host_id = Settings::host_id().await?; + + let var_store = VarStore::new(store, host_id, encryption_key); + + match self { + Self::Set { + name, + value, + no_export, + } => { + self.set(var_store, name.clone(), value.clone(), !no_export) + .await + } + Self::Delete { name } => self.delete(var_store, name.clone()).await, + Self::List { + sort_by, + reverse, + name, + value, + exports_only, + shell_only, + } => { + self.list( + var_store, + *sort_by, + *reverse, + name.clone(), + value.clone(), + *exports_only, + *shell_only, + ) + .await + } + } + } +} diff --git a/crates/atuin/src/command/client/history.rs b/crates/atuin/src/command/client/history.rs new file mode 100644 index 00000000..fe9a7e32 --- /dev/null +++ b/crates/atuin/src/command/client/history.rs @@ -0,0 +1,848 @@ +use std::{ + fmt::{self, Display}, + io::{self, IsTerminal, Write}, + path::PathBuf, + time::Duration, +}; + +use atuin_common::utils::{self, Escapable as _}; +use clap::Subcommand; +use eyre::{Context, Result}; +use runtime_format::{FormatKey, FormatKeyError, ParseSegment, ParsedFmt}; + +#[cfg(feature = "daemon")] +use atuin_daemon::emit_event; + +use atuin_client::{ + database::{Database, Sqlite, current_context}, + encryption, + history::{History, store::HistoryStore}, + record::sqlite_store::SqliteStore, + settings::{ + FilterMode::{Directory, Global, Session}, + Settings, Timezone, + }, +}; + +#[cfg(feature = "sync")] +use atuin_client::{record, sync}; + +use log::{debug, warn}; +use time::{OffsetDateTime, macros::format_description}; + +#[cfg(feature = "daemon")] +use super::daemon; +use super::search::format_duration_into; + +#[derive(Subcommand, Debug)] +#[command(infer_subcommands = true)] +pub enum Cmd { + /// Begins a new command in the history + Start { + /// Collects the command from the `ATUIN_COMMAND_LINE` environment variable, + /// which does not need escaping and is more compatible between OS and shells + #[arg(long = "command-from-env", hide = true)] + cmd_env: bool, + + /// Author of this command, eg `ellie`, `claude`, or `copilot` + #[arg(long)] + author: Option<String>, + + /// Optional intent/rationale for running this command + #[arg(long)] + intent: Option<String>, + + command: Vec<String>, + }, + + /// Finishes a new command in the history (adds time, exit code) + End { + id: String, + #[arg(long, short)] + exit: i64, + #[arg(long, short)] + duration: Option<u64>, + }, + + /// List all items in history + List { + #[arg(long, short)] + cwd: bool, + + #[arg(long, short)] + session: bool, + + #[arg(long)] + human: bool, + + /// Show only the text of the command + #[arg(long)] + cmd_only: bool, + + /// Terminate the output with a null, for better multiline support + #[arg(long)] + print0: bool, + + #[arg(long, short, default_value = "true")] + // accept no value + #[arg(num_args(0..=1), default_missing_value("true"))] + // accept a value + #[arg(action = clap::ArgAction::Set)] + reverse: bool, + + /// Display the command time in another timezone other than the configured default. + /// + /// This option takes one of the following kinds of values: + /// - the special value "local" (or "l") which refers to the system time zone + /// - an offset from UTC (e.g. "+9", "-2:30") + #[arg(long, visible_alias = "tz")] + timezone: Option<Timezone>, + + /// Available variables: {command}, {directory}, {duration}, {user}, {host}, {author}, {intent}, {exit}, {time}, {session}, and {uuid} + /// Example: --format "{time} - [{duration}] - {directory}$\t{command}" + #[arg(long, short)] + format: Option<String>, + }, + + /// Get the last command ran + Last { + #[arg(long)] + human: bool, + + /// Show only the text of the command + #[arg(long)] + cmd_only: bool, + + /// Display the command time in another timezone other than the configured default. + /// + /// This option takes one of the following kinds of values: + /// - the special value "local" (or "l") which refers to the system time zone + /// - an offset from UTC (e.g. "+9", "-2:30") + #[arg(long, visible_alias = "tz")] + timezone: Option<Timezone>, + + /// Available variables: {command}, {directory}, {duration}, {user}, {host}, {author}, {intent}, {time}, {session}, {uuid} and {relativetime}. + /// Example: --format "{time} - [{duration}] - {directory}$\t{command}" + #[arg(long, short)] + format: Option<String>, + }, + + InitStore, + + /// Delete history entries matching the configured exclusion filters + Prune { + /// List matching history lines without performing the actual deletion. + #[arg(short = 'n', long)] + dry_run: bool, + }, + + /// Delete duplicate history entries (that have the same command, cwd and hostname) + Dedup { + /// List matching history lines without performing the actual deletion. + #[arg(short = 'n', long)] + dry_run: bool, + + /// Only delete results added before this date + #[arg(long, short)] + before: String, + + /// How many recent duplicates to keep + #[arg(long)] + dupkeep: u32, + }, +} + +#[derive(Clone, Copy, Debug)] +pub enum ListMode { + Human, + CmdOnly, + Regular, +} + +impl ListMode { + pub const fn from_flags(human: bool, cmd_only: bool) -> Self { + if human { + ListMode::Human + } else if cmd_only { + ListMode::CmdOnly + } else { + ListMode::Regular + } + } +} + +#[allow(clippy::cast_sign_loss)] +pub fn print_list( + h: &[History], + list_mode: ListMode, + format: Option<&str>, + print0: bool, + reverse: bool, + tz: Timezone, +) { + let w = std::io::stdout(); + let mut w = w.lock(); + + let fmt_str = match list_mode { + ListMode::Human => format + .unwrap_or("{time} · {duration}\t{command}") + .replace("\\t", "\t"), + ListMode::Regular => format + .unwrap_or("{time}\t{command}\t{duration}") + .replace("\\t", "\t"), + // not used + ListMode::CmdOnly => String::new(), + }; + + let parsed_fmt = match list_mode { + ListMode::Human | ListMode::Regular => parse_fmt(&fmt_str), + ListMode::CmdOnly => std::iter::once(ParseSegment::Key("command")).collect(), + }; + + let iterator = if reverse { + Box::new(h.iter().rev()) as Box<dyn Iterator<Item = &History>> + } else { + Box::new(h.iter()) as Box<dyn Iterator<Item = &History>> + }; + + let entry_terminator = if print0 { "\0" } else { "\n" }; + let flush_each_line = print0; + + for history in iterator { + let fh = FmtHistory { + history, + cmd_format: CmdFormat::for_output(&w), + tz: &tz, + }; + let args = parsed_fmt.with_args(&fh); + + // Check for formatting errors before attempting to write + if let Err(err) = args.status() { + eprintln!("ERROR: history output failed with: {err}"); + std::process::exit(1); + } + + let write_result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + write!(w, "{args}{entry_terminator}") + })); + + match write_result { + Ok(Ok(())) => { + // Write succeeded + } + Ok(Err(err)) => { + if err.kind() != io::ErrorKind::BrokenPipe { + eprintln!("ERROR: Failed to write history output: {err}"); + std::process::exit(1); + } + } + Err(_) => { + eprintln!("ERROR: Format string caused a formatting error."); + eprintln!( + "This may be due to an unsupported format string containing special characters." + ); + eprintln!( + "Please check your format string syntax and ensure literal braces are properly escaped." + ); + std::process::exit(1); + } + } + if flush_each_line { + check_for_write_errors(w.flush()); + } + } + + if !flush_each_line { + check_for_write_errors(w.flush()); + } +} + +fn check_for_write_errors(write: Result<(), io::Error>) { + if let Err(err) = write { + // Ignore broken pipe (issue #626) + if err.kind() != io::ErrorKind::BrokenPipe { + eprintln!("ERROR: History output failed with the following error: {err}"); + std::process::exit(1); + } + } +} + +/// Type wrapper around `History` with formatting settings. +#[derive(Clone, Copy, Debug)] +struct FmtHistory<'a> { + history: &'a History, + cmd_format: CmdFormat, + tz: &'a Timezone, +} + +#[derive(Clone, Copy, Debug)] +enum CmdFormat { + Literal, + Escaped, +} +impl CmdFormat { + fn for_output<O: IsTerminal>(out: &O) -> Self { + if out.is_terminal() { + Self::Escaped + } else { + Self::Literal + } + } +} + +static TIME_FMT: &[time::format_description::FormatItem<'static>] = + format_description!("[year]-[month]-[day] [hour repr:24]:[minute]:[second]"); + +/// defines how to format the history +impl FormatKey for FmtHistory<'_> { + #[allow(clippy::cast_sign_loss)] + fn fmt(&self, key: &str, f: &mut fmt::Formatter<'_>) -> Result<(), FormatKeyError> { + match key { + "command" => match self.cmd_format { + CmdFormat::Literal => f.write_str(self.history.command.trim()), + CmdFormat::Escaped => f.write_str(&self.history.command.trim().escape_control()), + }?, + "directory" => f.write_str(self.history.cwd.trim())?, + "exit" => f.write_str(&self.history.exit.to_string())?, + "duration" => { + let dur = Duration::from_nanos(std::cmp::max(self.history.duration, 0) as u64); + format_duration_into(dur, f)?; + } + "time" => { + self.history + .timestamp + .to_offset(self.tz.0) + .format(TIME_FMT) + .map_err(|_| fmt::Error)? + .fmt(f)?; + } + "relativetime" => { + let since = OffsetDateTime::now_utc() - self.history.timestamp; + let d = Duration::try_from(since).unwrap_or_default(); + format_duration_into(d, f)?; + } + "host" => f.write_str( + self.history + .hostname + .split_once(':') + .map_or(&self.history.hostname, |(host, _)| host), + )?, + "author" => f.write_str(&self.history.author)?, + "intent" => f.write_str(self.history.intent.as_deref().unwrap_or_default())?, + "user" => f.write_str( + self.history + .hostname + .split_once(':') + .map_or("", |(_, user)| user), + )?, + "session" => f.write_str(&self.history.session)?, + "uuid" => f.write_str(&self.history.id.0)?, + _ => return Err(FormatKeyError::UnknownKey), + } + Ok(()) + } +} + +fn parse_fmt(format: &str) -> ParsedFmt<'_> { + match ParsedFmt::new(format) { + Ok(fmt) => fmt, + Err(err) => { + eprintln!("ERROR: History formatting failed with the following error: {err}"); + + if format.contains('"') && (format.contains(":{") || format.contains(",{")) { + eprintln!("It looks like you're trying to create JSON output."); + eprintln!("For JSON, you need to escape literal braces by doubling them:"); + eprintln!("Example: '{{\"command\":\"{{command}}\",\"time\":\"{{time}}\"}}'"); + } else { + eprintln!( + "If your formatting string contains literal curly braces, you need to escape them by doubling:" + ); + eprintln!("Use {{{{ for literal {{ and }}}} for literal }}"); + } + std::process::exit(1) + } + } +} + +impl Cmd { + fn apply_start_metadata(history: &mut History, author: Option<&str>, intent: Option<&str>) { + if let Some(author) = author.map(str::trim).filter(|author| !author.is_empty()) { + author.clone_into(&mut history.author); + } + + if let Some(intent) = intent.map(str::trim).filter(|intent| !intent.is_empty()) { + history.intent = Some(intent.to_owned()); + } else if intent.is_some() { + history.intent = None; + } + } + + #[allow(clippy::too_many_lines, clippy::cast_possible_truncation)] + async fn handle_start( + db: &impl Database, + settings: &Settings, + command: &str, + author: Option<&str>, + intent: Option<&str>, + ) -> Result<()> { + // It's better for atuin to silently fail here and attempt to + // store whatever is ran, than to throw an error to the terminal + let cwd = utils::get_current_dir(); + + let mut h: History = History::capture() + .timestamp(OffsetDateTime::now_utc()) + .command(command) + .cwd(cwd) + .build() + .into(); + Self::apply_start_metadata(&mut h, author, intent); + + if !h.should_save(settings) { + return Ok(()); + } + + // print the ID + // we use this as the key for calling end + println!("{}", h.id); + + // Silently ignore database errors to avoid breaking the shell + // This is important when disk is full or database is locked + if let Err(e) = db.save(&h).await { + debug!("failed to save history: {e}"); + } + + Ok(()) + } + + #[cfg(feature = "daemon")] + async fn handle_daemon_start( + settings: &Settings, + command: &str, + author: Option<&str>, + intent: Option<&str>, + ) -> Result<()> { + // It's better for atuin to silently fail here and attempt to + // store whatever is ran, than to throw an error to the terminal + let cwd = utils::get_current_dir(); + + let mut h: History = History::capture() + .timestamp(OffsetDateTime::now_utc()) + .command(command) + .cwd(cwd) + .build() + .into(); + Self::apply_start_metadata(&mut h, author, intent); + + if !h.should_save(settings) { + return Ok(()); + } + + // Attempt to start history via daemon, but silently ignore errors + // to avoid breaking the shell when the daemon is unavailable or disk is full + let resp = match daemon::start_history(settings, h.clone()).await { + Ok(id) => id, + Err(e) => { + debug!("failed to start history via daemon: {e}"); + h.id.0.clone() + } + }; + + // print the ID + // we use this as the key for calling end + println!("{resp}"); + + Ok(()) + } + + #[allow(unused_variables)] + async fn handle_end( + db: &impl Database, + store: SqliteStore, + history_store: HistoryStore, + settings: &Settings, + id: &str, + exit: i64, + duration: Option<u64>, + ) -> Result<()> { + if id.trim() == "" { + return Ok(()); + } + + let Some(mut h) = db.load(id).await? else { + warn!("history entry is missing"); + return Ok(()); + }; + + if h.duration > 0 { + debug!("cannot end history - already has duration"); + + // returning OK as this can occur if someone Ctrl-c a prompt + return Ok(()); + } + + if !settings.store_failed && exit > 0 { + debug!("history has non-zero exit code, and store_failed is false"); + + // the history has already been inserted half complete. remove it + db.delete(h).await?; + + return Ok(()); + } + + h.exit = exit; + h.duration = match duration { + Some(value) => i64::try_from(value).context("command took over 292 years")?, + None => i64::try_from((OffsetDateTime::now_utc() - h.timestamp).whole_nanoseconds()) + .context("command took over 292 years")?, + }; + + db.update(&h).await?; + history_store.push(h).await?; + + if settings.should_sync().await? { + #[cfg(feature = "sync")] + { + if settings.sync.records { + let (_, downloaded) = record::sync::sync(settings, &store).await?; + Settings::save_sync_time().await?; + + crate::sync::build(settings, &store, db, Some(&downloaded)).await?; + } else { + debug!("running periodic background sync"); + sync::sync(settings, false, db).await?; + } + } + #[cfg(not(feature = "sync"))] + debug!("not compiled with sync support"); + } else { + debug!("sync disabled! not syncing"); + } + + Ok(()) + } + + #[cfg(feature = "daemon")] + async fn handle_daemon_end( + settings: &Settings, + id: &str, + exit: i64, + duration: Option<u64>, + ) -> Result<()> { + daemon::end_history(settings, id.to_string(), duration.unwrap_or(0), exit).await?; + + Ok(()) + } + + #[allow(clippy::too_many_arguments)] + #[allow(clippy::fn_params_excessive_bools)] + async fn handle_list( + db: &impl Database, + settings: &Settings, + context: atuin_client::database::Context, + session: bool, + cwd: bool, + mode: ListMode, + format: Option<String>, + include_deleted: bool, + print0: bool, + reverse: bool, + tz: Timezone, + ) -> Result<()> { + let filters = match (session, cwd) { + (true, true) => [Session, Directory], + (true, false) => [Session, Global], + (false, true) => [Global, Directory], + (false, false) => [ + settings.default_filter_mode(context.git_root.is_some()), + Global, + ], + }; + + let history = db + .list(&filters, &context, None, false, include_deleted) + .await?; + + print_list( + &history, + mode, + match format { + None => Some(settings.history_format.as_str()), + _ => format.as_deref(), + }, + print0, + reverse, + tz, + ); + + Ok(()) + } + + async fn handle_prune( + db: &impl Database, + settings: &Settings, + store: SqliteStore, + context: atuin_client::database::Context, + dry_run: bool, + ) -> Result<()> { + // Grab all executed commands and filter them using History::should_save. + // We could iterate or paginate here if memory usage becomes an issue. + let matches: Vec<History> = db + .list(&[Global], &context, None, false, false) + .await? + .into_iter() + .filter(|h| !h.should_save(settings)) + .collect(); + + match matches.len() { + 0 => { + println!("No entries to prune."); + return Ok(()); + } + 1 => println!("Found 1 entry to prune."), + n => println!("Found {n} entries to prune."), + } + + if dry_run { + print_list( + &matches, + ListMode::Human, + Some(settings.history_format.as_str()), + false, + false, + settings.timezone, + ); + } else { + let encryption_key: [u8; 32] = encryption::load_key(settings) + .context("could not load encryption key")? + .into(); + let host_id = Settings::host_id().await?; + let history_store = HistoryStore::new(store.clone(), host_id, encryption_key); + + for entry in matches { + eprintln!("deleting {}", entry.id); + if settings.sync.records { + let (id, _) = history_store.delete(entry.id.clone()).await?; + history_store.incremental_build(db, &[id]).await?; + } else { + db.delete(entry.clone()).await?; + } + } + + #[cfg(feature = "daemon")] + let _ = emit_event(atuin_daemon::DaemonEvent::HistoryPruned).await; + } + Ok(()) + } + + async fn handle_dedup( + db: &impl Database, + settings: &Settings, + store: SqliteStore, + before: i64, + dupkeep: u32, + dry_run: bool, + ) -> Result<()> { + if dupkeep == 0 { + eprintln!( + "\"--dupkeep 0\" would keep 0 copies of duplicate commands and thus delete all of them! Use \"atuin search --delete ...\" if you really want that." + ); + std::process::exit(1); + } + + let matches: Vec<History> = db.get_dups(before, dupkeep).await?; + + match matches.len() { + 0 => { + println!("No duplicates to delete."); + return Ok(()); + } + 1 => println!("Found 1 duplicate to delete."), + n => println!("Found {n} duplicates to delete."), + } + + if dry_run { + print_list( + &matches, + ListMode::Human, + Some(settings.history_format.as_str()), + false, + false, + settings.timezone, + ); + } else { + let encryption_key: [u8; 32] = encryption::load_key(settings) + .context("could not load encryption key")? + .into(); + let host_id = Settings::host_id().await?; + let history_store = HistoryStore::new(store.clone(), host_id, encryption_key); + + #[cfg(feature = "daemon")] + let ids = matches.iter().map(|h| h.id.clone()).collect::<Vec<_>>(); + + for entry in matches { + eprintln!("deleting {}", entry.id); + if settings.sync.records { + let (id, _) = history_store.delete(entry.id).await?; + history_store.incremental_build(db, &[id]).await?; + } else { + db.delete(entry).await?; + } + } + + #[cfg(feature = "daemon")] + let _ = emit_event(atuin_daemon::DaemonEvent::HistoryDeleted { ids }).await; + } + Ok(()) + } + + pub async fn run(self, settings: &Settings) -> Result<()> { + let context = current_context().await?; + + #[cfg(feature = "daemon")] + // Skip initializing any databases for start/end, if the daemon is enabled + if settings.daemon.enabled { + match self { + Self::Start { .. } => { + let command = self.get_start_command().unwrap_or_default(); + let (author, intent) = self.get_start_metadata().unwrap_or_default(); + return Self::handle_daemon_start(settings, &command, author, intent).await; + } + + Self::End { id, exit, duration } => { + return Self::handle_daemon_end(settings, &id, exit, duration).await; + } + + _ => {} + } + } + + let db_path = PathBuf::from(settings.db_path.as_str()); + let record_store_path = PathBuf::from(settings.record_store_path.as_str()); + + let db = Sqlite::new(db_path, settings.local_timeout).await?; + let store = SqliteStore::new(record_store_path, settings.local_timeout).await?; + + let encryption_key: [u8; 32] = encryption::load_key(settings) + .context("could not load encryption key")? + .into(); + + let host_id = Settings::host_id().await?; + let history_store = HistoryStore::new(store.clone(), host_id, encryption_key); + + match self { + Self::Start { .. } => { + let command = self.get_start_command().unwrap_or_default(); + let (author, intent) = self.get_start_metadata().unwrap_or_default(); + Self::handle_start(&db, settings, &command, author, intent).await + } + Self::End { id, exit, duration } => { + Self::handle_end(&db, store, history_store, settings, &id, exit, duration).await + } + Self::List { + session, + cwd, + human, + cmd_only, + print0, + reverse, + timezone, + format, + } => { + let mode = ListMode::from_flags(human, cmd_only); + let tz = timezone.unwrap_or(settings.timezone); + Self::handle_list( + &db, settings, context, session, cwd, mode, format, false, print0, reverse, tz, + ) + .await + } + + Self::Last { + human, + cmd_only, + timezone, + format, + } => { + let last = db.last().await?; + let last = last.as_slice(); + let tz = timezone.unwrap_or(settings.timezone); + print_list( + last, + ListMode::from_flags(human, cmd_only), + match format { + None => Some(settings.history_format.as_str()), + _ => format.as_deref(), + }, + false, + true, + tz, + ); + + Ok(()) + } + + Self::InitStore => history_store.init_store(&db).await, + + Self::Prune { dry_run } => { + Self::handle_prune(&db, settings, store, context, dry_run).await + } + + Self::Dedup { + dry_run, + before, + dupkeep, + } => { + let before = i64::try_from( + interim::parse_date_string( + before.as_str(), + OffsetDateTime::now_utc(), + interim::Dialect::Uk, + )? + .unix_timestamp_nanos(), + )?; + Self::handle_dedup(&db, settings, store, before, dupkeep, dry_run).await + } + } + } + + /// Returns the command line to use for the `Start` variant. + /// Returns `None` for any other variant. + fn get_start_command(&self) -> Option<String> { + match self { + Self::Start { cmd_env: true, .. } => { + Some(std::env::var("ATUIN_COMMAND_LINE").unwrap_or_default()) + } + Self::Start { command, .. } => Some(command.join(" ")), + _ => None, + } + } + + /// Returns `(author, intent)` for the `Start` variant. + /// Returns `None` for any other variant. + fn get_start_metadata(&self) -> Option<(Option<&str>, Option<&str>)> { + match self { + Self::Start { author, intent, .. } => Some((author.as_deref(), intent.as_deref())), + _ => None, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_format_string_no_panic() { + // Don't panic but provide helpful output (issue #2776) + let malformed_json = r#"{"command":"{command}","key":"value"}"#; + + let result = std::panic::catch_unwind(|| parse_fmt(malformed_json)); + + assert!(result.is_ok()); + } + + #[test] + fn test_valid_formats_still_work() { + assert!(std::panic::catch_unwind(|| parse_fmt("{command}")).is_ok()); + assert!(std::panic::catch_unwind(|| parse_fmt("{time} - {command}")).is_ok()); + } +} diff --git a/crates/atuin/src/command/client/import.rs b/crates/atuin/src/command/client/import.rs new file mode 100644 index 00000000..4df14ce8 --- /dev/null +++ b/crates/atuin/src/command/client/import.rs @@ -0,0 +1,186 @@ +use std::env; + +use async_trait::async_trait; +use clap::Parser; +use eyre::Result; +use indicatif::ProgressBar; + +use atuin_client::{ + database::Database, + history::History, + import::{ + Importer, Loader, bash::Bash, fish::Fish, nu::Nu, nu_histdb::NuHistDb, + powershell::PowerShell, replxx::Replxx, resh::Resh, xonsh::Xonsh, + xonsh_sqlite::XonshSqlite, zsh::Zsh, zsh_histdb::ZshHistDb, + }, +}; + +#[derive(Parser, Debug)] +#[command(infer_subcommands = true)] +pub enum Cmd { + /// Import history for the current shell + Auto, + + /// Import history from the zsh history file + Zsh, + /// Import history from the zsh history file + ZshHistDb, + /// Import history from the bash history file + Bash, + /// Import history from the replxx history file + Replxx, + /// Import history from the resh history file + Resh, + /// Import history from the fish history file + Fish, + /// Import history from the nu history file + Nu, + /// Import history from the nu history file + NuHistDb, + /// Import history from xonsh json files + Xonsh, + /// Import history from xonsh sqlite db + XonshSqlite, + /// Import history from the powershell history file + Powershell, +} + +const BATCH_SIZE: usize = 100; + +impl Cmd { + #[allow(clippy::cognitive_complexity)] + pub async fn run<DB: Database>(&self, db: &DB) -> Result<()> { + println!(" Atuin "); + println!("======================"); + println!(" \u{1f30d} "); + println!(" \u{1f418}\u{1f418}\u{1f418}\u{1f418} "); + println!(" \u{1f422} "); + println!("======================"); + println!("Importing history..."); + + match self { + Self::Auto => { + if cfg!(windows) { + return if env::var("PSModulePath").is_ok() { + println!("Detected PowerShell"); + import::<PowerShell, DB>(db).await + } else { + println!("Could not detect the current shell."); + println!("Please run atuin import <SHELL>."); + println!("To view a list of shells, run atuin import."); + Ok(()) + }; + } + + // $XONSH_HISTORY_BACKEND isn't always set, but $XONSH_HISTORY_FILE is + let xonsh_histfile = + env::var("XONSH_HISTORY_FILE").unwrap_or_else(|_| String::new()); + let shell = env::var("SHELL").unwrap_or_else(|_| String::from("NO_SHELL")); + + if xonsh_histfile.to_lowercase().ends_with(".json") { + println!("Detected Xonsh",); + import::<Xonsh, DB>(db).await + } else if xonsh_histfile.to_lowercase().ends_with(".sqlite") { + println!("Detected Xonsh (SQLite backend)"); + import::<XonshSqlite, DB>(db).await + } else if shell.ends_with("/zsh") { + if ZshHistDb::histpath().is_ok() { + println!( + "Detected Zsh-HistDb, using :{}", + ZshHistDb::histpath().unwrap().to_str().unwrap() + ); + import::<ZshHistDb, DB>(db).await + } else { + println!("Detected ZSH"); + import::<Zsh, DB>(db).await + } + } else if shell.ends_with("/fish") { + println!("Detected Fish"); + import::<Fish, DB>(db).await + } else if shell.ends_with("/bash") { + println!("Detected Bash"); + import::<Bash, DB>(db).await + } else if shell.ends_with("/nu") { + if NuHistDb::histpath().is_ok() { + println!( + "Detected Nu-HistDb, using :{}", + NuHistDb::histpath().unwrap().to_str().unwrap() + ); + import::<NuHistDb, DB>(db).await + } else { + println!("Detected Nushell"); + import::<Nu, DB>(db).await + } + } else if shell.ends_with("/pwsh") { + println!("Detected PowerShell"); + import::<PowerShell, DB>(db).await + } else { + println!("cannot import {shell} history"); + Ok(()) + } + } + + Self::Zsh => import::<Zsh, DB>(db).await, + Self::ZshHistDb => import::<ZshHistDb, DB>(db).await, + Self::Bash => import::<Bash, DB>(db).await, + Self::Replxx => import::<Replxx, DB>(db).await, + Self::Resh => import::<Resh, DB>(db).await, + Self::Fish => import::<Fish, DB>(db).await, + Self::Nu => import::<Nu, DB>(db).await, + Self::NuHistDb => import::<NuHistDb, DB>(db).await, + Self::Xonsh => import::<Xonsh, DB>(db).await, + Self::XonshSqlite => import::<XonshSqlite, DB>(db).await, + Self::Powershell => import::<PowerShell, DB>(db).await, + } + } +} + +pub struct HistoryImporter<'db, DB: Database> { + pb: ProgressBar, + buf: Vec<History>, + db: &'db DB, +} + +impl<'db, DB: Database> HistoryImporter<'db, DB> { + fn new(db: &'db DB, len: usize) -> Self { + Self { + pb: ProgressBar::new(len as u64), + buf: Vec::with_capacity(BATCH_SIZE), + db, + } + } + + async fn flush(self) -> Result<()> { + if !self.buf.is_empty() { + self.db.save_bulk(&self.buf).await?; + } + self.pb.finish(); + Ok(()) + } +} + +#[async_trait] +impl<DB: Database> Loader for HistoryImporter<'_, DB> { + async fn push(&mut self, hist: History) -> Result<()> { + self.pb.inc(1); + self.buf.push(hist); + if self.buf.len() == self.buf.capacity() { + self.db.save_bulk(&self.buf).await?; + self.buf.clear(); + } + Ok(()) + } +} + +async fn import<I: Importer + Send, DB: Database>(db: &DB) -> Result<()> { + println!("Importing history from {}", I::NAME); + + let mut importer = I::new().await?; + let len = importer.entries().await.unwrap(); + let mut loader = HistoryImporter::new(db, len); + importer.load(&mut loader).await?; + loader.flush().await?; + + println!("Import complete!"); + Ok(()) +} diff --git a/crates/atuin/src/command/client/info.rs b/crates/atuin/src/command/client/info.rs new file mode 100644 index 00000000..a69f9b2f --- /dev/null +++ b/crates/atuin/src/command/client/info.rs @@ -0,0 +1,31 @@ +use atuin_client::settings::Settings;
+
+use crate::{SHA, VERSION};
+
+pub fn run(settings: &Settings) {
+ let config = atuin_common::utils::config_dir();
+ let mut config_file = config.clone();
+ config_file.push("config.toml");
+ let mut sever_config = config;
+ sever_config.push("server.toml");
+
+ let config_paths = format!(
+ "Config files:\nclient config: {:?}\nserver config: {:?}\nclient db path: {:?}\nkey path: {:?}\nmeta db path: {:?}",
+ config_file.to_string_lossy(),
+ sever_config.to_string_lossy(),
+ settings.db_path,
+ settings.key_path,
+ settings.meta.db_path
+ );
+
+ let env_vars = format!(
+ "Env Vars:\nATUIN_CONFIG_DIR = {:?}",
+ std::env::var("ATUIN_CONFIG_DIR").unwrap_or_else(|_| "None".into())
+ );
+
+ let general_info = format!("Version info:\nversion: {VERSION}\ncommit: {SHA}");
+
+ let print_out = format!("{config_paths}\n\n{env_vars}\n\n{general_info}");
+
+ println!("{print_out}");
+}
diff --git a/crates/atuin/src/command/client/init.rs b/crates/atuin/src/command/client/init.rs new file mode 100644 index 00000000..00c6c2fc --- /dev/null +++ b/crates/atuin/src/command/client/init.rs @@ -0,0 +1,227 @@ +use std::path::PathBuf; + +use atuin_client::{ + encryption, + record::sqlite_store::SqliteStore, + settings::{Settings, Tmux}, +}; +use atuin_dotfiles::store::{AliasStore, var::VarStore}; +use clap::{Parser, ValueEnum}; +use eyre::{Result, WrapErr}; + +mod bash; +mod fish; +mod powershell; +mod xonsh; +mod zsh; + +#[derive(Parser, Debug)] +pub struct Cmd { + shell: Shell, + + /// Disable the binding of CTRL-R to atuin + #[clap(long)] + disable_ctrl_r: bool, + + /// Disable the binding of the Up Arrow key to atuin + #[clap(long)] + disable_up_arrow: bool, + + /// Disable the binding of ? to Atuin AI + #[clap(long)] + disable_ai: bool, +} + +#[derive(Clone, Copy, ValueEnum, Debug)] +#[value(rename_all = "lower")] +#[allow(clippy::enum_variant_names, clippy::doc_markdown)] +pub enum Shell { + /// Zsh setup + Zsh, + /// Bash setup + Bash, + /// Fish setup + Fish, + /// Nu setup + Nu, + /// Xonsh setup + Xonsh, + /// PowerShell setup + PowerShell, +} + +impl Cmd { + fn init_nu(&self, _tmux: &Tmux) { + let full = include_str!("../../shell/atuin.nu"); + + // TODO: tmux popup for Nu + println!("{full}"); + + if std::env::var("ATUIN_NOBIND").is_err() { + const BIND_CTRL_R: &str = r"$env.config = ( + $env.config | upsert keybindings ( + $env.config.keybindings + | append { + name: atuin + modifier: control + keycode: char_r + mode: [emacs, vi_normal, vi_insert] + event: { send: executehostcommand cmd: (_atuin_search_cmd) } + } + ) +)"; + const BIND_UP_ARROW: &str = r" +$env.config = ( + $env.config | upsert keybindings ( + $env.config.keybindings + | append { + name: atuin + modifier: none + keycode: up + mode: [emacs, vi_normal, vi_insert] + event: { + until: [ + {send: menuup} + {send: executehostcommand cmd: (_atuin_search_cmd '--shell-up-key-binding') } + ] + } + } + ) +) +"; + if !self.disable_ctrl_r { + println!("{BIND_CTRL_R}"); + } + if !self.disable_up_arrow { + println!("{BIND_UP_ARROW}"); + } + } + } + + fn static_init(&self, tmux: &Tmux) { + match self.shell { + Shell::Zsh => { + zsh::init_static( + self.disable_up_arrow, + self.disable_ctrl_r, + self.disable_ai, + tmux, + ); + } + Shell::Bash => { + bash::init_static( + self.disable_up_arrow, + self.disable_ctrl_r, + self.disable_ai, + tmux, + ); + } + Shell::Fish => { + fish::init_static( + self.disable_up_arrow, + self.disable_ctrl_r, + self.disable_ai, + tmux, + ); + } + Shell::Nu => { + self.init_nu(tmux); + } + Shell::Xonsh => { + xonsh::init_static(self.disable_up_arrow, self.disable_ctrl_r, tmux); + } + Shell::PowerShell => { + powershell::init_static(self.disable_up_arrow, self.disable_ctrl_r, tmux); + } + } + } + + async fn dotfiles_init(&self, settings: &Settings) -> Result<()> { + let record_store_path = PathBuf::from(settings.record_store_path.as_str()); + let sqlite_store = SqliteStore::new(record_store_path, settings.local_timeout).await?; + + let encryption_key: [u8; 32] = encryption::load_key(settings) + .context("could not load encryption key")? + .into(); + let host_id = Settings::host_id().await?; + + let alias_store = AliasStore::new(sqlite_store.clone(), host_id, encryption_key); + let var_store = VarStore::new(sqlite_store.clone(), host_id, encryption_key); + + match self.shell { + Shell::Zsh => { + zsh::init( + alias_store, + var_store, + self.disable_up_arrow, + self.disable_ctrl_r, + self.disable_ai, + &settings.tmux, + ) + .await?; + } + Shell::Bash => { + bash::init( + alias_store, + var_store, + self.disable_up_arrow, + self.disable_ctrl_r, + self.disable_ai, + &settings.tmux, + ) + .await?; + } + Shell::Fish => { + fish::init( + alias_store, + var_store, + self.disable_up_arrow, + self.disable_ctrl_r, + self.disable_ai, + &settings.tmux, + ) + .await?; + } + Shell::Nu => self.init_nu(&settings.tmux), + Shell::Xonsh => { + xonsh::init( + alias_store, + var_store, + self.disable_up_arrow, + self.disable_ctrl_r, + &settings.tmux, + ) + .await?; + } + Shell::PowerShell => { + powershell::init( + alias_store, + var_store, + self.disable_up_arrow, + self.disable_ctrl_r, + &settings.tmux, + ) + .await?; + } + } + + Ok(()) + } + + pub async fn run(self, settings: &Settings) -> Result<()> { + if !settings.paths_ok() { + eprintln!( + "Atuin settings paths are broken. Disabling atuin shell hooks. Run `atuin doctor` to diagnose." + ); + return Ok(()); + } + + if settings.dotfiles.enabled { + self.dotfiles_init(settings).await?; + } else { + self.static_init(&settings.tmux); + } + + Ok(()) + } +} diff --git a/crates/atuin/src/command/client/init/bash.rs b/crates/atuin/src/command/client/init/bash.rs new file mode 100644 index 00000000..745c239a --- /dev/null +++ b/crates/atuin/src/command/client/init/bash.rs @@ -0,0 +1,52 @@ +use atuin_client::settings::Tmux; +use atuin_dotfiles::store::{AliasStore, var::VarStore}; +use eyre::Result; + +fn print_tmux_config(tmux: &Tmux) { + if tmux.enabled { + println!("export ATUIN_TMUX_POPUP_WIDTH='{}'", tmux.width); + println!("export ATUIN_TMUX_POPUP_HEIGHT='{}'", tmux.height); + } else { + println!("export ATUIN_TMUX_POPUP=false"); + } +} + +pub fn init_static(disable_up_arrow: bool, disable_ctrl_r: bool, disable_ai: bool, tmux: &Tmux) { + let base = include_str!("../../../shell/atuin.bash"); + + let (bind_ctrl_r, bind_up_arrow) = if std::env::var("ATUIN_NOBIND").is_ok() { + (false, false) + } else { + (!disable_ctrl_r, !disable_up_arrow) + }; + + print_tmux_config(tmux); + println!("__atuin_bind_ctrl_r={bind_ctrl_r}"); + println!("__atuin_bind_up_arrow={bind_up_arrow}"); + println!("{base}"); + + #[cfg(feature = "ai")] + if !disable_ai { + let bind_ai = atuin_ai::commands::init::generate_bash_integration(); + println!("{bind_ai}"); + } +} + +pub async fn init( + aliases: AliasStore, + vars: VarStore, + disable_up_arrow: bool, + disable_ctrl_r: bool, + disable_ai: bool, + tmux: &Tmux, +) -> Result<()> { + init_static(disable_up_arrow, disable_ctrl_r, disable_ai, tmux); + + let aliases = atuin_dotfiles::shell::bash::alias_config(&aliases).await; + let vars = atuin_dotfiles::shell::bash::var_config(&vars).await; + + println!("{aliases}"); + println!("{vars}"); + + Ok(()) +} diff --git a/crates/atuin/src/command/client/init/fish.rs b/crates/atuin/src/command/client/init/fish.rs new file mode 100644 index 00000000..6d6c8c23 --- /dev/null +++ b/crates/atuin/src/command/client/init/fish.rs @@ -0,0 +1,113 @@ +use atuin_client::settings::Tmux; +use atuin_dotfiles::store::{AliasStore, var::VarStore}; +use eyre::Result; + +fn print_tmux_config(tmux: &Tmux) { + if tmux.enabled { + println!("set -gx ATUIN_TMUX_POPUP_WIDTH '{}'", tmux.width); + println!("set -gx ATUIN_TMUX_POPUP_HEIGHT '{}'", tmux.height); + } else { + println!("set -gx ATUIN_TMUX_POPUP false"); + } +} + +fn print_bindings( + indent: &str, + disable_up_arrow: bool, + disable_ctrl_r: bool, + bind_ctrl_r: &str, + bind_up_arrow: &str, + bind_ctrl_r_ins: &str, + bind_up_arrow_ins: &str, +) { + if !disable_ctrl_r { + println!("{indent}{bind_ctrl_r}"); + } + if !disable_up_arrow { + println!("{indent}{bind_up_arrow}"); + } + + println!("{indent}if bind -M insert >/dev/null 2>&1"); + if !disable_ctrl_r { + println!("{indent}{indent}{bind_ctrl_r_ins}"); + } + if !disable_up_arrow { + println!("{indent}{indent}{bind_up_arrow_ins}"); + } + println!("{indent}end"); +} + +pub fn init_static(disable_up_arrow: bool, disable_ctrl_r: bool, disable_ai: bool, tmux: &Tmux) { + let indent = " ".repeat(4); + + let base = include_str!("../../../shell/atuin.fish"); + + print_tmux_config(tmux); + println!("{base}"); + + if std::env::var("ATUIN_NOBIND").is_err() { + println!("if string match -q '4.*' $version"); + + // In fish 4.0 and above the option bind -k doesn't exist anymore, + // instead we can use key names and modifiers directly. + print_bindings( + &indent, + disable_up_arrow, + disable_ctrl_r, + "bind ctrl-r _atuin_search", + "bind up _atuin_bind_up", + "bind -M insert ctrl-r _atuin_search", + "bind -M insert up _atuin_bind_up", + ); + + println!("else"); + + // We keep these for compatibility with fish 3.x + print_bindings( + &indent, + disable_up_arrow, + disable_ctrl_r, + r"bind \cr _atuin_search", + &[ + r"bind -k up _atuin_bind_up", + r"bind \eOA _atuin_bind_up", + r"bind \e\[A _atuin_bind_up", + ] + .join("; "), + r"bind -M insert \cr _atuin_search", + &[ + r"bind -M insert -k up _atuin_bind_up", + r"bind -M insert \eOA _atuin_bind_up", + r"bind -M insert \e\[A _atuin_bind_up", + ] + .join("; "), + ); + + println!("end"); + + #[cfg(feature = "ai")] + if !disable_ai { + let bind_ai = atuin_ai::commands::init::generate_fish_integration(); + println!("{bind_ai}"); + } + } +} + +pub async fn init( + aliases: AliasStore, + vars: VarStore, + disable_up_arrow: bool, + disable_ctrl_r: bool, + disable_ai: bool, + tmux: &Tmux, +) -> Result<()> { + init_static(disable_up_arrow, disable_ctrl_r, disable_ai, tmux); + + let aliases = atuin_dotfiles::shell::fish::alias_config(&aliases).await; + let vars = atuin_dotfiles::shell::fish::var_config(&vars).await; + + println!("{aliases}"); + println!("{vars}"); + + Ok(()) +} diff --git a/crates/atuin/src/command/client/init/powershell.rs b/crates/atuin/src/command/client/init/powershell.rs new file mode 100644 index 00000000..d3399404 --- /dev/null +++ b/crates/atuin/src/command/client/init/powershell.rs @@ -0,0 +1,42 @@ +use atuin_client::settings::Tmux; +use atuin_dotfiles::store::{AliasStore, var::VarStore}; + +pub fn init_static(disable_up_arrow: bool, disable_ctrl_r: bool, _tmux: &Tmux) { + let base = include_str!("../../../shell/atuin.ps1"); + + let (bind_ctrl_r, bind_up_arrow) = if std::env::var("ATUIN_NOBIND").is_ok() { + (false, false) + } else { + (!disable_ctrl_r, !disable_up_arrow) + }; + + // TODO: tmux popup for Powershell + println!("{base}"); + println!( + "Enable-AtuinSearchKeys -CtrlR {} -UpArrow {}", + ps_bool(bind_ctrl_r), + ps_bool(bind_up_arrow) + ); +} + +pub async fn init( + aliases: AliasStore, + vars: VarStore, + disable_up_arrow: bool, + disable_ctrl_r: bool, + tmux: &Tmux, +) -> eyre::Result<()> { + init_static(disable_up_arrow, disable_ctrl_r, tmux); + + let aliases = atuin_dotfiles::shell::powershell::alias_config(&aliases).await; + let vars = atuin_dotfiles::shell::powershell::var_config(&vars).await; + + println!("{aliases}"); + println!("{vars}"); + + Ok(()) +} + +fn ps_bool(value: bool) -> &'static str { + if value { "$true" } else { "$false" } +} diff --git a/crates/atuin/src/command/client/init/xonsh.rs b/crates/atuin/src/command/client/init/xonsh.rs new file mode 100644 index 00000000..8b9f1595 --- /dev/null +++ b/crates/atuin/src/command/client/init/xonsh.rs @@ -0,0 +1,42 @@ +use atuin_client::settings::Tmux; +use atuin_dotfiles::store::{AliasStore, var::VarStore}; +use eyre::Result; + +pub fn init_static(disable_up_arrow: bool, disable_ctrl_r: bool, _tmux: &Tmux) { + let base = include_str!("../../../shell/atuin.xsh"); + + let (bind_ctrl_r, bind_up_arrow) = if std::env::var("ATUIN_NOBIND").is_ok() { + (false, false) + } else { + (!disable_ctrl_r, !disable_up_arrow) + }; + + // TODO: tmux popup for xonsh + println!( + "_ATUIN_BIND_CTRL_R={}", + if bind_ctrl_r { "True" } else { "False" } + ); + println!( + "_ATUIN_BIND_UP_ARROW={}", + if bind_up_arrow { "True" } else { "False" } + ); + println!("{base}"); +} + +pub async fn init( + aliases: AliasStore, + vars: VarStore, + disable_up_arrow: bool, + disable_ctrl_r: bool, + tmux: &Tmux, +) -> Result<()> { + init_static(disable_up_arrow, disable_ctrl_r, tmux); + + let aliases = atuin_dotfiles::shell::xonsh::alias_config(&aliases).await; + let vars = atuin_dotfiles::shell::xonsh::var_config(&vars).await; + + println!("{aliases}"); + println!("{vars}"); + + Ok(()) +} diff --git a/crates/atuin/src/command/client/init/zsh.rs b/crates/atuin/src/command/client/init/zsh.rs new file mode 100644 index 00000000..5d588aa0 --- /dev/null +++ b/crates/atuin/src/command/client/init/zsh.rs @@ -0,0 +1,66 @@ +use atuin_client::settings::Tmux; +use atuin_dotfiles::store::{AliasStore, var::VarStore}; +use eyre::Result; + +fn print_tmux_config(tmux: &Tmux) { + if tmux.enabled { + println!("export ATUIN_TMUX_POPUP_WIDTH='{}'", tmux.width); + println!("export ATUIN_TMUX_POPUP_HEIGHT='{}'", tmux.height); + } else { + println!("export ATUIN_TMUX_POPUP=false"); + } +} + +pub fn init_static(disable_up_arrow: bool, disable_ctrl_r: bool, disable_ai: bool, tmux: &Tmux) { + let base = include_str!("../../../shell/atuin.zsh"); + + print_tmux_config(tmux); + println!("{base}"); + + if std::env::var("ATUIN_NOBIND").is_err() { + const BIND_CTRL_R: &str = r"bindkey -M emacs '^r' atuin-search +bindkey -M viins '^r' atuin-search-viins +bindkey -M vicmd '/' atuin-search"; + + const BIND_UP_ARROW: &str = r"bindkey -M emacs '^[[A' atuin-up-search +bindkey -M vicmd '^[[A' atuin-up-search-vicmd +bindkey -M viins '^[[A' atuin-up-search-viins +bindkey -M emacs '^[OA' atuin-up-search +bindkey -M vicmd '^[OA' atuin-up-search-vicmd +bindkey -M viins '^[OA' atuin-up-search-viins +bindkey -M vicmd 'k' atuin-up-search-vicmd"; + + if !disable_ctrl_r { + println!("{BIND_CTRL_R}"); + } + if !disable_up_arrow { + println!("{BIND_UP_ARROW}"); + } + + #[cfg(feature = "ai")] + if !disable_ai { + let bind_ai = atuin_ai::commands::init::generate_zsh_integration(); + + println!("{bind_ai}"); + } + } +} + +pub async fn init( + aliases: AliasStore, + vars: VarStore, + disable_up_arrow: bool, + disable_ctrl_r: bool, + disable_ai: bool, + tmux: &Tmux, +) -> Result<()> { + init_static(disable_up_arrow, disable_ctrl_r, disable_ai, tmux); + + let aliases = atuin_dotfiles::shell::zsh::alias_config(&aliases).await; + let vars = atuin_dotfiles::shell::zsh::var_config(&vars).await; + + println!("{aliases}"); + println!("{vars}"); + + Ok(()) +} diff --git a/crates/atuin/src/command/client/kv.rs b/crates/atuin/src/command/client/kv.rs new file mode 100644 index 00000000..de487b89 --- /dev/null +++ b/crates/atuin/src/command/client/kv.rs @@ -0,0 +1,122 @@ +use clap::Subcommand; +use eyre::{Context, Result, eyre}; + +use atuin_client::{encryption, record::sqlite_store::SqliteStore, settings::Settings}; +use atuin_kv::store::KvStore; + +#[derive(Subcommand, Debug)] +#[command(infer_subcommands = true)] +pub enum Cmd { + /// Set a key-value pair + Set { + /// Key to set + #[arg(long, short)] + key: String, + + /// Value to store + value: String, + + /// Namespace for the key-value pair + #[arg(long, short, default_value = "default")] + namespace: String, + }, + + /// Delete one or more key-value pairs + #[command(alias = "rm")] + Delete { + /// Keys to delete + #[arg(required = true)] + keys: Vec<String>, + + /// Namespace for the key-value pair + #[arg(long, short, default_value = "default")] + namespace: String, + }, + + /// Retrieve a saved value + Get { + /// Key to retrieve + key: String, + + /// Namespace for the key-value pair + #[arg(long, short, default_value = "default")] + namespace: String, + }, + + /// List all keys in a namespace, or in all namespaces + #[command(alias = "ls")] + List { + /// Namespace to list keys from + #[arg(long, short, default_value = "default")] + namespace: String, + + /// List all keys in all namespaces + #[arg(long, short, alias = "all")] + all_namespaces: bool, + }, + + /// Rebuild the KV store + Rebuild, +} + +impl Cmd { + pub async fn run(&self, settings: &Settings, store: &SqliteStore) -> Result<()> { + let encryption_key: [u8; 32] = encryption::load_key(settings) + .context("could not load encryption key")? + .into(); + + let host_id = Settings::host_id().await?; + + let kv_db = atuin_kv::database::Database::new(settings.kv.db_path.clone(), 1.0).await?; + let kv_store = KvStore::new(store.clone(), kv_db, host_id, encryption_key); + + match self { + Self::Set { + key, + value, + namespace, + } => { + if namespace.is_empty() { + return Err(eyre!("namespace cannot be empty")); + } + + kv_store.set(namespace, key, value).await + } + + Self::Delete { keys, namespace } => kv_store.delete(namespace, keys).await, + + Self::Get { key, namespace } => { + let kv = kv_store.get(namespace, key).await?; + + if let Some(val) = kv { + println!("{val}"); + } + + Ok(()) + } + + Self::List { + namespace, + all_namespaces, + } => { + let entries = if *all_namespaces { + kv_store.list(None).await? + } else { + kv_store.list(Some(namespace)).await? + }; + + for entry in entries { + if *all_namespaces { + println!("{}.{}", entry.namespace, entry.key); + } else { + println!("{}", entry.key); + } + } + + Ok(()) + } + + Self::Rebuild {} => kv_store.build().await, + } + } +} diff --git a/crates/atuin/src/command/client/scripts.rs b/crates/atuin/src/command/client/scripts.rs new file mode 100644 index 00000000..e5adacc4 --- /dev/null +++ b/crates/atuin/src/command/client/scripts.rs @@ -0,0 +1,590 @@ +use std::collections::HashMap; +use std::collections::HashSet; +use std::io::IsTerminal; +use std::io::Read; +use std::path::PathBuf; + +use atuin_scripts::execution::template_script; +use atuin_scripts::{ + execution::{build_executable_script, execute_script_interactive, template_variables}, + store::{ScriptStore, script::Script}, +}; +use clap::{Parser, Subcommand}; +use eyre::OptionExt; +use eyre::{Result, bail}; +use tempfile::NamedTempFile; + +use atuin_client::{database::Database, record::sqlite_store::SqliteStore, settings::Settings}; +use tracing::debug; + +#[derive(Parser, Debug)] +pub struct NewScript { + pub name: String, + + #[arg(short, long)] + pub description: Option<String>, + + #[arg(short, long)] + pub tags: Vec<String>, + + #[arg(short, long)] + pub shebang: Option<String>, + + #[arg(long)] + pub script: Option<PathBuf>, + + #[allow(clippy::option_option)] + #[arg(long)] + /// Use the last command as the script content + /// Optionally specify a number to use the last N commands + pub last: Option<Option<usize>>, + + #[arg(long)] + /// Skip opening editor when using --last + pub no_edit: bool, +} + +#[derive(Parser, Debug)] +pub struct Run { + pub name: String, + + /// Specify template variables in the format KEY=VALUE + /// Example: -v name=John -v greeting="Hello there" + #[arg(short, long = "var")] + pub var: Vec<String>, +} + +#[derive(Parser, Debug)] +pub struct List {} + +#[derive(Parser, Debug)] +pub struct Get { + pub name: String, + + #[arg(short, long)] + /// Display only the executable script with shebang + pub script: bool, +} + +#[derive(Parser, Debug)] +pub struct Edit { + pub name: String, + + #[arg(short, long)] + pub description: Option<String>, + + /// Replace all existing tags with these new tags + #[arg(short, long)] + pub tags: Vec<String>, + + /// Remove all tags from the script + #[arg(long)] + pub no_tags: bool, + + /// Rename the script + #[arg(long)] + pub rename: Option<String>, + + #[arg(short, long)] + pub shebang: Option<String>, + + #[arg(long)] + pub script: Option<PathBuf>, + + #[allow(clippy::struct_field_names)] + /// Skip opening editor + #[arg(long)] + pub no_edit: bool, +} + +#[derive(Parser, Debug)] +pub struct Delete { + pub name: String, + + #[arg(short, long)] + pub force: bool, +} + +#[derive(Subcommand, Debug)] +#[command(infer_subcommands = true)] +pub enum Cmd { + New(NewScript), + Run(Run), + #[command(alias = "ls")] + List(List), + + Get(Get), + Edit(Edit), + #[command(alias = "rm")] + Delete(Delete), +} + +impl Cmd { + // Helper function to open an editor with optional initial content + fn open_editor(initial_content: Option<&str>) -> Result<String> { + // Create a temporary file + let temp_file = NamedTempFile::new()?; + let path = temp_file.into_temp_path(); + + // Write initial content to the temp file if provided + if let Some(content) = initial_content { + std::fs::write(&path, content)?; + } + + // Open the file in the user's preferred editor + let editor_str = std::env::var("EDITOR").unwrap_or_else(|_| "vi".to_string()); + + // Use shlex to safely split the string into shell-like parts. + let parts = shlex::split(&editor_str).ok_or_eyre("Failed to parse editor command")?; + let (command, args) = parts.split_first().ok_or_eyre("No editor command found")?; + + let status = std::process::Command::new(command) + .args(args) + .arg(&path) + .status()?; + if !status.success() { + bail!("failed to open editor"); + } + + // Read back the edited content + let content = std::fs::read_to_string(&path)?; + path.close()?; + + Ok(content) + } + + // Helper function to execute a script and manage stdin/stdout/stderr + async fn execute_script(script_content: String, shebang: String) -> Result<i32> { + let mut session = execute_script_interactive(script_content, shebang) + .await + .expect("failed to execute script"); + + // Create a channel to signal when the process exits + let (exit_tx, mut exit_rx) = tokio::sync::oneshot::channel(); + + // Set up a task to read from stdin and forward to the script + let sender = session.stdin_tx.clone(); + let stdin_task = tokio::spawn(async move { + use tokio::io::AsyncReadExt; + use tokio::select; + + let stdin = tokio::io::stdin(); + let mut reader = tokio::io::BufReader::new(stdin); + let mut buffer = vec![0u8; 1024]; // Read in chunks for efficiency + + loop { + // Use select to either read from stdin or detect when the process exits + select! { + // Check if the script process has exited + _ = &mut exit_rx => { + break; + } + // Try to read from stdin + read_result = reader.read(&mut buffer) => { + match read_result { + Ok(0) => break, // EOF + Ok(n) => { + // Convert the bytes to a string and forward to script + let input = String::from_utf8_lossy(&buffer[0..n]).to_string(); + if let Err(e) = sender.send(input).await { + eprintln!("Error sending input to script: {e}"); + break; + } + }, + Err(e) => { + eprintln!("Error reading from stdin: {e}"); + break; + } + } + } + } + } + }); + + // Wait for the script to complete + let exit_code = session.wait_for_exit().await; + + // Signal the stdin task to stop + let _ = exit_tx.send(()); + let _ = stdin_task.await; + + let code = exit_code.unwrap_or(-1); + if code != 0 { + eprintln!("Script exited with code {code}"); + } + + Ok(code) + } + + async fn handle_new_script( + settings: &Settings, + new_script: NewScript, + script_store: ScriptStore, + script_db: atuin_scripts::database::Database, + history_db: &impl Database, + ) -> Result<()> { + let mut stdin = std::io::stdin(); + let script_content = if let Some(count_opt) = new_script.last { + // Get the last N commands from history, plus 1 to exclude the command that runs this script + let count = count_opt.unwrap_or(1) + 1; // Add 1 to the count to exclude the current command + let context = atuin_client::database::current_context().await?; + + // Get the last N+1 commands, filtering by the default mode + let filters = [settings.default_filter_mode(context.git_root.is_some())]; + + let mut history = history_db + .list(&filters, &context, Some(count), false, false) + .await?; + + // Reverse to get chronological order + history.reverse(); + + // Skip the most recent command (which would be the atuin scripts new command itself) + if !history.is_empty() { + history.pop(); // Remove the most recent command + } + + // Format the commands into a script + let commands: Vec<String> = history.iter().map(|h| h.command.clone()).collect(); + + if commands.is_empty() { + bail!("No commands found in history"); + } + + let script_text = commands.join("\n"); + + // Only open editor if --no-edit is not specified + if new_script.no_edit { + Some(script_text) + } else { + // Open the editor with the commands pre-loaded + Some(Self::open_editor(Some(&script_text))?) + } + } else if let Some(script_path) = new_script.script { + let script_content = std::fs::read_to_string(script_path)?; + Some(script_content) + } else if !stdin.is_terminal() { + let mut buffer = String::new(); + stdin.read_to_string(&mut buffer)?; + Some(buffer) + } else { + // Open editor with empty file + Some(Self::open_editor(None)?) + }; + + let script = Script::builder() + .name(new_script.name) + .description(new_script.description.unwrap_or_default()) + .shebang(new_script.shebang.unwrap_or_default()) + .tags(new_script.tags) + .script(script_content.unwrap_or_default()) + .build(); + + script_store.create(script).await?; + + script_store.build(script_db).await?; + + Ok(()) + } + + async fn handle_run( + _settings: &Settings, + run: Run, + script_db: atuin_scripts::database::Database, + ) -> Result<()> { + let script = script_db.get_by_name(&run.name).await?; + + if let Some(script) = script { + // Get variables used in the template + let variables = template_variables(&script)?; + + // Create a hashmap to store variable values + let mut variable_values: HashMap<String, serde_json::Value> = HashMap::new(); + + // Parse variables from command-line arguments first + for var_str in &run.var { + if let Some((key, value)) = var_str.split_once('=') { + // Add to variable values + variable_values.insert( + key.to_string(), + serde_json::Value::String(value.to_string()), + ); + debug!("Using CLI variable: {}={}", key, value); + } else { + eprintln!("Warning: Ignoring malformed variable specification: {var_str}"); + eprintln!("Variables should be specified as KEY=VALUE"); + } + } + + // Collect variables that are still needed (not specified via CLI) + let remaining_vars: HashSet<String> = variables + .into_iter() + .filter(|var| !variable_values.contains_key(var)) + .collect(); + + // If there are variables in the template that weren't specified on the command line, prompt for them + if !remaining_vars.is_empty() { + println!("This script contains template variables that need values:"); + + let stdin = std::io::stdin(); + let mut input = String::new(); + + for var in remaining_vars { + input.clear(); + + println!("Enter value for '{var}': "); + + if stdin.read_line(&mut input).is_err() { + eprintln!("Failed to read input for variable '{var}'"); + // Provide an empty string as fallback + variable_values.insert(var, serde_json::Value::String(String::new())); + continue; + } + + let value = input.trim().to_string(); + variable_values.insert(var, serde_json::Value::String(value)); + } + } + + let final_script = if variable_values.is_empty() { + // No variables to template, just use the original script + script.script.clone() + } else { + // If we have variables, we need to template the script + debug!("Templating script with variables: {:?}", variable_values); + template_script(&script, &variable_values)? + }; + + // Execute the script (either templated or original) + Self::execute_script(final_script, script.shebang.clone()).await?; + } else { + bail!("script not found"); + } + Ok(()) + } + + async fn handle_list( + _settings: &Settings, + _list: List, + script_db: atuin_scripts::database::Database, + ) -> Result<()> { + let scripts = script_db.list().await?; + + if scripts.is_empty() { + println!("No scripts found"); + } else { + println!("Available scripts:"); + for script in scripts { + if script.tags.is_empty() { + println!("- {} ", script.name); + } else { + println!("- {} [tags: {}]", script.name, script.tags.join(", ")); + } + + // Print description if it's not empty + if !script.description.is_empty() { + println!(" Description: {}", script.description); + } + } + } + + Ok(()) + } + + async fn handle_get( + _settings: &Settings, + get: Get, + script_db: atuin_scripts::database::Database, + ) -> Result<()> { + let script = script_db.get_by_name(&get.name).await?; + + if let Some(script) = script { + if get.script { + // Just print the executable script with shebang + print!( + "{}", + build_executable_script(script.script.clone(), script.shebang) + ); + return Ok(()); + } + + // Create a YAML representation of the script + println!("---"); + println!("name: {}", script.name); + println!("id: {}", script.id); + + if script.description.is_empty() { + println!("description: \"\""); + } else { + println!("description: |"); + // Indent multiline descriptions properly for YAML + for line in script.description.lines() { + println!(" {line}"); + } + } + + if script.tags.is_empty() { + println!("tags: []"); + } else { + println!("tags:"); + for tag in &script.tags { + println!(" - {tag}"); + } + } + + println!("shebang: {}", script.shebang); + + println!("script: |"); + // Indent the script content for proper YAML multiline format + for line in script.script.lines() { + println!(" {line}"); + } + + Ok(()) + } else { + bail!("script '{}' not found", get.name); + } + } + + #[allow(clippy::cognitive_complexity)] + async fn handle_edit( + _settings: &Settings, + edit: Edit, + script_store: ScriptStore, + script_db: atuin_scripts::database::Database, + ) -> Result<()> { + debug!("editing script {:?}", edit); + // Find the existing script + let existing_script = script_db.get_by_name(&edit.name).await?; + debug!("existing script {:?}", existing_script); + + if let Some(mut script) = existing_script { + // Update the script with new values if provided + if let Some(description) = edit.description { + script.description = description; + } + + // Handle renaming if requested + if let Some(new_name) = edit.rename { + // Check if a script with the new name already exists + if (script_db.get_by_name(&new_name).await?).is_some() { + bail!("A script named '{}' already exists", new_name); + } + + // Update the name + script.name = new_name; + } + + // Handle tag updates with priority: + // 1. If --no-tags is provided, clear all tags + // 2. If --tags is provided, replace all tags + // 3. If neither is provided, tags remain unchanged + if edit.no_tags { + // Clear all tags + script.tags.clear(); + } else if !edit.tags.is_empty() { + // Replace all tags + script.tags = edit.tags; + } + // If none of the above conditions are met, tags remain unchanged + + if let Some(shebang) = edit.shebang { + script.shebang = shebang; + } + + // Handle script content update + let script_content = if let Some(script_path) = edit.script { + // Load script from provided file + std::fs::read_to_string(script_path)? + } else if !edit.no_edit { + // Open the script in editor for interactive editing if --no-edit is not specified + Self::open_editor(Some(&script.script))? + } else { + // If --no-edit is specified, keep the existing script content + script.script.clone() + }; + + // Update the script content + script.script = script_content; + + // Update the script in the store + script_store.update(script).await?; + + // Rebuild the database to apply changes + script_store.build(script_db).await?; + + println!("Script '{}' updated successfully!", edit.name); + + Ok(()) + } else { + bail!("script '{}' not found", edit.name); + } + } + + async fn handle_delete( + _settings: &Settings, + delete: Delete, + script_store: ScriptStore, + script_db: atuin_scripts::database::Database, + ) -> Result<()> { + // Find the script by name + let script = script_db.get_by_name(&delete.name).await?; + + if let Some(script) = script { + // If not force, confirm deletion + if !delete.force { + println!( + "Are you sure you want to delete script '{}'? [y/N]", + delete.name + ); + let mut input = String::new(); + std::io::stdin().read_line(&mut input)?; + + let input = input.trim().to_lowercase(); + if input != "y" && input != "yes" { + println!("Deletion cancelled"); + return Ok(()); + } + } + + // Delete the script + script_store.delete(script.id).await?; + + // Rebuild the database to apply changes + script_store.build(script_db).await?; + + println!("Script '{}' deleted successfully", delete.name); + Ok(()) + } else { + bail!("script '{}' not found", delete.name); + } + } + + pub async fn run( + self, + settings: &Settings, + store: SqliteStore, + history_db: &impl Database, + ) -> Result<()> { + let host_id = Settings::host_id().await?; + let encryption_key: [u8; 32] = atuin_client::encryption::load_key(settings)?.into(); + + let script_store = ScriptStore::new(store, host_id, encryption_key); + let script_db = + atuin_scripts::database::Database::new(settings.scripts.db_path.clone(), 1.0).await?; + + match self { + Self::New(new_script) => { + Self::handle_new_script(settings, new_script, script_store, script_db, history_db) + .await + } + Self::Run(run) => Self::handle_run(settings, run, script_db).await, + Self::List(list) => Self::handle_list(settings, list, script_db).await, + Self::Get(get) => Self::handle_get(settings, get, script_db).await, + Self::Edit(edit) => Self::handle_edit(settings, edit, script_store, script_db).await, + Self::Delete(delete) => { + Self::handle_delete(settings, delete, script_store, script_db).await + } + } + } +} diff --git a/crates/atuin/src/command/client/search.rs b/crates/atuin/src/command/client/search.rs new file mode 100644 index 00000000..d05b1c24 --- /dev/null +++ b/crates/atuin/src/command/client/search.rs @@ -0,0 +1,362 @@ +use std::fs::File; +use std::io::{IsTerminal as _, Write, stderr, stdout}; + +use atuin_common::utils::{self, Escapable as _}; +use clap::Parser; +use eyre::Result; + +use atuin_client::{ + database::Database, + database::{OptFilters, current_context}, + encryption, + history::{History, store::HistoryStore}, + record::sqlite_store::SqliteStore, + settings::{FilterMode, KeymapMode, SearchMode, Settings, Timezone}, + theme::Theme, +}; + +use super::history::ListMode; + +mod cursor; +mod duration; +mod engines; +mod history_list; +mod inspector; +mod interactive; +pub mod keybindings; + +pub use duration::format_duration_into; + +#[allow(clippy::struct_excessive_bools, clippy::struct_field_names)] +#[derive(Parser, Debug)] +pub struct Cmd { + /// Filter search result by directory + #[arg(long, short)] + cwd: Option<String>, + + /// Exclude directory from results + #[arg(long = "exclude-cwd")] + exclude_cwd: Option<String>, + + /// Filter search result by exit code + #[arg(long, short)] + exit: Option<i64>, + + /// Exclude results with this exit code + #[arg(long = "exclude-exit")] + exclude_exit: Option<i64>, + + /// Only include results added before this date + #[arg(long, short)] + before: Option<String>, + + /// Only include results after this date + #[arg(long)] + after: Option<String>, + + /// How many entries to return at most + #[arg(long)] + limit: Option<i64>, + + /// Offset from the start of the results + #[arg(long)] + offset: Option<i64>, + + /// Open interactive search UI + #[arg(long, short)] + interactive: bool, + + /// Allow overriding filter mode over config + #[arg(long = "filter-mode")] + filter_mode: Option<FilterMode>, + + /// Allow overriding search mode over config + #[arg(long = "search-mode")] + search_mode: Option<SearchMode>, + + /// Marker argument used to inform atuin that it was invoked from a shell up-key binding (hidden from help to avoid confusion) + #[arg(long = "shell-up-key-binding", hide = true)] + shell_up_key_binding: bool, + + /// Notify the keymap at the shell's side + #[arg(long = "keymap-mode", default_value = "auto")] + keymap_mode: KeymapMode, + + /// Use human-readable formatting for time + #[arg(long)] + human: bool, + + #[arg(allow_hyphen_values = true)] + query: Option<Vec<String>>, + + /// Show only the text of the command + #[arg(long)] + cmd_only: bool, + + /// Terminate the output with a null, for better multiline handling + #[arg(long)] + print0: bool, + + /// Delete anything matching this query. Will not print out the match + #[arg(long)] + delete: bool, + + /// Delete EVERYTHING! + #[arg(long)] + delete_it_all: bool, + + /// Reverse the order of results, oldest first + #[arg(long, short)] + reverse: bool, + + /// Display the command time in another timezone other than the configured default. + /// + /// This option takes one of the following kinds of values: + /// - the special value "local" (or "l") which refers to the system time zone + /// - an offset from UTC (e.g. "+9", "-2:30") + #[arg(long, visible_alias = "tz")] + #[arg(allow_hyphen_values = true)] + // Clippy warns about `Option<Option<T>>`, but we suppress it because we need + // this distinction for proper argument handling. + #[allow(clippy::option_option)] + timezone: Option<Option<Timezone>>, + + /// Available variables: {command}, {directory}, {duration}, {user}, {host}, {time}, {exit} and + /// {relativetime}. + /// Example: --format "{time} - [{duration}] - {directory}$\t{command}" + #[arg(long, short)] + format: Option<String>, + + /// Set the maximum number of lines Atuin's interface should take up. + #[arg(long = "inline-height")] + inline_height: Option<u16>, + + /// Include duplicate commands in the output (non-interactive only) + #[arg(long)] + include_duplicates: bool, + + /// File name to write the result to (hidden from help as this is meant to be used from a script) + #[arg(long = "result-file", hide = true)] + result_file: Option<String>, +} + +impl Cmd { + /// Returns true if this search command will run in interactive (TUI) mode + pub fn is_interactive(&self) -> bool { + self.interactive + } + + // clippy: please write this instead + // clippy: now it has too many lines + // me: I'll do it later OKAY + #[allow(clippy::too_many_lines)] + pub async fn run( + self, + db: impl Database, + settings: &mut Settings, + store: SqliteStore, + theme: &Theme, + ) -> Result<()> { + let query = self.query.unwrap_or_else(|| { + std::env::var("ATUIN_QUERY").map_or_else( + |_| vec![], + |query| { + query + .split(' ') + .map(std::string::ToString::to_string) + .collect() + }, + ) + }); + + if (self.delete_it_all || self.delete) && self.limit.is_some() { + // Because of how deletion is implemented, it will always delete all matches + // and disregard the limit option. It is also not clear what deletion with a + // limit would even mean. Deleting the LIMIT most recent entries that match + // the search query would make sense, but that wouldn't match what's displayed + // when running the equivalent search, but deleting those entries that are + // displayed with the search would leave any duplicates of those lines which may + // or may not have been intended to be deleted. + eprintln!("\"--limit\" is not compatible with deletion."); + return Ok(()); + } + + if self.delete && query.is_empty() { + eprintln!( + "Please specify a query to match the items you wish to delete. If you wish to delete all history, pass --delete-it-all" + ); + return Ok(()); + } + + if self.delete_it_all && !query.is_empty() { + eprintln!( + "--delete-it-all will delete ALL of your history! It does not require a query." + ); + return Ok(()); + } + + if let Some(search_mode) = self.search_mode { + settings.search_mode = search_mode; + } + if let Some(filter_mode) = self.filter_mode { + settings.filter_mode = Some(filter_mode); + } + if let Some(inline_height) = self.inline_height { + settings.inline_height = inline_height; + } + + settings.shell_up_key_binding = self.shell_up_key_binding; + + // `keymap_mode` specified in config.toml overrides the `--keymap-mode` + // option specified in the keybindings. + settings.keymap_mode = match settings.keymap_mode { + KeymapMode::Auto => self.keymap_mode, + value => value, + }; + settings.keymap_mode_shell = self.keymap_mode; + + let encryption_key: [u8; 32] = encryption::load_key(settings)?.into(); + + let host_id = Settings::host_id().await?; + let history_store = HistoryStore::new(store.clone(), host_id, encryption_key); + + if self.interactive { + let item = interactive::history(&query, settings, db, &history_store, theme).await?; + + if let Some(result_file) = self.result_file { + let mut file = File::create(result_file)?; + write!(file, "{item}")?; + } else if !stdout().is_terminal() { + // stdout is not a terminal - likely command substitution like VAR=$(atuin search -i) + // Write to stdout so it gets captured + println!("{item}"); + } else if stderr().is_terminal() { + eprintln!("{}", item.escape_control()); + } else { + eprintln!("{item}"); + } + } else { + let opt_filter = OptFilters { + exit: self.exit, + exclude_exit: self.exclude_exit, + cwd: self.cwd, + exclude_cwd: self.exclude_cwd, + before: self.before, + after: self.after, + limit: self.limit, + offset: self.offset, + reverse: self.reverse, + include_duplicates: self.include_duplicates, + }; + + let mut entries = + run_non_interactive(settings, opt_filter.clone(), &query, &db).await?; + + if entries.is_empty() { + std::process::exit(1) + } + + // if we aren't deleting, print it all + if self.delete || self.delete_it_all { + // delete it + // it only took me _years_ to add this + // sorry + while !entries.is_empty() { + for entry in &entries { + eprintln!("deleting {}", entry.id); + + if settings.sync.records { + let (id, _) = history_store.delete(entry.id.clone()).await?; + history_store.incremental_build(&db, &[id]).await?; + } else { + db.delete(entry.clone()).await?; + } + } + + entries = + run_non_interactive(settings, opt_filter.clone(), &query, &db).await?; + } + } else { + let format = match self.format { + None => Some(settings.history_format.as_str()), + _ => self.format.as_deref(), + }; + let tz = match self.timezone { + Some(Some(tz)) => tz, // User provided a value + Some(None) | None => settings.timezone, // No value was provided + }; + + super::history::print_list( + &entries, + ListMode::from_flags(self.human, self.cmd_only), + format, + self.print0, + true, + tz, + ); + } + } + Ok(()) + } +} + +// This is supposed to more-or-less mirror the command line version, so ofc +// it is going to have a lot of args +#[allow(clippy::too_many_arguments, clippy::cast_possible_truncation)] +async fn run_non_interactive( + settings: &Settings, + filter_options: OptFilters, + query: &[String], + db: &impl Database, +) -> Result<Vec<History>> { + let dir = if filter_options.cwd.as_deref() == Some(".") { + Some(utils::get_current_dir()) + } else { + filter_options.cwd + }; + + let context = current_context().await?; + + let opt_filter = OptFilters { + cwd: dir.clone(), + ..filter_options + }; + + let filter_mode = settings.default_filter_mode(context.git_root.is_some()); + + let results = db + .search( + settings.search_mode, + filter_mode, + &context, + query.join(" ").as_str(), + opt_filter, + ) + .await?; + + Ok(results) +} + +#[cfg(test)] +mod tests { + use super::Cmd; + use clap::Parser; + + #[test] + fn search_for_triple_dash() { + // Issue #3028: searching for `---` should not be treated as a CLI flag + let cmd = Cmd::try_parse_from(["search", "---"]); + assert!(cmd.is_ok(), "Failed to parse '---' as a query: {cmd:?}"); + let cmd = cmd.unwrap(); + assert_eq!(cmd.query, Some(vec!["---".to_string()])); + } + + #[test] + fn search_for_double_dash_value() { + // Searching for strings starting with -- should also work + let cmd = Cmd::try_parse_from(["search", "--", "--foo"]); + assert!(cmd.is_ok()); + let cmd = cmd.unwrap(); + assert_eq!(cmd.query, Some(vec!["--foo".to_string()])); + } +} diff --git a/crates/atuin/src/command/client/search/cursor.rs b/crates/atuin/src/command/client/search/cursor.rs new file mode 100644 index 00000000..c1cdfee4 --- /dev/null +++ b/crates/atuin/src/command/client/search/cursor.rs @@ -0,0 +1,405 @@ +use atuin_client::settings::WordJumpMode; + +pub struct Cursor { + source: String, + index: usize, +} + +impl From<String> for Cursor { + fn from(source: String) -> Self { + Self { source, index: 0 } + } +} + +pub struct WordJumper<'a> { + word_chars: &'a str, + word_jump_mode: WordJumpMode, +} + +impl WordJumper<'_> { + fn is_word_boundary(&self, c: char, next_c: char) -> bool { + (c.is_whitespace() && !next_c.is_whitespace()) + || (!c.is_whitespace() && next_c.is_whitespace()) + || (self.word_chars.contains(c) && !self.word_chars.contains(next_c)) + || (!self.word_chars.contains(c) && self.word_chars.contains(next_c)) + } + + fn emacs_get_next_word_pos(&self, source: &str, index: usize) -> usize { + let index = (index + 1..source.len().saturating_sub(1)) + .find(|&i| self.word_chars.contains(source.chars().nth(i).unwrap())) + .unwrap_or(source.len()); + (index + 1..source.len().saturating_sub(1)) + .find(|&i| !self.word_chars.contains(source.chars().nth(i).unwrap())) + .unwrap_or(source.len()) + } + + fn emacs_get_prev_word_pos(&self, source: &str, index: usize) -> usize { + let index = (1..index) + .rev() + .find(|&i| self.word_chars.contains(source.chars().nth(i).unwrap())) + .unwrap_or(0); + (1..index) + .rev() + .find(|&i| !self.word_chars.contains(source.chars().nth(i).unwrap())) + .map_or(0, |i| i + 1) + } + + fn subl_get_next_word_pos(&self, source: &str, index: usize) -> usize { + let index = (index..source.len().saturating_sub(1)).find(|&i| { + self.is_word_boundary( + source.chars().nth(i).unwrap(), + source.chars().nth(i + 1).unwrap(), + ) + }); + if index.is_none() { + return source.len(); + } + (index.unwrap() + 1..source.len()) + .find(|&i| !source.chars().nth(i).unwrap().is_whitespace()) + .unwrap_or(source.len()) + } + + fn subl_get_prev_word_pos(&self, source: &str, index: usize) -> usize { + let index = (1..index) + .rev() + .find(|&i| !source.chars().nth(i).unwrap().is_whitespace()); + if index.is_none() { + return 0; + } + (1..index.unwrap()) + .rev() + .find(|&i| { + self.is_word_boundary( + source.chars().nth(i - 1).unwrap(), + source.chars().nth(i).unwrap(), + ) + }) + .unwrap_or(0) + } + + fn get_next_word_pos(&self, source: &str, index: usize) -> usize { + match self.word_jump_mode { + WordJumpMode::Emacs => self.emacs_get_next_word_pos(source, index), + WordJumpMode::Subl => self.subl_get_next_word_pos(source, index), + } + } + + fn get_prev_word_pos(&self, source: &str, index: usize) -> usize { + match self.word_jump_mode { + WordJumpMode::Emacs => self.emacs_get_prev_word_pos(source, index), + WordJumpMode::Subl => self.subl_get_prev_word_pos(source, index), + } + } +} + +impl Cursor { + pub fn as_str(&self) -> &str { + self.source.as_str() + } + + pub fn into_inner(self) -> String { + self.source + } + + /// Returns the string before the cursor + pub fn substring(&self) -> &str { + &self.source[..self.index] + } + + /// Returns the currently selected [`char`] + pub fn char(&self) -> Option<char> { + self.source[self.index..].chars().next() + } + + pub fn right(&mut self) { + if self.index < self.source.len() { + loop { + self.index += 1; + if self.source.is_char_boundary(self.index) { + break; + } + } + } + } + + pub fn left(&mut self) -> bool { + if self.index > 0 { + loop { + self.index -= 1; + if self.source.is_char_boundary(self.index) { + break true; + } + } + } else { + false + } + } + + pub fn next_word(&mut self, word_chars: &str, word_jump_mode: WordJumpMode) { + let word_jumper = WordJumper { + word_chars, + word_jump_mode, + }; + self.index = word_jumper.get_next_word_pos(&self.source, self.index); + } + + pub fn prev_word(&mut self, word_chars: &str, word_jump_mode: WordJumpMode) { + let word_jumper = WordJumper { + word_chars, + word_jump_mode, + }; + self.index = word_jumper.get_prev_word_pos(&self.source, self.index); + } + + /// Move cursor to the end of the current/next word (vim `e` motion). + /// + /// If cursor is in the middle of a word, moves to the end of that word. + /// If cursor is at the end of a word (or on whitespace), moves to the + /// end of the next word. + pub fn word_end(&mut self, word_chars: &str) { + let len = self.source.len(); + if self.index >= len { + return; + } + + let chars: Vec<char> = self.source.chars().collect(); + let mut char_idx = self.source[..self.index].chars().count(); + + if char_idx >= chars.len() { + return; + } + + let current = chars[char_idx]; + + // Check if we're at a word boundary (end of current word or on whitespace) + let at_word_boundary = current.is_whitespace() || char_idx + 1 >= chars.len() || { + let next = chars[char_idx + 1]; + next.is_whitespace() || (word_chars.contains(current) != word_chars.contains(next)) + }; + + // If at word boundary, advance past it and skip whitespace to find next word + if at_word_boundary { + char_idx += 1; + while char_idx < chars.len() && chars[char_idx].is_whitespace() { + char_idx += 1; + } + } + + // If we've gone past end, go to end of string + if char_idx >= chars.len() { + self.index = len; + return; + } + + // Find end of word: advance until next char is whitespace or different word type + let in_word_chars = word_chars.contains(chars[char_idx]); + while char_idx < chars.len() { + let next_idx = char_idx + 1; + if next_idx >= chars.len() { + // At last char, move past it + char_idx = next_idx; + break; + } + let next_c = chars[next_idx]; + if next_c.is_whitespace() || (word_chars.contains(next_c) != in_word_chars) { + // Next char is start of new word/whitespace, so current char is end + char_idx = next_idx; + break; + } + char_idx += 1; + } + + // Convert char index back to byte index + self.index = chars.iter().take(char_idx).map(|c| c.len_utf8()).sum(); + } + + pub fn insert(&mut self, c: char) { + self.source.insert(self.index, c); + self.index += c.len_utf8(); + } + + pub fn remove(&mut self) -> Option<char> { + if self.index < self.source.len() { + Some(self.source.remove(self.index)) + } else { + None + } + } + + pub fn remove_next_word(&mut self, word_chars: &str, word_jump_mode: WordJumpMode) { + let word_jumper = WordJumper { + word_chars, + word_jump_mode, + }; + let next_index = word_jumper.get_next_word_pos(&self.source, self.index); + self.source.replace_range(self.index..next_index, ""); + } + + pub fn remove_prev_word(&mut self, word_chars: &str, word_jump_mode: WordJumpMode) { + let word_jumper = WordJumper { + word_chars, + word_jump_mode, + }; + let next_index = word_jumper.get_prev_word_pos(&self.source, self.index); + self.source.replace_range(next_index..self.index, ""); + self.index = next_index; + } + + pub fn back(&mut self) -> Option<char> { + if self.left() { self.remove() } else { None } + } + + pub fn clear(&mut self) { + self.source.clear(); + self.index = 0; + } + + pub fn clear_to_start(&mut self) { + self.source.replace_range(..self.index, ""); + self.index = 0; + } + + pub fn clear_to_end(&mut self) { + self.source.replace_range(self.index.., ""); + self.index = self.source.len(); + } + + pub fn end(&mut self) { + self.index = self.source.len(); + } + + pub fn start(&mut self) { + self.index = 0; + } + + pub fn position(&self) -> usize { + self.index + } +} + +#[cfg(test)] +mod cursor_tests { + use super::Cursor; + use super::*; + + static EMACS_WORD_JUMPER: WordJumper = WordJumper { + word_chars: "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789", + word_jump_mode: WordJumpMode::Emacs, + }; + + static SUBL_WORD_JUMPER: WordJumper = WordJumper { + word_chars: "./\\()\"'-:,.;<>~!@#$%^&*|+=[]{}`~?", + word_jump_mode: WordJumpMode::Subl, + }; + + #[test] + fn right() { + // ö is 2 bytes + let mut c = Cursor::from(String::from("öaöböcödöeöfö")); + let indices = [0, 2, 3, 5, 6, 8, 9, 11, 12, 14, 15, 17, 18, 20, 20, 20, 20]; + for i in indices { + assert_eq!(c.index, i); + c.right(); + } + } + + #[test] + fn left() { + // ö is 2 bytes + let mut c = Cursor::from(String::from("öaöböcödöeöfö")); + c.end(); + let indices = [20, 18, 17, 15, 14, 12, 11, 9, 8, 6, 5, 3, 2, 0, 0, 0, 0]; + for i in indices { + assert_eq!(c.index, i); + c.left(); + } + } + + #[test] + fn test_emacs_get_next_word_pos() { + let s = String::from(" aaa ((()))bbb ((())) "); + let indices = [(0, 6), (3, 6), (7, 18), (19, 30)]; + for (i_src, i_dest) in indices { + assert_eq!(EMACS_WORD_JUMPER.get_next_word_pos(&s, i_src), i_dest); + } + assert_eq!(EMACS_WORD_JUMPER.get_next_word_pos("", 0), 0); + } + + #[test] + fn test_emacs_get_prev_word_pos() { + let s = String::from(" aaa ((()))bbb ((())) "); + let indices = [(30, 15), (29, 15), (15, 3), (3, 0)]; + for (i_src, i_dest) in indices { + assert_eq!(EMACS_WORD_JUMPER.get_prev_word_pos(&s, i_src), i_dest); + } + assert_eq!(EMACS_WORD_JUMPER.get_prev_word_pos("", 0), 0); + } + + #[test] + fn test_subl_get_next_word_pos() { + let s = String::from(" aaa ((()))bbb ((())) "); + let indices = [(0, 3), (1, 3), (3, 9), (9, 15), (15, 21), (21, 30)]; + for (i_src, i_dest) in indices { + assert_eq!(SUBL_WORD_JUMPER.get_next_word_pos(&s, i_src), i_dest); + } + assert_eq!(SUBL_WORD_JUMPER.get_next_word_pos("", 0), 0); + } + + #[test] + fn test_subl_get_prev_word_pos() { + let s = String::from(" aaa ((()))bbb ((())) "); + let indices = [(30, 21), (21, 15), (15, 9), (9, 3), (3, 0)]; + for (i_src, i_dest) in indices { + assert_eq!(SUBL_WORD_JUMPER.get_prev_word_pos(&s, i_src), i_dest); + } + assert_eq!(SUBL_WORD_JUMPER.get_prev_word_pos("", 0), 0); + } + + #[test] + fn pop() { + let mut s = String::from("öaöböcödöeöfö"); + let mut c = Cursor::from(s.clone()); + c.end(); + while !s.is_empty() { + let c1 = s.pop(); + let c2 = c.back(); + assert_eq!(c1, c2); + assert_eq!(s.as_str(), c.substring()); + } + let c1 = s.pop(); + let c2 = c.back(); + assert_eq!(c1, c2); + } + + #[test] + fn back() { + let mut c = Cursor::from(String::from("öaöböcödöeöfö")); + // move to ^ + for _ in 0..4 { + c.right(); + } + assert_eq!(c.substring(), "öaöb"); + assert_eq!(c.back(), Some('b')); + assert_eq!(c.back(), Some('ö')); + assert_eq!(c.back(), Some('a')); + assert_eq!(c.back(), Some('ö')); + assert_eq!(c.back(), None); + assert_eq!(c.as_str(), "öcödöeöfö"); + } + + #[test] + fn insert() { + let mut c = Cursor::from(String::from("öaöböcödöeöfö")); + // move to ^ + for _ in 0..4 { + c.right(); + } + assert_eq!(c.substring(), "öaöb"); + c.insert('ö'); + c.insert('g'); + c.insert('ö'); + c.insert('h'); + assert_eq!(c.substring(), "öaöbögöh"); + assert_eq!(c.as_str(), "öaöbögöhöcödöeöfö"); + } +} diff --git a/crates/atuin/src/command/client/search/duration.rs b/crates/atuin/src/command/client/search/duration.rs new file mode 100644 index 00000000..dfa9426b --- /dev/null +++ b/crates/atuin/src/command/client/search/duration.rs @@ -0,0 +1,65 @@ +use core::fmt; +use std::{ops::ControlFlow, time::Duration}; + +#[allow(clippy::module_name_repetitions)] +pub fn format_duration_into(dur: Duration, f: &mut fmt::Formatter<'_>) -> fmt::Result { + fn item(unit: &'static str, value: u64) -> ControlFlow<(&'static str, u64)> { + if value > 0 { + ControlFlow::Break((unit, value)) + } else { + ControlFlow::Continue(()) + } + } + + // impl taken and modified from + // https://github.com/tailhook/humantime/blob/master/src/duration.rs#L295-L331 + // Copyright (c) 2016 The humantime Developers + fn fmt(f: Duration) -> ControlFlow<(&'static str, u64), ()> { + let secs = f.as_secs(); + let nanos = f.subsec_nanos(); + + let years = secs / 31_557_600; // 365.25d + let year_days = secs % 31_557_600; + let months = year_days / 2_630_016; // 30.44d + let month_days = year_days % 2_630_016; + let days = month_days / 86400; + let day_secs = month_days % 86400; + let hours = day_secs / 3600; + let minutes = day_secs % 3600 / 60; + let seconds = day_secs % 60; + + let millis = nanos / 1_000_000; + let micros = nanos / 1_000; + + // a difference from our impl than the original is that + // we only care about the most-significant segment of the duration. + // If the item call returns `Break`, then the `?` will early-return. + // This allows for a very consise impl + item("y", years)?; + item("mo", months)?; + item("d", days)?; + item("h", hours)?; + item("m", minutes)?; + item("s", seconds)?; + item("ms", u64::from(millis))?; + item("us", u64::from(micros))?; + item("ns", u64::from(nanos))?; + ControlFlow::Continue(()) + } + + match fmt(dur) { + ControlFlow::Break((unit, value)) => write!(f, "{value}{unit}"), + ControlFlow::Continue(()) => write!(f, "0s"), + } +} + +#[allow(clippy::module_name_repetitions)] +pub fn format_duration(f: Duration) -> String { + struct F(Duration); + impl fmt::Display for F { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + format_duration_into(self.0, f) + } + } + F(f).to_string() +} diff --git a/crates/atuin/src/command/client/search/engines.rs b/crates/atuin/src/command/client/search/engines.rs new file mode 100644 index 00000000..8cbee0c3 --- /dev/null +++ b/crates/atuin/src/command/client/search/engines.rs @@ -0,0 +1,85 @@ +use async_trait::async_trait; +use atuin_client::{ + database::{Context, Database}, + history::{History, HistoryId}, + settings::{FilterMode, SearchMode, Settings}, +}; +use eyre::Result; + +use super::cursor::Cursor; + +#[cfg(feature = "daemon")] +pub mod daemon; +pub mod db; +pub mod skim; + +#[allow(unused)] // settings is only used if daemon feature is enabled +pub fn engine(search_mode: SearchMode, settings: &Settings) -> Box<dyn SearchEngine> { + match search_mode { + SearchMode::Skim => Box::new(skim::Search::new()) as Box<_>, + #[cfg(feature = "daemon")] + SearchMode::DaemonFuzzy => Box::new(daemon::Search::new(settings)) as Box<_>, + #[cfg(not(feature = "daemon"))] + SearchMode::DaemonFuzzy => { + // Fall back to fuzzy mode if daemon feature is not enabled + Box::new(db::Search(SearchMode::Fuzzy)) as Box<_> + } + mode => Box::new(db::Search(mode)) as Box<_>, + } +} + +pub struct SearchState { + pub input: Cursor, + pub filter_mode: FilterMode, + pub context: Context, + pub custom_context: Option<HistoryId>, +} + +impl SearchState { + pub(crate) fn rotate_filter_mode(&mut self, settings: &Settings, offset: isize) { + let mut i = settings + .search + .filters + .iter() + .position(|&m| m == self.filter_mode) + .unwrap_or_default(); + for _ in 0..settings.search.filters.len() { + i = (i.wrapping_add_signed(offset)) % settings.search.filters.len(); + let mode = settings.search.filters[i]; + if self.filter_mode_available(mode, settings) { + self.filter_mode = mode; + break; + } + } + } + + fn filter_mode_available(&self, mode: FilterMode, settings: &Settings) -> bool { + match mode { + FilterMode::Global | FilterMode::SessionPreload => self.custom_context.is_none(), + FilterMode::Workspace => settings.workspaces && self.context.git_root.is_some(), + _ => true, + } + } +} + +#[async_trait] +pub trait SearchEngine: Send + Sync + 'static { + async fn full_query( + &mut self, + state: &SearchState, + db: &mut dyn Database, + ) -> Result<Vec<History>>; + + async fn query(&mut self, state: &SearchState, db: &mut dyn Database) -> Result<Vec<History>> { + if state.input.as_str().is_empty() { + Ok(db + .list(&[state.filter_mode], &state.context, Some(200), true, false) + .await? + .into_iter() + .collect::<Vec<_>>()) + } else { + self.full_query(state, db).await + } + } + fn get_highlight_indices(&self, command: &str, search_input: &str) -> Vec<usize>; +} diff --git a/crates/atuin/src/command/client/search/engines/daemon.rs b/crates/atuin/src/command/client/search/engines/daemon.rs new file mode 100644 index 00000000..d317a4f6 --- /dev/null +++ b/crates/atuin/src/command/client/search/engines/daemon.rs @@ -0,0 +1,206 @@ +use async_trait::async_trait; +use atuin_client::{ + database::{Database, OptFilters}, + history::History, + settings::{SearchMode, Settings}, +}; +use atuin_daemon::client::SearchClient; +use eyre::Result; +use nucleo_matcher::{ + Config, Matcher, Utf32Str, + pattern::{CaseMatching, Normalization, Pattern}, +}; +use tracing::{Level, debug, instrument, span}; +use uuid::Uuid; + +use super::{SearchEngine, SearchState}; + +pub struct Search { + client: Option<SearchClient>, + query_id: u64, + socket_path: String, + #[cfg(not(unix))] + tcp_port: u64, +} + +impl Search { + pub fn new(settings: &Settings) -> Self { + Search { + client: None, + query_id: 0, + socket_path: settings.daemon.socket_path.clone(), + #[cfg(not(unix))] + tcp_port: settings.daemon.tcp_port, + } + } + + #[instrument(skip_all, level = Level::TRACE, name = "get_daemon_client")] + async fn get_client(&mut self) -> Result<&mut SearchClient> { + if self.client.is_none() { + #[cfg(unix)] + let client = SearchClient::new(self.socket_path.clone()).await?; + + #[cfg(not(unix))] + let client = SearchClient::new(self.tcp_port).await?; + + self.client = Some(client); + } + Ok(self.client.as_mut().unwrap()) + } + + fn next_query_id(&mut self) -> u64 { + self.query_id += 1; + self.query_id + } + + /// Check if query contains regex pattern (r/.../) + /// Nucleo doesn't support regex, so we fall back to database search + fn contains_regex_pattern(query: &str) -> bool { + query.starts_with("r/") || query.contains(" r/") + } + + #[instrument(skip_all, level = Level::TRACE, name = "daemon_db_fallback")] + async fn fallback_to_db_search( + &self, + state: &SearchState, + db: &dyn Database, + ) -> Result<Vec<History>> { + let results = db + .search( + SearchMode::FullText, + state.filter_mode, + &state.context, + state.input.as_str(), + OptFilters { + limit: Some(200), + ..Default::default() + }, + ) + .await + .map_or(Vec::new(), |r| r.into_iter().collect()); + Ok(results) + } + + #[instrument(skip_all, level = Level::TRACE, name = "hydrate_from_db", fields(count = ids.len()))] + async fn hydrate_from_db(&self, db: &dyn Database, ids: &[String]) -> Result<Vec<History>> { + let placeholders: Vec<String> = ids.iter().map(|id| format!("'{id}'")).collect(); + let sql_query = format!( + "SELECT * FROM history WHERE id IN ({}) ORDER BY timestamp DESC", + placeholders.join(",") + ); + Ok(db.query_history(&sql_query).await?) + } +} + +#[async_trait] +impl SearchEngine for Search { + #[instrument(skip_all, level = Level::TRACE, name = "daemon_search", fields(query = %state.input.as_str()))] + async fn full_query( + &mut self, + state: &SearchState, + db: &mut dyn Database, + ) -> Result<Vec<History>> { + let query = state.input.as_str().to_string(); + + // Fall back to database for regex queries (Nucleo doesn't support regex) + if Self::contains_regex_pattern(&query) { + debug!(query = %query, "[daemon-client] regex detected, falling back to db"); + return self.fallback_to_db_search(state, db).await; + } + + let query_id = self.next_query_id(); + + let span = + span!(Level::TRACE, "daemon_search.req_resp", query = %query, query_id = query_id); + + let client = self.get_client().await?; + + let _span = span.enter(); + let mut stream = client + .search( + query.clone(), + query_id, + state.filter_mode, + Some(state.context.clone()), + ) + .await?; + + let mut ids = Vec::with_capacity(200); + span!(Level::TRACE, "daemon_search.resp") + .in_scope(async || { + while let Ok(Some(response)) = stream.message().await { + let span2 = span!( + Level::TRACE, + "daemon_search.resp.item", + query_id = response.query_id + ); + let _span2 = span2.enter(); + // Only process if the query_id matches (prevents stale responses) + if response.query_id == query_id { + let uuids = response + .ids + .iter() + .map(|id| { + let bytes: [u8; 16] = + id.as_slice().try_into().expect("id should be 16 bytes"); + Uuid::from_bytes(bytes).as_simple().to_string() + }) + .collect::<Vec<_>>(); + ids.extend(uuids); + } + drop(_span2); + drop(span2); + } + }) + .await; + drop(_span); + drop(span); + + if ids.is_empty() { + debug!(query = %query, results = 0, "[daemon-client] empty results"); + return Ok(Vec::new()); + } + + // // Hydrate from local database + let results = self.hydrate_from_db(db, &ids).await?; + + // // Reorder results to match the order from the daemon (which is ranked by relevance) + let ordered_results = span!(Level::TRACE, "reorder_results").in_scope(|| { + let mut ordered_results = Vec::with_capacity(results.len()); + for id in &ids { + if let Some(history) = results.iter().find(|h| h.id.0 == *id) { + ordered_results.push(history.clone()); + } + } + ordered_results + }); + + debug!( + query = %query, + results = results.len(), + "[daemon-client]" + ); + + Ok(ordered_results) + } + + #[instrument(skip_all, level = Level::TRACE, name = "daemon_highlight")] + fn get_highlight_indices(&self, command: &str, search_input: &str) -> Vec<usize> { + // Use fulltext highlighting for regex queries + if Self::contains_regex_pattern(search_input) { + return super::db::get_highlight_indices_fulltext(command, search_input); + } + + let mut matcher = Matcher::new(Config::DEFAULT); + let pattern = Pattern::parse(search_input, CaseMatching::Smart, Normalization::Smart); + + let mut indices: Vec<u32> = Vec::new(); + let mut haystack_buf = Vec::new(); + + let haystack = Utf32Str::new(command, &mut haystack_buf); + pattern.indices(haystack, &mut matcher, &mut indices); + + // Convert u32 indices to usize + indices.into_iter().map(|i| i as usize).collect() + } +} diff --git a/crates/atuin/src/command/client/search/engines/db.rs b/crates/atuin/src/command/client/search/engines/db.rs new file mode 100644 index 00000000..476462f5 --- /dev/null +++ b/crates/atuin/src/command/client/search/engines/db.rs @@ -0,0 +1,109 @@ +use super::{SearchEngine, SearchState}; +use async_trait::async_trait; +use atuin_client::{ + database::Database, + database::OptFilters, + database::{QueryToken, QueryTokenizer}, + history::History, + settings::SearchMode, +}; +use eyre::Result; +use norm::Metric; +use norm::fzf::{FzfParser, FzfV2}; +use std::ops::Range; +use tracing::{Level, instrument}; + +pub struct Search(pub SearchMode); + +#[async_trait] +impl SearchEngine for Search { + #[instrument(skip_all, level = Level::TRACE, name = "db_search", fields(mode = ?self.0, query = %state.input.as_str()))] + async fn full_query( + &mut self, + state: &SearchState, + db: &mut dyn Database, + ) -> Result<Vec<History>> { + let results = db + .search( + self.0, + state.filter_mode, + &state.context, + state.input.as_str(), + OptFilters { + limit: Some(200), + ..Default::default() + }, + ) + .await + // ignore errors as it may be caused by incomplete regex + .map_or(Vec::new(), |r| r.into_iter().collect()); + Ok(results) + } + + #[instrument(skip_all, level = Level::TRACE, name = "db_highlight")] + fn get_highlight_indices(&self, command: &str, search_input: &str) -> Vec<usize> { + if self.0 == SearchMode::Prefix { + return vec![]; + } else if self.0 == SearchMode::FullText { + return get_highlight_indices_fulltext(command, search_input); + } + let mut fzf = FzfV2::new(); + let mut parser = FzfParser::new(); + let query = parser.parse(search_input); + let mut ranges: Vec<Range<usize>> = Vec::new(); + let _ = fzf.distance_and_ranges(query, command, &mut ranges); + + // convert ranges to all indices + ranges.into_iter().flatten().collect() + } +} + +#[instrument(skip_all, level = Level::TRACE, name = "db_highlight_fulltext")] +pub fn get_highlight_indices_fulltext(command: &str, search_input: &str) -> Vec<usize> { + let mut ranges = vec![]; + let lower_command = command.to_ascii_lowercase(); + + for token in QueryTokenizer::new(search_input) { + let matchee = if token.has_uppercase() { + command + } else { + &lower_command + }; + + if token.is_inverse() { + continue; + } + + match token { + QueryToken::Or => {} + QueryToken::Regex(r) => { + if let Ok(re) = regex::Regex::new(r) { + for m in re.find_iter(command) { + ranges.push(m.range()); + } + } + } + QueryToken::MatchStart(term, _) => { + if matchee.starts_with(term) { + ranges.push(0..term.len()); + } + } + QueryToken::MatchEnd(term, _) => { + if matchee.ends_with(term) { + let l = matchee.len(); + ranges.push((l - term.len())..l); + } + } + QueryToken::Match(term, _) | QueryToken::MatchFull(term, _) => { + for (idx, m) in matchee.match_indices(term) { + ranges.push(idx..(idx + m.len())); + } + } + } + } + + let mut ret: Vec<_> = ranges.into_iter().flatten().collect(); + ret.sort_unstable(); + ret.dedup(); + ret +} diff --git a/crates/atuin/src/command/client/search/engines/skim.rs b/crates/atuin/src/command/client/search/engines/skim.rs new file mode 100644 index 00000000..7d9feb40 --- /dev/null +++ b/crates/atuin/src/command/client/search/engines/skim.rs @@ -0,0 +1,222 @@ +use std::path::Path; + +use async_trait::async_trait; +use atuin_client::{database::Database, history::History, settings::FilterMode}; +use eyre::Result; +use fuzzy_matcher::{FuzzyMatcher, skim::SkimMatcherV2}; +use itertools::Itertools; +use time::OffsetDateTime; +use tokio::task::yield_now; +use tracing::{Level, instrument, warn}; +use uuid; + +use super::{SearchEngine, SearchState}; + +pub struct Search { + all_history: Vec<(History, i32)>, + engine: SkimMatcherV2, +} + +impl Search { + pub fn new() -> Self { + Search { + all_history: vec![], + engine: SkimMatcherV2::default(), + } + } +} + +#[async_trait] +impl SearchEngine for Search { + #[instrument(skip_all, level = Level::TRACE, name = "skim_search", fields(query = %state.input.as_str()))] + async fn full_query( + &mut self, + state: &SearchState, + db: &mut dyn Database, + ) -> Result<Vec<History>> { + if self.all_history.is_empty() { + self.all_history = load_all_history(db).await; + } + + Ok(fuzzy_search(&self.engine, state, &self.all_history).await) + } + + #[instrument(skip_all, level = Level::TRACE, name = "skim_highlight")] + fn get_highlight_indices(&self, command: &str, search_input: &str) -> Vec<usize> { + let (_, indices) = self + .engine + .fuzzy_indices(command, search_input) + .unwrap_or_default(); + indices + } +} + +#[instrument(skip_all, level = Level::TRACE, name = "load_all_history")] +async fn load_all_history(db: &dyn Database) -> Vec<(History, i32)> { + db.all_with_count().await.unwrap() +} + +#[allow(clippy::too_many_lines)] +#[instrument(skip_all, level = Level::TRACE, name = "fuzzy_match", fields(history_count = all_history.len()))] +async fn fuzzy_search( + engine: &SkimMatcherV2, + state: &SearchState, + all_history: &[(History, i32)], +) -> Vec<History> { + let mut set = Vec::with_capacity(200); + let mut ranks = Vec::with_capacity(200); + let query = state.input.as_str(); + let now = OffsetDateTime::now_utc(); + + for (i, (history, count)) in all_history.iter().enumerate() { + if i % 256 == 0 { + yield_now().await; + } + let context = &state.context; + let git_root = context + .git_root + .as_ref() + .and_then(|git_root| git_root.to_str()) + .unwrap_or(&context.cwd); + match state.filter_mode { + FilterMode::Global => {} + // we aggregate host by ',' separating them + FilterMode::Host + if history + .hostname + .split(',') + .contains(&context.hostname.as_str()) => {} + // we aggregate session by concattenating them. + // sessions are 32 byte simple uuid formats + FilterMode::Session + if history + .session + .as_bytes() + .chunks(32) + .contains(&context.session.as_bytes()) => {} + // SessionPreload: include current session + global history from before session start + FilterMode::SessionPreload => { + let is_current_session = { + history + .session + .as_bytes() + .chunks(32) + .any(|chunk| chunk == context.session.as_bytes()) + }; + + if !is_current_session { + let Ok(uuid) = uuid::Uuid::parse_str(&context.session) else { + warn!("failed to parse session id '{}'", context.session); + continue; + }; + let Some(timestamp) = uuid.get_timestamp() else { + warn!( + "failed to get timestamp from uuid '{}'", + uuid.as_hyphenated() + ); + continue; + }; + let (seconds, nanos) = timestamp.to_unix(); + let Ok(session_start) = time::OffsetDateTime::from_unix_timestamp_nanos( + i128::from(seconds) * 1_000_000_000 + i128::from(nanos), + ) else { + warn!( + "failed to create OffsetDateTime from second: {seconds}, nanosecond: {nanos}" + ); + continue; + }; + + if history.timestamp >= session_start { + continue; + } + } + } + // we aggregate directory by ':' separating them + FilterMode::Directory if history.cwd.split(':').contains(&context.cwd.as_str()) => {} + FilterMode::Workspace if history.cwd.split(':').contains(&git_root) => {} + _ => continue, + } + #[allow(clippy::cast_lossless, clippy::cast_precision_loss)] + if let Some((score, indices)) = engine.fuzzy_indices(&history.command, query) { + let begin = indices.first().copied().unwrap_or_default(); + + let mut duration = (now - history.timestamp).as_seconds_f64().log2(); + if !duration.is_finite() || duration <= 1.0 { + duration = 1.0; + } + // these + X.0 just make the log result a bit smoother. + // log is very spiky towards 1-4, but I want a gradual decay. + // eg: + // log2(4) = 2, log2(5) = 2.3 (16% increase) + // log2(8) = 3, log2(9) = 3.16 (5% increase) + // log2(16) = 4, log2(17) = 4.08 (2% increase) + let count = (*count as f64 + 8.0).log2(); + let begin = (begin as f64 + 16.0).log2(); + let path = path_dist(history.cwd.as_ref(), state.context.cwd.as_ref()); + let path = (path as f64 + 8.0).log2(); + + // reduce longer durations, raise higher counts, raise matches close to the start + let score = (-score as f64) * count / path / duration / begin; + + 'insert: { + // algorithm: + // 1. find either the position that this command ranks + // 2. find the same command positioned better than our rank. + for i in 0..set.len() { + // do we out score the current position? + if ranks[i] > score { + ranks.insert(i, score); + set.insert(i, history.clone()); + let mut j = i + 1; + while j < set.len() { + // remove duplicates that have a worse score + if set[j].command == history.command { + ranks.remove(j); + set.remove(j); + + // break this while loop because there won't be any other + // duplicates. + break; + } + j += 1; + } + + // keep it limited + if ranks.len() > 200 { + ranks.pop(); + set.pop(); + } + + break 'insert; + } + // don't continue if this command has a better score already + if set[i].command == history.command { + break 'insert; + } + } + + if set.len() < 200 { + ranks.push(score); + set.push(history.clone()); + } + } + } + } + + set +} + +fn path_dist(a: &Path, b: &Path) -> usize { + let mut a: Vec<_> = a.components().collect(); + let b: Vec<_> = b.components().collect(); + + let mut dist = 0; + + // pop a until there's a common ancestor + while !b.starts_with(&a) { + dist += 1; + a.pop(); + } + + b.len() - a.len() + dist +} diff --git a/crates/atuin/src/command/client/search/history_list.rs b/crates/atuin/src/command/client/search/history_list.rs new file mode 100644 index 00000000..23c23548 --- /dev/null +++ b/crates/atuin/src/command/client/search/history_list.rs @@ -0,0 +1,429 @@ +use std::time::Duration; + +use super::duration::format_duration; +use super::engines::SearchEngine; +use atuin_client::{ + history::History, + settings::{UiColumn, UiColumnType}, + theme::{Meaning, Theme}, +}; +use atuin_common::utils::Escapable as _; +use itertools::Itertools; +use ratatui::{ + backend::FromCrossterm, + buffer::Buffer, + crossterm::style, + layout::Rect, + style::{Modifier, Style}, + widgets::{Block, StatefulWidget, Widget}, +}; +use time::OffsetDateTime; + +pub struct HistoryHighlighter<'a> { + pub engine: &'a dyn SearchEngine, + pub search_input: &'a str, +} + +impl HistoryHighlighter<'_> { + pub fn get_highlight_indices(&self, command: &str) -> Vec<usize> { + self.engine + .get_highlight_indices(command, self.search_input) + } +} + +pub struct HistoryList<'a> { + history: &'a [History], + block: Option<Block<'a>>, + inverted: bool, + /// Apply an alternative highlighting to the selected row + alternate_highlight: bool, + now: &'a dyn Fn() -> OffsetDateTime, + indicator: &'a str, + theme: &'a Theme, + history_highlighter: HistoryHighlighter<'a>, + show_numeric_shortcuts: bool, + /// Columns to display (in order, after the indicator) + columns: &'a [UiColumn], +} + +#[derive(Default)] +pub struct ListState { + offset: usize, + selected: usize, + max_entries: usize, +} + +impl ListState { + pub fn selected(&self) -> usize { + self.selected + } + + pub fn max_entries(&self) -> usize { + self.max_entries + } + + pub fn offset(&self) -> usize { + self.offset + } + + pub fn select(&mut self, index: usize) { + self.selected = index; + } +} + +impl StatefulWidget for HistoryList<'_> { + type State = ListState; + + fn render(mut self, area: Rect, buf: &mut Buffer, state: &mut Self::State) { + let list_area = self.block.take().map_or(area, |b| { + let inner_area = b.inner(area); + b.render(area, buf); + inner_area + }); + + if list_area.width < 1 || list_area.height < 1 || self.history.is_empty() { + return; + } + let list_height = list_area.height as usize; + + let (start, end) = self.get_items_bounds(state.selected, state.offset, list_height); + state.offset = start; + state.max_entries = end - start; + + let mut s = DrawState { + buf, + list_area, + x: 0, + y: 0, + state, + inverted: self.inverted, + alternate_highlight: self.alternate_highlight, + now: &self.now, + indicator: self.indicator, + theme: self.theme, + history_highlighter: self.history_highlighter, + show_numeric_shortcuts: self.show_numeric_shortcuts, + columns: self.columns, + }; + + for item in self.history.iter().skip(state.offset).take(end - start) { + s.render_row(item); + + // reset line + s.y += 1; + s.x = 0; + } + } +} + +impl<'a> HistoryList<'a> { + #[allow(clippy::too_many_arguments)] + pub fn new( + history: &'a [History], + inverted: bool, + alternate_highlight: bool, + now: &'a dyn Fn() -> OffsetDateTime, + indicator: &'a str, + theme: &'a Theme, + history_highlighter: HistoryHighlighter<'a>, + show_numeric_shortcuts: bool, + columns: &'a [UiColumn], + ) -> Self { + Self { + history, + block: None, + inverted, + alternate_highlight, + now, + indicator, + theme, + history_highlighter, + show_numeric_shortcuts, + columns, + } + } + + pub fn block(mut self, block: Block<'a>) -> Self { + self.block = Some(block); + self + } + + fn get_items_bounds(&self, selected: usize, offset: usize, height: usize) -> (usize, usize) { + let offset = offset.min(self.history.len().saturating_sub(1)); + + let max_scroll_space = height.min(10).min(self.history.len() - selected); + if offset + height < selected + max_scroll_space { + let end = selected + max_scroll_space; + (end - height, end) + } else if selected < offset { + (selected, selected + height) + } else { + (offset, offset + height) + } + } +} + +struct DrawState<'a> { + buf: &'a mut Buffer, + list_area: Rect, + x: u16, + y: u16, + state: &'a ListState, + inverted: bool, + alternate_highlight: bool, + now: &'a dyn Fn() -> OffsetDateTime, + indicator: &'a str, + theme: &'a Theme, + history_highlighter: HistoryHighlighter<'a>, + show_numeric_shortcuts: bool, + columns: &'a [UiColumn], +} + +// these encode the slices of `" > "`, `" {n} "`, or `" "` in a compact form. +// Yes, this is a hack, but it makes me feel happy +static SLICES: &str = " > 1 2 3 4 5 6 7 8 9 "; + +impl DrawState<'_> { + /// Render a complete row for a history item based on configured columns. + fn render_row(&mut self, h: &History) { + // Always render the indicator first (width 3) + self.index(); + + // Calculate the width for the expanding column + // Fixed columns use their configured width + 1 (trailing space) + let indicator_width: u16 = 3; + let fixed_width: u16 = self + .columns + .iter() + .filter(|c| !c.expand) + .map(|c| c.width + 1) + .sum(); + let expand_width = self + .list_area + .width + .saturating_sub(indicator_width + fixed_width); + + let style = self.theme.as_style(Meaning::Base); + // Render each configured column + for (idx, column) in self.columns.iter().enumerate() { + if idx != 0 { + self.draw(" ", Style::from_crossterm(style)); + } + let width = if column.expand { + expand_width + } else { + column.width + }; + match column.column_type { + UiColumnType::Duration => self.duration(h, width), + UiColumnType::Time => self.time(h, width), + UiColumnType::Datetime => self.datetime(h, width), + UiColumnType::Directory => self.directory(h, width), + UiColumnType::Host => self.host(h, width), + UiColumnType::User => self.user(h, width), + UiColumnType::Exit => self.exit_code(h, width), + UiColumnType::Command => self.command(h), + } + } + } + + fn index(&mut self) { + if !self.show_numeric_shortcuts { + let i = self.y as usize + self.state.offset; + let is_selected = i == self.state.selected(); + let prompt: &str = if is_selected { self.indicator } else { " " }; + self.draw(prompt, Style::default()); + return; + } + + // these encode the slices of `" > "`, `" {n} "`, or `" "` in a compact form. + // Yes, this is a hack, but it makes me feel happy + + let i = self.y as usize + self.state.offset; + let i = i.checked_sub(self.state.selected); + let i = i.unwrap_or(10).min(10) * 2; + let prompt: &str = if i == 0 { + self.indicator + } else { + &SLICES[i..i + 3] + }; + self.draw(prompt, Style::default()); + } + + fn duration(&mut self, h: &History, width: u16) { + let style = self.theme.as_style(if h.success() { + Meaning::AlertInfo + } else { + Meaning::AlertError + }); + let duration = Duration::from_nanos(u64::try_from(h.duration).unwrap_or(0)); + let formatted = format_duration(duration); + let w = width as usize; + // Right-align duration within its column width, plus trailing space + let display = format!("{formatted:>w$}"); + self.draw(&display, Style::from_crossterm(style)); + } + + fn time(&mut self, h: &History, width: u16) { + let style = self.theme.as_style(Meaning::Guidance); + + // Account for the chance that h.timestamp is "in the future" + // This would mean that "since" is negative, and the unwrap here + // would fail. + // If the timestamp would otherwise be in the future, display + // the time since as 0. + let since = (self.now)() - h.timestamp; + let time = format_duration(since.try_into().unwrap_or_default()); + + // Format as "Xs ago" right-aligned within column width + let w = width as usize; + let time_str = format!("{time} ago"); + + let display = format!("{time_str:>w$}"); + self.draw(&display, Style::from_crossterm(style)); + } + + fn command(&mut self, h: &History) { + let mut style = self.theme.as_style(Meaning::Base); + let mut row_highlighted = false; + if !self.alternate_highlight && (self.y as usize + self.state.offset == self.state.selected) + { + row_highlighted = true; + // if not applying alternative highlighting to the whole row, color the command + style = self.theme.as_style(Meaning::AlertError); + style.attributes.set(style::Attribute::Bold); + } + + let highlight_indices = self.history_highlighter.get_highlight_indices( + h.command + .escape_control() + .split_ascii_whitespace() + .join(" ") + .as_str(), + ); + + let mut pos = 0; + for section in h.command.escape_control().split_ascii_whitespace() { + if pos != 0 { + self.draw(" ", Style::from_crossterm(style)); + } + for ch in section.chars() { + if self.x > self.list_area.width { + // Avoid attempting to draw a command section beyond the width + // of the list + return; + } + let mut style = style; + if highlight_indices.contains(&pos) { + if row_highlighted { + // if the row is highlighted bold is not enough as the whole row is bold + // change the color too + style = self.theme.as_style(Meaning::AlertWarn); + } + style.attributes.set(style::Attribute::Bold); + } + let s = ch.to_string(); + self.draw(&s, Style::from_crossterm(style)); + pos += s.len(); + } + pos += 1; + } + } + + /// Render the absolute datetime column (e.g., "2025-01-22 14:35") + fn datetime(&mut self, h: &History, width: u16) { + let style = self.theme.as_style(Meaning::Annotation); + // Format: YYYY-MM-DD HH:MM + let formatted = h + .timestamp + .format( + &time::format_description::parse("[year]-[month]-[day] [hour]:[minute]") + .expect("valid format"), + ) + .unwrap_or_else(|_| "????-??-?? ??:??".to_string()); + let w = width as usize; + let display = format!("{formatted:w$}"); + self.draw(&display, Style::from_crossterm(style)); + } + + /// Render the directory column (working directory, truncated) + fn directory(&mut self, h: &History, width: u16) { + let style = self.theme.as_style(Meaning::Annotation); + let w = width as usize; + let cwd = &h.cwd; + let char_count = cwd.chars().count(); + // Truncate from the left with "..." if too long, plus trailing space + // Use character count for comparison and skip for UTF-8 safety + let display = if char_count > w && w >= 4 { + let truncated: String = cwd.chars().skip(char_count - (w - 3)).collect(); + format!("...{truncated}") + } else { + format!("{cwd:w$}") + }; + self.draw(&display, Style::from_crossterm(style)); + } + + /// Render the host column (just the hostname) + fn host(&mut self, h: &History, width: u16) { + let style = self.theme.as_style(Meaning::Annotation); + let w = width as usize; + // Database stores hostname as "hostname:username" + let host = h.hostname.split(':').next().unwrap_or(&h.hostname); + let char_count = host.chars().count(); + // Use character count for comparison and take for UTF-8 safety + let display = if char_count > w && w >= 4 { + let truncated: String = host.chars().take(w.saturating_sub(4)).collect(); + format!("{truncated}...") + } else { + format!("{host:w$}") + }; + self.draw(&display, Style::from_crossterm(style)); + } + + /// Render the user column + fn user(&mut self, h: &History, width: u16) { + let style = self.theme.as_style(Meaning::Annotation); + let w = width as usize; + // Database stores hostname as "hostname:username" + let user = h.hostname.split(':').nth(1).unwrap_or(""); + let char_count = user.chars().count(); + // Use character count for comparison and take for UTF-8 safety + let display = if char_count > w && w >= 4 { + let truncated: String = user.chars().take(w.saturating_sub(4)).collect(); + format!("{truncated}...") + } else { + format!("{user:w$}") + }; + self.draw(&display, Style::from_crossterm(style)); + } + + /// Render the exit code column + fn exit_code(&mut self, h: &History, width: u16) { + let style = if h.success() { + self.theme.as_style(Meaning::AlertInfo) + } else { + self.theme.as_style(Meaning::AlertError) + }; + let w = width as usize; + let display = format!("{:>w$}", h.exit); + self.draw(&display, Style::from_crossterm(style)); + } + + fn draw(&mut self, s: &str, mut style: Style) { + let cx = self.list_area.left() + self.x; + + let cy = if self.inverted { + self.list_area.top() + self.y + } else { + self.list_area.bottom() - self.y - 1 + }; + + if self.alternate_highlight && (self.y as usize + self.state.offset == self.state.selected) + { + style = style.add_modifier(Modifier::REVERSED); + } + + let w = (self.list_area.width - self.x) as usize; + self.x += self.buf.set_stringn(cx, cy, s, w, style).0 - cx; + } +} diff --git a/crates/atuin/src/command/client/search/inspector.rs b/crates/atuin/src/command/client/search/inspector.rs new file mode 100644 index 00000000..151e1354 --- /dev/null +++ b/crates/atuin/src/command/client/search/inspector.rs @@ -0,0 +1,421 @@ +use std::time::Duration; +use time::macros::format_description; + +use atuin_client::{ + history::{History, HistoryStats}, + settings::{Settings, Timezone}, +}; +use ratatui::{ + Frame, + backend::FromCrossterm, + layout::Rect, + prelude::{Constraint, Direction, Layout}, + style::Style, + text::{Span, Text}, + widgets::{Bar, BarChart, BarGroup, Block, Borders, Padding, Paragraph, Row, Table}, +}; + +use super::duration::format_duration; + +use super::super::theme::{Meaning, Theme}; +use super::interactive::{Compactness, to_compactness}; + +#[allow(clippy::cast_sign_loss)] +fn u64_or_zero(num: i64) -> u64 { + if num < 0 { 0 } else { num as u64 } +} + +pub fn draw_commands( + f: &mut Frame<'_>, + parent: Rect, + history: &History, + stats: &HistoryStats, + compact: bool, + theme: &Theme, +) { + let commands = Layout::default() + .direction(if compact { + Direction::Vertical + } else { + Direction::Horizontal + }) + .constraints(if compact { + [ + Constraint::Length(1), + Constraint::Length(1), + Constraint::Min(0), + ] + } else { + [ + Constraint::Ratio(1, 4), + Constraint::Ratio(1, 2), + Constraint::Ratio(1, 4), + ] + }) + .split(parent); + + let command = Paragraph::new(Text::from(Span::styled( + history.command.clone(), + Style::from_crossterm(theme.as_style(Meaning::Important)), + ))) + .block(if compact { + Block::new() + .borders(Borders::NONE) + .style(Style::from_crossterm(theme.as_style(Meaning::Base))) + } else { + Block::new() + .borders(Borders::ALL) + .style(Style::from_crossterm(theme.as_style(Meaning::Base))) + .title("Command") + .padding(Padding::horizontal(1)) + }); + + let previous = Paragraph::new( + stats + .previous + .clone() + .map_or_else(|| "[No previous command]".to_string(), |prev| prev.command), + ) + .block(if compact { + Block::new() + .borders(Borders::NONE) + .style(Style::from_crossterm(theme.as_style(Meaning::Annotation))) + } else { + Block::new() + .borders(Borders::ALL) + .style(Style::from_crossterm(theme.as_style(Meaning::Annotation))) + .title("Previous command") + .padding(Padding::horizontal(1)) + }); + + // Add [] around blank text, as when this is shown in a list + // compacted, it makes it more obviously control text. + let next = Paragraph::new( + stats + .next + .clone() + .map_or_else(|| "[No next command]".to_string(), |next| next.command), + ) + .block(if compact { + Block::new() + .borders(Borders::NONE) + .style(Style::from_crossterm(theme.as_style(Meaning::Annotation))) + } else { + Block::new() + .borders(Borders::ALL) + .title("Next command") + .padding(Padding::horizontal(1)) + .style(Style::from_crossterm(theme.as_style(Meaning::Annotation))) + }); + + f.render_widget(previous, commands[0]); + f.render_widget(command, commands[1]); + f.render_widget(next, commands[2]); +} + +pub fn draw_stats_table( + f: &mut Frame<'_>, + parent: Rect, + history: &History, + tz: Timezone, + stats: &HistoryStats, + theme: &Theme, +) { + let duration = Duration::from_nanos(u64_or_zero(history.duration)); + let avg_duration = Duration::from_nanos(stats.average_duration); + let (host, user) = history.hostname.split_once(':').unwrap_or(("", "")); + + let rows = [ + Row::new(vec!["Host".to_string(), host.to_string()]), + Row::new(vec!["User".to_string(), user.to_string()]), + Row::new(vec![ + "Time".to_string(), + history.timestamp.to_offset(tz.0).to_string(), + ]), + Row::new(vec!["Duration".to_string(), format_duration(duration)]), + Row::new(vec![ + "Avg duration".to_string(), + format_duration(avg_duration), + ]), + Row::new(vec!["Exit".to_string(), history.exit.to_string()]), + Row::new(vec!["Directory".to_string(), history.cwd.clone()]), + Row::new(vec!["Session".to_string(), history.session.clone()]), + Row::new(vec!["Total runs".to_string(), stats.total.to_string()]), + ]; + + let widths = [Constraint::Ratio(1, 5), Constraint::Ratio(4, 5)]; + + let table = Table::new(rows, widths).column_spacing(1).block( + Block::default() + .title("Command stats") + .borders(Borders::ALL) + .style(Style::from_crossterm(theme.as_style(Meaning::Base))) + .padding(Padding::vertical(1)), + ); + + f.render_widget(table, parent); +} + +fn num_to_day(num: &str) -> String { + match num { + "0" => "Sunday".to_string(), + "1" => "Monday".to_string(), + "2" => "Tuesday".to_string(), + "3" => "Wednesday".to_string(), + "4" => "Thursday".to_string(), + "5" => "Friday".to_string(), + "6" => "Saturday".to_string(), + _ => "Invalid day".to_string(), + } +} + +fn sort_duration_over_time(durations: &[(String, i64)]) -> Vec<(String, i64)> { + let format = format_description!("[day]-[month]-[year]"); + let output = format_description!("[month]/[year repr:last_two]"); + + let mut durations: Vec<(time::Date, i64)> = durations + .iter() + .map(|d| { + ( + time::Date::parse(d.0.as_str(), &format).expect("invalid date string from sqlite"), + d.1, + ) + }) + .collect(); + + durations.sort_by(|a, b| a.0.cmp(&b.0)); + + durations + .iter() + .map(|(date, duration)| { + ( + date.format(output).expect("failed to format sqlite date"), + *duration, + ) + }) + .collect() +} + +fn draw_stats_charts(f: &mut Frame<'_>, parent: Rect, stats: &HistoryStats, theme: &Theme) { + let exits: Vec<Bar> = stats + .exits + .iter() + .map(|(exit, count)| { + Bar::default() + .label(exit.to_string()) + .value(u64_or_zero(*count)) + }) + .collect(); + + let exits = BarChart::default() + .block( + Block::default() + .title("Exit code distribution") + .style(Style::from_crossterm(theme.as_style(Meaning::Base))) + .borders(Borders::ALL), + ) + .bar_width(3) + .bar_gap(1) + .bar_style(Style::default()) + .value_style(Style::default()) + .label_style(Style::default()) + .data(BarGroup::default().bars(&exits)); + + let day_of_week: Vec<Bar> = stats + .day_of_week + .iter() + .map(|(day, count)| { + Bar::default() + .label(num_to_day(day.as_str())) + .value(u64_or_zero(*count)) + }) + .collect(); + + let day_of_week = BarChart::default() + .block( + Block::default() + .title("Runs per day") + .style(Style::from_crossterm(theme.as_style(Meaning::Base))) + .borders(Borders::ALL), + ) + .bar_width(3) + .bar_gap(1) + .bar_style(Style::default()) + .value_style(Style::default()) + .label_style(Style::default()) + .data(BarGroup::default().bars(&day_of_week)); + + let duration_over_time = sort_duration_over_time(&stats.duration_over_time); + let duration_over_time: Vec<Bar> = duration_over_time + .iter() + .map(|(date, duration)| { + let d = Duration::from_nanos(u64_or_zero(*duration)); + Bar::default() + .label(date.clone()) + .value(u64_or_zero(*duration)) + .text_value(format_duration(d)) + }) + .collect(); + + let duration_over_time = BarChart::default() + .block( + Block::default() + .title("Duration over time") + .style(Style::from_crossterm(theme.as_style(Meaning::Base))) + .borders(Borders::ALL), + ) + .bar_width(5) + .bar_gap(1) + .bar_style(Style::default()) + .value_style(Style::default()) + .label_style(Style::default()) + .data(BarGroup::default().bars(&duration_over_time)); + + let layout = Layout::default() + .direction(Direction::Vertical) + .constraints([ + Constraint::Ratio(1, 3), + Constraint::Ratio(1, 3), + Constraint::Ratio(1, 3), + ]) + .split(parent); + + f.render_widget(exits, layout[0]); + f.render_widget(day_of_week, layout[1]); + f.render_widget(duration_over_time, layout[2]); +} + +pub fn draw( + f: &mut Frame<'_>, + chunk: Rect, + history: &History, + stats: &HistoryStats, + settings: &Settings, + theme: &Theme, + tz: Timezone, +) { + let compactness = to_compactness(f, settings); + + match compactness { + Compactness::Ultracompact => draw_ultracompact(f, chunk, history, stats, theme), + _ => draw_full(f, chunk, history, stats, theme, tz), + } +} + +pub fn draw_ultracompact( + f: &mut Frame<'_>, + chunk: Rect, + history: &History, + stats: &HistoryStats, + theme: &Theme, +) { + draw_commands(f, chunk, history, stats, true, theme); +} + +pub fn draw_full( + f: &mut Frame<'_>, + chunk: Rect, + history: &History, + stats: &HistoryStats, + theme: &Theme, + tz: Timezone, +) { + let vert_layout = Layout::default() + .direction(Direction::Vertical) + .constraints([Constraint::Ratio(1, 5), Constraint::Ratio(4, 5)]) + .split(chunk); + + let stats_layout = Layout::default() + .direction(Direction::Horizontal) + .constraints([Constraint::Ratio(1, 3), Constraint::Ratio(2, 3)]) + .split(vert_layout[1]); + + draw_commands(f, vert_layout[0], history, stats, false, theme); + draw_stats_table(f, stats_layout[0], history, tz, stats, theme); + draw_stats_charts(f, stats_layout[1], stats, theme); +} + +#[cfg(test)] +mod tests { + use super::draw_ultracompact; + use atuin_client::{ + history::{History, HistoryId, HistoryStats}, + theme::ThemeManager, + }; + use ratatui::{backend::TestBackend, prelude::*}; + use time::OffsetDateTime; + + fn mock_history_stats() -> (History, HistoryStats) { + let history = History { + id: HistoryId::from("test1".to_string()), + timestamp: OffsetDateTime::now_utc(), + duration: 3, + exit: 0, + command: "/bin/cmd".to_string(), + cwd: "/toot".to_string(), + session: "sesh1".to_string(), + hostname: "hostn".to_string(), + author: "hostn".to_string(), + intent: None, + deleted_at: None, + }; + let next = History { + id: HistoryId::from("test2".to_string()), + timestamp: OffsetDateTime::now_utc(), + duration: 2, + exit: 0, + command: "/bin/cmd -os".to_string(), + cwd: "/toot".to_string(), + session: "sesh1".to_string(), + hostname: "hostn".to_string(), + author: "hostn".to_string(), + intent: None, + deleted_at: None, + }; + let prev = History { + id: HistoryId::from("test3".to_string()), + timestamp: OffsetDateTime::now_utc(), + duration: 1, + exit: 0, + command: "/bin/cmd -a".to_string(), + cwd: "/toot".to_string(), + session: "sesh1".to_string(), + hostname: "hostn".to_string(), + author: "hostn".to_string(), + intent: None, + deleted_at: None, + }; + let stats = HistoryStats { + next: Some(next.clone()), + previous: Some(prev.clone()), + total: 2, + average_duration: 3, + exits: Vec::new(), + day_of_week: Vec::new(), + duration_over_time: Vec::new(), + }; + (history, stats) + } + + #[test] + fn test_output_looks_correct_for_ultracompact() { + let backend = TestBackend::new(22, 5); + let mut terminal = Terminal::new(backend).expect("Could not create terminal"); + let chunk = Rect::new(0, 0, 22, 5); + let (history, stats) = mock_history_stats(); + let prev = stats.previous.clone().unwrap(); + let next = stats.next.clone().unwrap(); + + let mut manager = ThemeManager::new(Some(true), Some("".to_string())); + let theme = manager.load_theme("(none)", None); + let _ = terminal.draw(|f| draw_ultracompact(f, chunk, &history, &stats, &theme)); + let mut lines = [" "; 5].map(|l| Line::from(l)); + for (n, entry) in [prev, history, next].iter().enumerate() { + let mut l = lines[n].to_string(); + l.replace_range(0..entry.command.len(), &entry.command); + lines[n] = Line::from(l); + } + + terminal.backend().assert_buffer_lines(lines); + } +} diff --git a/crates/atuin/src/command/client/search/interactive.rs b/crates/atuin/src/command/client/search/interactive.rs new file mode 100644 index 00000000..4acf7be1 --- /dev/null +++ b/crates/atuin/src/command/client/search/interactive.rs @@ -0,0 +1,3035 @@ +use std::{ + io::{IsTerminal, Write, stdout}, + time::Duration, +}; + +#[cfg(unix)] +use std::io::Read as _; + +use atuin_common::{shell::Shell, utils::Escapable as _}; +use eyre::Result; +use futures_util::FutureExt; +use semver::Version; +use time::OffsetDateTime; +use unicode_width::UnicodeWidthStr; + +use super::{ + cursor::Cursor, + engines::{SearchEngine, SearchState}, + history_list::{HistoryList, ListState}, +}; +use atuin_client::{ + database::{Context, Database, current_context}, + history::{History, HistoryId, HistoryStats, store::HistoryStore}, + settings::{ + CursorStyle, ExitMode, FilterMode, KeymapMode, PreviewStrategy, SearchMode, Settings, + UiColumn, + }, +}; + +use crate::command::client::search::history_list::HistoryHighlighter; +use crate::command::client::search::keybindings::KeymapSet; +use crate::command::client::theme::{Meaning, Theme}; +use crate::{VERSION, command::client::search::engines}; + +use ratatui::{ + Frame, Terminal, TerminalOptions, Viewport, + backend::{CrosstermBackend, FromCrossterm}, + crossterm::{ + cursor::SetCursorStyle, + event::{self, DisableMouseCapture, EnableMouseCapture, Event, KeyEvent, MouseEvent}, + execute, queue, terminal, + }, + layout::{Alignment, Constraint, Direction, Layout}, + prelude::*, + style::{Modifier, Style}, + text::{Line, Span, Text}, + widgets::{Block, BorderType, Borders, Clear, Padding, Paragraph, Tabs}, +}; + +#[cfg(not(target_os = "windows"))] +use ratatui::crossterm::event::{ + KeyboardEnhancementFlags, PopKeyboardEnhancementFlags, PushKeyboardEnhancementFlags, +}; + +const TAB_TITLES: [&str; 2] = ["Search", "Inspect"]; + +pub enum InputAction { + Accept(usize), + AcceptInspecting, + Copy(usize), + Delete(usize), + ReturnOriginal, + ReturnQuery, + Continue, + Redraw, + SwitchContext(Option<usize>), +} + +#[derive(Clone)] +pub struct InspectingState { + current: Option<HistoryId>, + next: Option<HistoryId>, + previous: Option<HistoryId>, +} + +impl InspectingState { + pub fn move_to_previous(&mut self) { + let previous = self.previous.clone(); + self.reset(); + self.current = previous; + } + + pub fn move_to_next(&mut self) { + let next = self.next.clone(); + self.reset(); + self.current = next; + } + + pub fn reset(&mut self) { + self.current = None; + self.next = None; + self.previous = None; + } +} + +pub fn to_compactness(f: &Frame, settings: &Settings) -> Compactness { + if match settings.style { + atuin_client::settings::Style::Auto => f.area().height < 14, + atuin_client::settings::Style::Compact => true, + atuin_client::settings::Style::Full => false, + } { + if settings.auto_hide_height != 0 && f.area().height <= settings.auto_hide_height { + Compactness::Ultracompact + } else { + Compactness::Compact + } + } else { + Compactness::Full + } +} + +#[allow(clippy::struct_field_names)] +#[allow(clippy::struct_excessive_bools)] +pub struct State { + history_count: i64, + update_needed: Option<Version>, + results_state: ListState, + switched_search_mode: bool, + search_mode: SearchMode, + results_len: usize, + accept: bool, + keymap_mode: KeymapMode, + prefix: bool, + current_cursor: Option<CursorStyle>, + tab_index: usize, + pending_vim_key: Option<char>, + original_input_empty: bool, + + pub inspecting_state: InspectingState, + + keymaps: KeymapSet, + search: SearchState, + engine: Box<dyn SearchEngine>, + now: Box<dyn Fn() -> OffsetDateTime + Send>, +} + +#[derive(Clone, Copy)] +pub enum Compactness { + Ultracompact, + Compact, + Full, +} + +#[derive(Clone, Copy)] +struct StyleState { + compactness: Compactness, + invert: bool, + inner_width: usize, +} + +impl State { + async fn query_results( + &mut self, + db: &mut dyn Database, + smart_sort: bool, + ) -> Result<Vec<History>> { + let results = self.engine.query(&self.search, db).await?; + + self.inspecting_state = InspectingState { + current: None, + next: None, + previous: None, + }; + self.results_state.select(0); + self.results_len = results.len(); + + if smart_sort { + Ok(atuin_history::sort::sort( + self.search.input.as_str(), + results, + )) + } else { + Ok(results) + } + } + + fn handle_input<W>( + &mut self, + settings: &Settings, + input: &Event, + w: &mut W, + ) -> Result<InputAction> + where + W: Write, + { + execute!(w, EnableMouseCapture)?; + let r = match input { + Event::Key(k) => self.handle_key_input(settings, k), + Event::Mouse(m) => self.handle_mouse_input(*m), + Event::Paste(d) => self.handle_paste_input(d), + _ => InputAction::Continue, + }; + execute!(w, DisableMouseCapture)?; + Ok(r) + } + + fn handle_mouse_input(&mut self, input: MouseEvent) -> InputAction { + match input.kind { + event::MouseEventKind::ScrollDown => { + self.scroll_down(1); + } + event::MouseEventKind::ScrollUp => { + self.scroll_up(1); + } + _ => {} + } + InputAction::Continue + } + + fn handle_paste_input(&mut self, input: &str) -> InputAction { + for i in input.chars() { + self.search.input.insert(i); + } + InputAction::Continue + } + + fn cast_cursor_style(style: CursorStyle) -> SetCursorStyle { + match style { + CursorStyle::DefaultUserShape => SetCursorStyle::DefaultUserShape, + CursorStyle::BlinkingBlock => SetCursorStyle::BlinkingBlock, + CursorStyle::SteadyBlock => SetCursorStyle::SteadyBlock, + CursorStyle::BlinkingUnderScore => SetCursorStyle::BlinkingUnderScore, + CursorStyle::SteadyUnderScore => SetCursorStyle::SteadyUnderScore, + CursorStyle::BlinkingBar => SetCursorStyle::BlinkingBar, + CursorStyle::SteadyBar => SetCursorStyle::SteadyBar, + } + } + + fn set_keymap_cursor(&mut self, settings: &Settings, keymap_name: &str) { + let cursor_style = if keymap_name == "__clear__" { + None + } else { + settings.keymap_cursor.get(keymap_name).copied() + } + .or_else(|| self.current_cursor.map(|_| CursorStyle::DefaultUserShape)); + + if cursor_style != self.current_cursor + && let Some(style) = cursor_style + { + self.current_cursor = cursor_style; + let _ = execute!(stdout(), Self::cast_cursor_style(style)); + } + } + + pub fn initialize_keymap_cursor(&mut self, settings: &Settings) { + match self.keymap_mode { + KeymapMode::Emacs => self.set_keymap_cursor(settings, "emacs"), + KeymapMode::VimNormal => self.set_keymap_cursor(settings, "vim_normal"), + KeymapMode::VimInsert => self.set_keymap_cursor(settings, "vim_insert"), + KeymapMode::Auto => {} + } + } + + pub fn finalize_keymap_cursor(&mut self, settings: &Settings) { + match settings.keymap_mode_shell { + KeymapMode::Emacs => self.set_keymap_cursor(settings, "emacs"), + KeymapMode::VimNormal => self.set_keymap_cursor(settings, "vim_normal"), + KeymapMode::VimInsert => self.set_keymap_cursor(settings, "vim_insert"), + KeymapMode::Auto => self.set_keymap_cursor(settings, "__clear__"), + } + } + + fn handle_key_exit(settings: &Settings) -> InputAction { + match settings.exit_mode { + ExitMode::ReturnOriginal => InputAction::ReturnOriginal, + ExitMode::ReturnQuery => InputAction::ReturnQuery, + } + } + + /// Select the keymap for the current mode (ignoring prefix). + fn mode_keymap(&self) -> &super::keybindings::Keymap { + if self.tab_index == 1 { + &self.keymaps.inspector + } else { + match self.keymap_mode { + KeymapMode::Emacs | KeymapMode::Auto => &self.keymaps.emacs, + KeymapMode::VimNormal => &self.keymaps.vim_normal, + KeymapMode::VimInsert => &self.keymaps.vim_insert, + } + } + } + + /// Whether the current mode supports character insertion on unmatched keys. + fn is_insert_mode(&self) -> bool { + matches!( + self.keymap_mode, + KeymapMode::Emacs | KeymapMode::Auto | KeymapMode::VimInsert + ) + } + + fn handle_key_input(&mut self, settings: &Settings, input: &KeyEvent) -> InputAction { + use super::keybindings::Action; + use super::keybindings::EvalContext; + use super::keybindings::key::{KeyCodeValue, KeyInput, SingleKey}; + + // Skip release events + if input.kind == event::KeyEventKind::Release { + return InputAction::Continue; + } + + // Reset switched_search_mode at start of each key event + self.switched_search_mode = false; + + // Build evaluation context from current state + let ctx = EvalContext { + cursor_position: self.search.input.position(), + input_width: UnicodeWidthStr::width(self.search.input.as_str()), + input_byte_len: self.search.input.as_str().len(), + selected_index: self.results_state.selected(), + results_len: self.results_len, + original_input_empty: self.original_input_empty, + has_context: self.search.custom_context.is_some(), + }; + + // Convert KeyEvent to SingleKey + let Some(single) = SingleKey::from_event(input) else { + return InputAction::Continue; + }; + + // --- Phase 1: Resolve (take pending key first, then immutable borrows) --- + + // Take pending key before any immutable borrows of self + let pending = self.pending_vim_key.take(); + + // If in prefix mode, try prefix keymap first (single keys only) + let prefix_action = if self.prefix { + let ki = KeyInput::Single(single.clone()); + self.keymaps.prefix.resolve(&ki, &ctx) + } else { + None + }; + + // The if-let/else-if chain here is clearer than map_or_else with nested closures. + #[allow(clippy::option_if_let_else)] + let (action, new_pending) = if prefix_action.is_some() { + (prefix_action, None) + } else { + // Use mode keymap (handles both single and multi-key sequences) + let keymap = self.mode_keymap(); + + if let Some(pending_char) = pending { + // We have a pending key from a previous press (e.g., first 'g' of 'gg') + let pending_single = SingleKey { + code: KeyCodeValue::Char(pending_char), + ctrl: false, + alt: false, + shift: false, + super_key: false, + }; + let seq = KeyInput::Sequence(vec![pending_single, single.clone()]); + let action = keymap + .resolve(&seq, &ctx) + .or_else(|| keymap.resolve(&KeyInput::Single(single.clone()), &ctx)); + (action, None) + } else if keymap.has_sequence_starting_with(&single) + && matches!(single.code, KeyCodeValue::Char(_)) + && !single.ctrl + && !single.alt + { + // This key starts a multi-key sequence; wait for next key + let KeyCodeValue::Char(c) = single.code else { + unreachable!() + }; + (Some(Action::Noop), Some(c)) + } else { + ( + keymap.resolve(&KeyInput::Single(single.clone()), &ctx), + None, + ) + } + }; + + // --- Phase 2: Apply mutations --- + self.pending_vim_key = new_pending; + + // Reset prefix (before execute, so EnterPrefixMode can re-set it) + self.prefix = false; + + if let Some(action) = action { + self.execute_action(&action, settings) + } else { + // No action matched. In insert-capable modes, insert the character. + if self.is_insert_mode() && !single.ctrl && !single.alt { + match single.code { + KeyCodeValue::Char(c) => { + self.search.input.insert(c); + } + KeyCodeValue::Space => { + self.search.input.insert(' '); + } + _ => {} + } + } + InputAction::Continue + } + } + + fn scroll_down(&mut self, scroll_len: usize) { + let i = self.results_state.selected().saturating_sub(scroll_len); + self.inspecting_state.reset(); + self.results_state.select(i); + } + + fn scroll_up(&mut self, scroll_len: usize) { + let i = self.results_state.selected() + scroll_len; + self.results_state + .select(i.min(self.results_len.saturating_sub(1))); + self.inspecting_state.reset(); + } + + /// Execute a resolved action, performing all side effects and returning the + /// appropriate `InputAction` for the event loop. + /// + /// This is the "do it" half of the resolve+execute pipeline. The resolver + /// decides *what* to do (which `Action`), and this function carries it out. + /// + /// Invert handling: scroll actions (`SelectNext`, `ScrollPageDown`, etc.) account + /// for `settings.invert` so that keybindings are always in "visual" terms — + /// users never need to think about invert in their keybinding config. + #[allow(clippy::too_many_lines)] + pub(crate) fn execute_action( + &mut self, + action: &super::keybindings::Action, + settings: &Settings, + ) -> InputAction { + use crate::command::client::search::keybindings::Action; + + match action { + // -- Cursor movement -- + Action::CursorLeft => { + self.search.input.left(); + InputAction::Continue + } + Action::CursorRight => { + self.search.input.right(); + InputAction::Continue + } + Action::CursorWordLeft => { + self.search + .input + .prev_word(&settings.word_chars, settings.word_jump_mode); + InputAction::Continue + } + Action::CursorWordRight => { + self.search + .input + .next_word(&settings.word_chars, settings.word_jump_mode); + InputAction::Continue + } + Action::CursorWordEnd => { + self.search.input.word_end(&settings.word_chars); + InputAction::Continue + } + Action::CursorStart => { + self.search.input.start(); + InputAction::Continue + } + Action::CursorEnd => { + self.search.input.end(); + InputAction::Continue + } + + // -- Editing -- + Action::DeleteCharBefore => { + self.search.input.back(); + InputAction::Continue + } + Action::DeleteCharAfter => { + self.search.input.remove(); + InputAction::Continue + } + Action::DeleteWordBefore => { + self.search + .input + .remove_prev_word(&settings.word_chars, settings.word_jump_mode); + InputAction::Continue + } + Action::DeleteWordAfter => { + self.search + .input + .remove_next_word(&settings.word_chars, settings.word_jump_mode); + InputAction::Continue + } + Action::DeleteToWordBoundary => { + // ctrl-w: remove trailing whitespace, then delete to word boundary + while matches!(self.search.input.back(), Some(c) if c.is_whitespace()) {} + while self.search.input.left() { + if self.search.input.char().unwrap().is_whitespace() { + self.search.input.right(); + break; + } + self.search.input.remove(); + } + InputAction::Continue + } + Action::ClearLine => { + self.search.input.clear(); + InputAction::Continue + } + Action::ClearToStart => { + self.search.input.clear_to_start(); + InputAction::Continue + } + Action::ClearToEnd => { + self.search.input.clear_to_end(); + InputAction::Continue + } + + // -- List navigation (invert-aware) -- + Action::SelectNext => { + if settings.invert { + self.scroll_up(1); + } else { + self.scroll_down(1); + } + InputAction::Continue + } + Action::SelectPrevious => { + if settings.invert { + self.scroll_down(1); + } else { + self.scroll_up(1); + } + InputAction::Continue + } + // -- Page/half-page scroll (invert-aware) -- + Action::ScrollHalfPageUp => { + let scroll_len = self + .results_state + .max_entries() + .saturating_sub(settings.scroll_context_lines) + / 2; + if settings.invert { + self.scroll_down(scroll_len); + } else { + self.scroll_up(scroll_len); + } + InputAction::Continue + } + Action::ScrollHalfPageDown => { + let scroll_len = self + .results_state + .max_entries() + .saturating_sub(settings.scroll_context_lines) + / 2; + if settings.invert { + self.scroll_up(scroll_len); + } else { + self.scroll_down(scroll_len); + } + InputAction::Continue + } + Action::ScrollPageUp => { + let scroll_len = self + .results_state + .max_entries() + .saturating_sub(settings.scroll_context_lines); + if settings.invert { + self.scroll_down(scroll_len); + } else { + self.scroll_up(scroll_len); + } + InputAction::Continue + } + Action::ScrollPageDown => { + let scroll_len = self + .results_state + .max_entries() + .saturating_sub(settings.scroll_context_lines); + if settings.invert { + self.scroll_up(scroll_len); + } else { + self.scroll_down(scroll_len); + } + InputAction::Continue + } + + // -- Absolute jumps (invert-aware) -- + Action::ScrollToTop => { + // Visual top of history + if settings.invert { + self.results_state.select(0); + } else { + let last_idx = self.results_len.saturating_sub(1); + self.results_state.select(last_idx); + } + self.inspecting_state.reset(); + InputAction::Continue + } + Action::ScrollToBottom => { + // Visual bottom of history + if settings.invert { + let last_idx = self.results_len.saturating_sub(1); + self.results_state.select(last_idx); + } else { + self.results_state.select(0); + } + self.inspecting_state.reset(); + InputAction::Continue + } + Action::ScrollToScreenTop => { + // H — jump to top of visible screen + let top = self.results_state.offset(); + let visible = self.results_state.max_entries().min(self.results_len); + let bottom = top + visible.saturating_sub(1); + self.results_state + .select(bottom.min(self.results_len.saturating_sub(1))); + self.inspecting_state.reset(); + InputAction::Continue + } + Action::ScrollToScreenMiddle => { + // M — jump to middle of visible screen + let top = self.results_state.offset(); + let visible = self.results_state.max_entries().min(self.results_len); + let middle = top + visible / 2; + self.results_state + .select(middle.min(self.results_len.saturating_sub(1))); + self.inspecting_state.reset(); + InputAction::Continue + } + Action::ScrollToScreenBottom => { + // L — jump to bottom of visible screen + let top_visible = self.results_state.offset(); + self.results_state.select(top_visible); + self.inspecting_state.reset(); + InputAction::Continue + } + + // -- Commands -- + Action::Accept => { + if self.tab_index == 1 { + return InputAction::AcceptInspecting; + } + self.accept = true; + InputAction::Accept(self.results_state.selected()) + } + Action::AcceptNth(n) => { + self.accept = true; + InputAction::Accept(self.results_state.selected() + *n as usize) + } + Action::ReturnSelection => { + if self.tab_index == 1 { + return InputAction::AcceptInspecting; + } + InputAction::Accept(self.results_state.selected()) + } + Action::ReturnSelectionNth(n) => { + InputAction::Accept(self.results_state.selected() + *n as usize) + } + Action::Copy => InputAction::Copy(self.results_state.selected()), + Action::Delete => InputAction::Delete(self.results_state.selected()), + Action::ReturnOriginal => InputAction::ReturnOriginal, + Action::ReturnQuery => InputAction::ReturnQuery, + Action::Exit => Self::handle_key_exit(settings), + Action::Redraw => InputAction::Redraw, + Action::CycleFilterMode => { + self.search.rotate_filter_mode(settings, 1); + InputAction::Continue + } + Action::CycleSearchMode => { + self.switched_search_mode = true; + self.search_mode = self.search_mode.next(settings); + self.engine = engines::engine(self.search_mode, settings); + InputAction::Continue + } + Action::SwitchContext => { + InputAction::SwitchContext(Some(self.results_state.selected())) + } + Action::ClearContext => InputAction::SwitchContext(None), + Action::ToggleTab => { + self.tab_index = (self.tab_index + 1) % TAB_TITLES.len(); + InputAction::Continue + } + + // -- Mode changes -- + Action::VimEnterNormal => { + self.set_keymap_cursor(settings, "vim_normal"); + self.keymap_mode = KeymapMode::VimNormal; + InputAction::Continue + } + Action::VimEnterInsert => { + self.set_keymap_cursor(settings, "vim_insert"); + self.keymap_mode = KeymapMode::VimInsert; + InputAction::Continue + } + Action::VimEnterInsertAfter => { + self.search.input.right(); + self.set_keymap_cursor(settings, "vim_insert"); + self.keymap_mode = KeymapMode::VimInsert; + InputAction::Continue + } + Action::VimEnterInsertAtStart => { + self.search.input.start(); + self.set_keymap_cursor(settings, "vim_insert"); + self.keymap_mode = KeymapMode::VimInsert; + InputAction::Continue + } + Action::VimEnterInsertAtEnd => { + self.search.input.end(); + self.set_keymap_cursor(settings, "vim_insert"); + self.keymap_mode = KeymapMode::VimInsert; + InputAction::Continue + } + Action::VimSearchInsert => { + self.search.input.clear(); + self.set_keymap_cursor(settings, "vim_insert"); + self.keymap_mode = KeymapMode::VimInsert; + InputAction::Continue + } + Action::VimChangeToEnd => { + self.search.input.clear_to_end(); + self.set_keymap_cursor(settings, "vim_insert"); + self.keymap_mode = KeymapMode::VimInsert; + InputAction::Continue + } + Action::EnterPrefixMode => { + self.prefix = true; + InputAction::Continue + } + + // -- Inspector -- + Action::InspectPrevious => { + self.inspecting_state.move_to_previous(); + InputAction::Redraw + } + Action::InspectNext => { + self.inspecting_state.move_to_next(); + InputAction::Redraw + } + + // -- Special -- + Action::Noop => InputAction::Continue, + } + } + + #[allow(clippy::cast_possible_truncation)] + #[allow(clippy::bool_to_int_with_if)] + fn calc_preview_height( + settings: &Settings, + results: &[History], + selected: usize, + tab_index: usize, + compactness: Compactness, + border_size: u16, + preview_width: u16, + ) -> u16 { + if settings.show_preview + && settings.preview.strategy == PreviewStrategy::Auto + && tab_index == 0 + && !results.is_empty() + { + let length_current_cmd = results[selected].command.len() as u16; + // calculate the number of newlines in the command + let num_newlines = results[selected] + .command + .chars() + .filter(|&c| c == '\n') + .count() as u16; + if num_newlines > 0 { + std::cmp::min( + settings.max_preview_height, + results[selected] + .command + .split('\n') + .map(|line| { + (line.len() as u16 + preview_width - 1 - border_size) + / (preview_width - border_size) + }) + .sum(), + ) + border_size * 2 + } + // The '- 19' takes the characters before the command (duration and time) into account + else if length_current_cmd > preview_width - 19 { + std::cmp::min( + settings.max_preview_height, + (length_current_cmd + preview_width - 1 - border_size) + / (preview_width - border_size), + ) + border_size * 2 + } else { + 1 + } + } else if settings.show_preview + && settings.preview.strategy == PreviewStrategy::Static + && tab_index == 0 + { + let longest_command = results + .iter() + .max_by(|h1, h2| h1.command.len().cmp(&h2.command.len())); + longest_command.map_or(0, |v| { + std::cmp::min( + settings.max_preview_height, + v.command + .split('\n') + .map(|line| { + (line.len() as u16 + preview_width - 1 - border_size) + / (preview_width - border_size) + }) + .sum(), + ) + }) + border_size * 2 + } else if settings.show_preview && settings.preview.strategy == PreviewStrategy::Fixed { + settings.max_preview_height + border_size * 2 + } else if !matches!(compactness, Compactness::Full) || tab_index == 1 { + 0 + } else { + 1 + } + } + + #[allow(clippy::bool_to_int_with_if)] + #[allow(clippy::too_many_lines)] + #[allow(clippy::too_many_arguments)] + fn draw( + &mut self, + f: &mut Frame, + results: &[History], + stats: Option<HistoryStats>, + inspecting: Option<&History>, + settings: &Settings, + theme: &Theme, + popup_mode: bool, + ) { + let area = f.area(); + if popup_mode { + f.render_widget(Clear, area); + } + self.draw_inner(f, area, results, stats, inspecting, settings, theme); + } + + #[allow(clippy::too_many_arguments)] + #[allow(clippy::too_many_lines)] + #[allow(clippy::bool_to_int_with_if)] + fn draw_inner( + &mut self, + f: &mut Frame, + area: Rect, + results: &[History], + stats: Option<HistoryStats>, + inspecting: Option<&History>, + settings: &Settings, + theme: &Theme, + ) { + let compactness = to_compactness(f, settings); + let invert = settings.invert; + let border_size = match compactness { + Compactness::Full => 1, + _ => 0, + }; + let preview_width = area.width.saturating_sub(2); + let preview_height = Self::calc_preview_height( + settings, + results, + self.results_state.selected(), + self.tab_index, + compactness, + border_size, + preview_width, + ); + let show_help = + settings.show_help && (matches!(compactness, Compactness::Full) || area.height > 1); + // This is an OR, as it seems more likely for someone to wish to override + // tabs unexpectedly being missed, than unexpectedly present. + let show_tabs = settings.show_tabs && !matches!(compactness, Compactness::Ultracompact); + let chunks = Layout::default() + .direction(Direction::Vertical) + .margin(0) + .horizontal_margin(1) + .constraints::<&[Constraint]>( + if invert { + [ + Constraint::Length(1 + border_size), // input + Constraint::Min(1), // results list + Constraint::Length(preview_height), // preview + Constraint::Length(if show_tabs { 1 } else { 0 }), // tabs + Constraint::Length(if show_help { 1 } else { 0 }), // header (sic) + ] + } else { + match compactness { + Compactness::Ultracompact => [ + Constraint::Length(if show_help { 1 } else { 0 }), // header + Constraint::Length(0), // tabs + Constraint::Min(1), // results list + Constraint::Length(0), + Constraint::Length(0), + ], + _ => [ + Constraint::Length(if show_help { 1 } else { 0 }), // header + Constraint::Length(if show_tabs { 1 } else { 0 }), // tabs + Constraint::Min(1), // results list + Constraint::Length(1 + border_size), // input + Constraint::Length(preview_height), // preview + ], + } + } + .as_ref(), + ) + .split(area); + + let input_chunk = if invert { chunks[0] } else { chunks[3] }; + let results_list_chunk = if invert { chunks[1] } else { chunks[2] }; + let preview_chunk = if invert { chunks[2] } else { chunks[4] }; + let tabs_chunk = if invert { chunks[3] } else { chunks[1] }; + let header_chunk = if invert { chunks[4] } else { chunks[0] }; + + // TODO: this should be split so that we have one interactive search container that is + // EITHER a search box or an inspector. But I'm not doing that now, way too much atm. + // also allocate less 🙈 + let titles: Vec<_> = TAB_TITLES.iter().copied().map(Line::from).collect(); + + if show_tabs { + let tabs = Tabs::new(titles) + .block(Block::default().borders(Borders::NONE)) + .select(self.tab_index) + .style(Style::default()) + .highlight_style(Style::from_crossterm(theme.as_style(Meaning::Important))); + + f.render_widget(tabs, tabs_chunk); + } + + let style = StyleState { + compactness, + invert, + inner_width: input_chunk.width.into(), + }; + + let header_chunks = Layout::default() + .direction(Direction::Horizontal) + .constraints::<&[Constraint]>( + [ + Constraint::Ratio(1, 5), + Constraint::Ratio(3, 5), + Constraint::Ratio(1, 5), + ] + .as_ref(), + ) + .split(header_chunk); + + let title = self.build_title(theme); + f.render_widget(title, header_chunks[0]); + + let help = self.build_help(settings, theme); + f.render_widget(help, header_chunks[1]); + + let stats_tab = self.build_stats(theme); + f.render_widget(stats_tab, header_chunks[2]); + + let indicator: String = match compactness { + Compactness::Ultracompact => { + if self.switched_search_mode { + format!("S{}>", self.search_mode.as_str().chars().next().unwrap()) + } else if self.search.custom_context.is_some() { + format!( + "C{}>", + self.search.filter_mode.as_str().chars().next().unwrap() + ) + } else { + format!( + "{}> ", + self.search.filter_mode.as_str().chars().next().unwrap() + ) + } + } + _ => " > ".to_string(), + }; + + match self.tab_index { + 0 => { + let history_highlighter = HistoryHighlighter { + engine: self.engine.as_ref(), + search_input: self.search.input.as_str(), + }; + let results_list = Self::build_results_list( + style, + results, + self.keymap_mode, + &self.now, + indicator.as_str(), + theme, + history_highlighter, + settings.show_numeric_shortcuts, + &settings.ui.columns, + ); + f.render_stateful_widget(results_list, results_list_chunk, &mut self.results_state); + } + + 1 => { + if results.is_empty() { + let message = Paragraph::new("Nothing to inspect") + .block( + Block::new() + .title(Line::from(" Info ".to_string())) + .title_alignment(Alignment::Center) + .borders(Borders::ALL) + .padding(Padding::vertical(2)), + ) + .alignment(Alignment::Center); + f.render_widget(message, results_list_chunk); + } else { + let inspecting = match inspecting { + Some(inspecting) => inspecting, + None => &results[self.results_state.selected()], + }; + super::inspector::draw( + f, + results_list_chunk, + inspecting, + &stats.expect("Drawing inspector, but no stats"), + settings, + theme, + settings.timezone, + ); + } + + // HACK: I'm following up with abstracting this into the UI container, with a + // sub-widget for search + for inspector + let feedback = Paragraph::new( + "The inspector is new - please give feedback (good, or bad) at https://forum.atuin.sh", + ); + f.render_widget(feedback, input_chunk); + + return; + } + + _ => { + panic!("invalid tab index"); + } + } + + if !matches!(compactness, Compactness::Ultracompact) { + let preview_width = match compactness { + Compactness::Full => preview_width - 2, + _ => preview_width, + }; + let preview = self.build_preview( + results, + compactness, + preview_width, + preview_chunk.width.into(), + theme, + ); + #[allow(clippy::cast_possible_truncation)] + let prefix_width = settings + .ui + .columns + .iter() + .take_while(|col| !col.expand) + .map(|col| col.width + 1) + .sum::<u16>() + + " > ".len() as u16; + #[allow(clippy::cast_possible_truncation)] + let min_prefix_width = "[ SRCH: FULLTXT ] ".len() as u16; + self.draw_preview( + f, + style, + input_chunk, + compactness, + preview_chunk, + preview, + std::cmp::max(prefix_width, min_prefix_width), + ); + } + } + + #[allow(clippy::cast_possible_truncation, clippy::too_many_arguments)] + fn draw_preview( + &self, + f: &mut Frame, + style: StyleState, + input_chunk: Rect, + compactness: Compactness, + preview_chunk: Rect, + preview: Paragraph, + prefix_width: u16, + ) { + let input = self.build_input(style, prefix_width); + f.render_widget(input, input_chunk); + + f.render_widget(preview, preview_chunk); + + let extra_width = UnicodeWidthStr::width(self.search.input.substring()); + + let cursor_offset = match compactness { + Compactness::Full => 1, + _ => 0, + }; + f.set_cursor_position(( + // Put cursor past the end of the input text + input_chunk.x + extra_width as u16 + prefix_width + cursor_offset, + input_chunk.y + cursor_offset, + )); + } + + fn build_title(&self, theme: &Theme) -> Paragraph<'_> { + let title = if self.update_needed.is_some() { + let error_style: Style = Style::from_crossterm(theme.get_error()); + Paragraph::new(Text::from(Span::styled( + format!("Atuin v{VERSION} - UPDATE"), + error_style.add_modifier(Modifier::BOLD), + ))) + } else { + let style: Style = Style::from_crossterm(theme.as_style(Meaning::Base)); + Paragraph::new(Text::from(Span::styled( + format!("Atuin v{VERSION}"), + style.add_modifier(Modifier::BOLD), + ))) + }; + title.alignment(Alignment::Left) + } + + #[allow(clippy::unused_self)] + fn build_help(&self, settings: &Settings, theme: &Theme) -> Paragraph<'_> { + match self.tab_index { + // search + 0 => Paragraph::new(Text::from(Line::from(vec![ + Span::styled("<esc>", Style::default().add_modifier(Modifier::BOLD)), + Span::raw(": exit"), + Span::raw(", "), + Span::styled("<tab>", Style::default().add_modifier(Modifier::BOLD)), + Span::raw(": edit"), + Span::raw(", "), + Span::styled("<enter>", Style::default().add_modifier(Modifier::BOLD)), + Span::raw(if settings.enter_accept { + ": run" + } else { + ": edit" + }), + Span::raw(", "), + Span::styled("<ctrl-o>", Style::default().add_modifier(Modifier::BOLD)), + Span::raw(": inspect"), + ]))), + + 1 => Paragraph::new(Text::from(Line::from(vec![ + Span::styled("<esc>", Style::default().add_modifier(Modifier::BOLD)), + Span::raw(": exit"), + Span::raw(", "), + Span::styled("<ctrl-o>", Style::default().add_modifier(Modifier::BOLD)), + Span::raw(": search"), + Span::raw(", "), + Span::styled("<ctrl-d>", Style::default().add_modifier(Modifier::BOLD)), + Span::raw(": delete"), + ]))), + + _ => unreachable!("invalid tab index"), + } + .style(Style::from_crossterm(theme.as_style(Meaning::Annotation))) + .alignment(Alignment::Center) + } + + fn build_stats(&self, theme: &Theme) -> Paragraph<'_> { + Paragraph::new(Text::from(Span::raw(format!( + "history count: {}", + self.history_count, + )))) + .style(Style::from_crossterm(theme.as_style(Meaning::Annotation))) + .alignment(Alignment::Right) + } + + #[allow(clippy::too_many_arguments)] + fn build_results_list<'a>( + style: StyleState, + results: &'a [History], + keymap_mode: KeymapMode, + now: &'a dyn Fn() -> OffsetDateTime, + indicator: &'a str, + theme: &'a Theme, + history_highlighter: HistoryHighlighter<'a>, + show_numeric_shortcuts: bool, + columns: &'a [UiColumn], + ) -> HistoryList<'a> { + let results_list = HistoryList::new( + results, + style.invert, + keymap_mode == KeymapMode::VimNormal, + now, + indicator, + theme, + history_highlighter, + show_numeric_shortcuts, + columns, + ); + + match style.compactness { + Compactness::Full => { + if style.invert { + results_list.block( + Block::default() + .borders(Borders::LEFT | Borders::RIGHT) + .border_type(BorderType::Rounded) + .title(format!("{:─>width$}", "", width = style.inner_width - 2)), + ) + } else { + results_list.block( + Block::default() + .borders(Borders::TOP | Borders::LEFT | Borders::RIGHT) + .border_type(BorderType::Rounded), + ) + } + } + _ => results_list, + } + } + + fn build_input(&self, style: StyleState, prefix_width: u16) -> Paragraph<'_> { + let (pref, mode) = if self.switched_search_mode { + (" SRCH:", self.search_mode.as_str()) + } else if self.search.custom_context.is_some() { + (" CTX:", self.search.filter_mode.as_str()) + } else { + ("", self.search.filter_mode.as_str()) + }; + // 3: surrounding "[" "] " + let mode_width = usize::from(prefix_width) - pref.len() - 3; + // sanity check to ensure we don't exceed the layout limits + debug_assert!(mode_width >= mode.len(), "mode name '{mode}' is too long!"); + let input = format!("[{pref}{mode:^mode_width$}] {}", self.search.input.as_str(),); + let input = Paragraph::new(input); + match style.compactness { + Compactness::Full => { + if style.invert { + input.block( + Block::default() + .borders(Borders::LEFT | Borders::RIGHT | Borders::TOP) + .border_type(BorderType::Rounded), + ) + } else { + input.block( + Block::default() + .borders(Borders::LEFT | Borders::RIGHT) + .border_type(BorderType::Rounded) + .title(format!("{:─>width$}", "", width = style.inner_width - 2)), + ) + } + } + _ => input, + } + } + + fn build_preview( + &self, + results: &[History], + compactness: Compactness, + preview_width: u16, + chunk_width: usize, + theme: &Theme, + ) -> Paragraph<'_> { + let selected = self.results_state.selected(); + let command = if results.is_empty() { + String::new() + } else { + use itertools::Itertools as _; + let s = &results[selected].command; + s.split('\n') + .flat_map(|line| { + line.char_indices() + .step_by(preview_width.into()) + .map(|(i, _)| i) + .chain(Some(line.len())) + .tuple_windows() + .map(|(a, b)| (&line[a..b]).escape_control().to_string()) + }) + .join("\n") + }; + + match compactness { + Compactness::Full => Paragraph::new(command).block( + Block::default() + .borders(Borders::BOTTOM | Borders::LEFT | Borders::RIGHT) + .border_type(BorderType::Rounded) + .title(format!("{:─>width$}", "", width = chunk_width - 2)), + ), + _ => Paragraph::new(command) + .style(Style::from_crossterm(theme.as_style(Meaning::Annotation))), + } + } +} + +/// The writer used for terminal output - either stdout or /dev/tty +enum TerminalWriter { + Stdout(std::io::Stdout), + #[cfg(unix)] + Tty(std::fs::File), +} + +impl Write for TerminalWriter { + fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> { + match self { + TerminalWriter::Stdout(stdout) => stdout.write(buf), + #[cfg(unix)] + TerminalWriter::Tty(file) => file.write(buf), + } + } + + fn flush(&mut self) -> std::io::Result<()> { + match self { + TerminalWriter::Stdout(stdout) => stdout.flush(), + #[cfg(unix)] + TerminalWriter::Tty(file) => file.flush(), + } + } +} + +/// Screen state captured from atuin-hex's screen server. +#[cfg(unix)] +struct SavedScreen { + #[allow(dead_code)] + rows: u16, + #[allow(dead_code)] + cols: u16, + cursor_row: u16, + cursor_col: u16, + /// Pre-formatted ANSI bytes for each screen row, ready to write to stdout. + rows_data: Vec<Vec<u8>>, +} + +/// Connect to atuin-hex's Unix socket and fetch the current screen state. +/// +/// The wire format is: +/// ```text +/// [rows: u16 BE][cols: u16 BE][cursor_row: u16 BE][cursor_col: u16 BE] +/// [row_0_len: u32 BE][row_0_bytes...] +/// [row_1_len: u32 BE][row_1_bytes...] +/// ... +/// ``` +#[cfg(unix)] +fn fetch_screen_state(socket_path: &str) -> Option<SavedScreen> { + use std::os::unix::net::UnixStream; + + let mut stream = UnixStream::connect(socket_path).ok()?; + stream.set_read_timeout(Some(Duration::from_secs(2))).ok()?; + + let mut data = Vec::new(); + stream.read_to_end(&mut data).ok()?; + + if data.len() < 8 { + return None; + } + + let rows = u16::from_be_bytes([data[0], data[1]]); + let cols = u16::from_be_bytes([data[2], data[3]]); + let cursor_row = u16::from_be_bytes([data[4], data[5]]); + let cursor_col = u16::from_be_bytes([data[6], data[7]]); + + // Parse length-prefixed rows + let mut rows_data = Vec::with_capacity(rows as usize); + let mut offset = 8; + while offset + 4 <= data.len() { + let row_len = u32::from_be_bytes([ + data[offset], + data[offset + 1], + data[offset + 2], + data[offset + 3], + ]) as usize; + offset += 4; + if offset + row_len > data.len() { + break; + } + rows_data.push(data[offset..offset + row_len].to_vec()); + offset += row_len; + } + + Some(SavedScreen { + rows, + cols, + cursor_row, + cursor_col, + rows_data, + }) +} + +/// Restore the screen area that was covered by the popup. +/// +/// Writes the pre-formatted per-row ANSI bytes received from atuin-hex +/// directly to stdout, which correctly handles wide characters, colors, and +/// all text attributes without needing a client-side vt100 parser. +#[cfg(unix)] +fn restore_popup_area(saved: &SavedScreen, popup_rect: Rect, scroll_offset: u16) { + use ratatui::crossterm::cursor::MoveTo; + + let mut stdout = stdout(); + + for dy in 0..popup_rect.height { + let target_row = popup_rect.y + dy; + let source_row = (target_row + scroll_offset) as usize; + + // Clear only the popup region. The server-side rows_formatted() skips + // default cells (spaces with default attributes) using cursor jumps, so + // any popup content at those positions would remain if not cleared + // beforehand. We write `popup_rect.width` spaces instead of + // ClearType::CurrentLine so that only the popup area is cleared, not + // the entire terminal line. + let _ = execute!( + stdout, + MoveTo(popup_rect.x, target_row), + ratatui::crossterm::style::SetAttribute(ratatui::crossterm::style::Attribute::Reset), + ); + let _ = write!(stdout, "{:width$}", "", width = popup_rect.width as usize); + let _ = execute!(stdout, MoveTo(popup_rect.x, target_row)); + + if let Some(row_bytes) = saved.rows_data.get(source_row) { + let _ = stdout.write_all(row_bytes); + } + } + + let _ = execute!( + stdout, + MoveTo( + saved.cursor_col, + saved.cursor_row.saturating_sub(scroll_offset) + ) + ); + let _ = stdout.flush(); +} + +struct Stdout { + writer: TerminalWriter, + inline_mode: bool, +} + +impl Stdout { + pub fn new(inline_mode: bool, stdout_is_terminal: bool) -> std::io::Result<Self> { + terminal::enable_raw_mode()?; + + // If stdout is not a terminal (e.g., captured by command substitution), + // fall back to /dev/tty so the TUI can still render. + // This allows usage like: VAR=$(atuin search -i) + let mut writer = if stdout_is_terminal { + TerminalWriter::Stdout(stdout()) + } else { + #[cfg(unix)] + { + TerminalWriter::Tty( + std::fs::File::options() + .read(true) + .write(true) + .open("/dev/tty")?, + ) + } + #[cfg(not(unix))] + { + return Err(std::io::Error::new( + std::io::ErrorKind::Unsupported, + "Interactive mode requires a terminal", + )); + } + }; + + if !inline_mode { + execute!(writer, terminal::EnterAlternateScreen)?; + } + + execute!( + writer, + event::EnableMouseCapture, + event::EnableBracketedPaste, + )?; + + #[cfg(not(target_os = "windows"))] + execute!( + writer, + PushKeyboardEnhancementFlags( + KeyboardEnhancementFlags::DISAMBIGUATE_ESCAPE_CODES + | KeyboardEnhancementFlags::REPORT_ALL_KEYS_AS_ESCAPE_CODES + | KeyboardEnhancementFlags::REPORT_ALTERNATE_KEYS + ), + )?; + + Ok(Self { + writer, + inline_mode, + }) + } +} + +impl Drop for Stdout { + fn drop(&mut self) { + #[cfg(not(target_os = "windows"))] + execute!(self.writer, PopKeyboardEnhancementFlags).unwrap(); + + if !self.inline_mode { + execute!(self.writer, terminal::LeaveAlternateScreen).unwrap(); + } + execute!( + self.writer, + event::DisableMouseCapture, + event::DisableBracketedPaste, + ) + .unwrap(); + + terminal::disable_raw_mode().unwrap(); + } +} + +impl Write for Stdout { + fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> { + self.writer.write(buf) + } + + fn flush(&mut self) -> std::io::Result<()> { + self.writer.flush() + } +} + +// this is a big blob of horrible! clean it up! +/// Compute the popup position and any scroll offset needed to make room. +/// +/// Given the cursor row, terminal dimensions, and desired popup height, +/// returns `(popup_rect, scroll_offset)` where `scroll_offset` is the number +/// of lines the caller should scroll the terminal up before rendering. +/// +/// This function performs no I/O — it is a pure computation. +fn compute_popup_placement( + cursor_row: u16, + term_rows: u16, + term_cols: u16, + inline_height: u16, +) -> (Rect, u16) { + let popup_w = term_cols; + let popup_h = inline_height.min(term_rows); + let space_below = term_rows.saturating_sub(cursor_row); + + let (popup_y, scroll) = if popup_h <= space_below { + // Fits below cursor + (cursor_row, 0u16) + } else if cursor_row >= term_rows / 2 { + // Bottom half — render above cursor (overlay on existing text) + (cursor_row.saturating_sub(popup_h), 0u16) + } else { + // Top half, not enough space — scroll terminal to make room + let scroll = popup_h.saturating_sub(space_below); + let popup_y = cursor_row.saturating_sub(scroll); + (popup_y, scroll) + }; + + (Rect::new(0, popup_y, popup_w, popup_h), scroll) +} + +// for now, it works. But it'd be great if it were more easily readable, and +// modular. I'd like to add some more stats and stuff at some point +#[allow( + clippy::cast_possible_truncation, + clippy::too_many_lines, + clippy::cognitive_complexity +)] +pub async fn history( + query: &[String], + settings: &Settings, + mut db: impl Database, + history_store: &HistoryStore, + theme: &Theme, +) -> Result<String> { + let inline_height = if settings.shell_up_key_binding { + settings + .inline_height_shell_up_key_binding + .unwrap_or(settings.inline_height) + } else { + settings.inline_height + }; + + // Check if stdout is a terminal - if not (e.g., command substitution like VAR=$(atuin search -i)), + // we need to use /dev/tty for the TUI and force fullscreen mode (inline mode requires + // cursor position queries that don't work when stdout is captured) + let stdout_is_terminal = stdout().is_terminal(); + + // Use fullscreen mode if the inline height doesn't fit in the terminal, + // this will preserve the scroll position upon exit. + // Also force fullscreen when stdout isn't a terminal (inline mode won't work). + let inline_height = if !stdout_is_terminal { + 0 + } else if let Ok(size) = terminal::size() + && inline_height >= size.1 + { + 0 + } else { + inline_height + }; + + // Popup mode: if running under atuin-hex and inline mode is requested, + // fetch the screen state and render as a centered overlay. + #[cfg(unix)] + let (saved_screen, popup_rect, popup_scroll_offset) = { + let socket_path = std::env::var("ATUIN_HEX_SOCKET").ok(); + if let Some(ref path) = socket_path + && inline_height > 0 + { + let saved = fetch_screen_state(path); + if let Some(ref s) = saved { + let (term_cols, term_rows) = terminal::size().unwrap_or((s.cols, s.rows)); + let (popup_rect, scroll) = + compute_popup_placement(s.cursor_row, term_rows, term_cols, inline_height); + + // Scroll terminal content up to make room if needed + if scroll > 0 { + use ratatui::crossterm::cursor::MoveTo; + let mut stdout = stdout(); + let _ = execute!(stdout, MoveTo(0, term_rows - 1)); + for _ in 0..scroll { + let _ = writeln!(stdout); + } + let _ = stdout.flush(); + } + + (saved, popup_rect, scroll) + } else { + (None, Rect::default(), 0u16) + } + } else { + (None, Rect::default(), 0u16) + } + }; + + #[cfg(not(unix))] + let (saved_screen, popup_rect, popup_scroll_offset): (Option<()>, Rect, u16) = + (None, Rect::default(), 0); + + let popup_mode = saved_screen.is_some(); + + let stdout = Stdout::new(inline_height > 0, stdout_is_terminal)?; + + // In popup mode, clear the popup region on the physical terminal before + // ratatui takes over. Ratatui's diff-based rendering compares against an + // initially-empty buffer, so cells that remain "empty" (spaces with default + // style) won't be written — leaving underlying terminal text visible. + // By pre-clearing with spaces, those cells are already correct on screen. + if popup_mode { + use ratatui::crossterm::cursor::MoveTo; + let mut raw_stdout = std::io::stdout(); + // Queue all commands without flushing so the terminal receives them + // as a single write — no intermediate cursor positions are visible. + let _ = queue!( + raw_stdout, + ratatui::crossterm::style::SetAttribute(ratatui::crossterm::style::Attribute::Reset) + ); + for row in popup_rect.y..popup_rect.y.saturating_add(popup_rect.height) { + let _ = queue!(raw_stdout, MoveTo(popup_rect.x, row)); + let _ = write!( + raw_stdout, + "{:width$}", + "", + width = popup_rect.width as usize + ); + } + let _ = raw_stdout.flush(); + } + + let backend = CrosstermBackend::new(stdout); + let mut terminal = Terminal::with_options( + backend, + TerminalOptions { + viewport: if popup_mode { + Viewport::Fixed(popup_rect) + } else if inline_height > 0 { + Viewport::Inline(inline_height) + } else { + Viewport::Fullscreen + }, + }, + )?; + + let original_query = query.join(" "); + + // Check if this is a command chaining scenario + let is_command_chaining = if settings.command_chaining { + let trimmed = original_query.trim_end(); + trimmed.ends_with("&&") || trimmed.ends_with('|') + } else { + false + }; + + // For command chaining, start with empty input to allow searching for new commands + let search_input = if is_command_chaining { + String::new() + } else { + original_query.clone() + }; + + let mut input = Cursor::from(search_input); + // Put the cursor at the end of the query by default + input.end(); + + let settings2 = settings.clone(); + let update_needed = tokio::spawn(async move { settings2.needs_update().await }).fuse(); + tokio::pin!(update_needed); + + let initial_context = current_context().await?; + + let history_count = db.history_count(false).await?; + let search_mode = if settings.shell_up_key_binding { + settings + .search_mode_shell_up_key_binding + .unwrap_or(settings.search_mode) + } else { + settings.search_mode + }; + let default_filter_mode = settings + .filter_mode_shell_up_key_binding + .filter(|_| settings.shell_up_key_binding) + .unwrap_or_else(|| settings.default_filter_mode(initial_context.git_root.is_some())); + let mut app = State { + history_count, + results_state: ListState::default(), + update_needed: None, + switched_search_mode: false, + search_mode, + tab_index: 0, + inspecting_state: InspectingState { + current: None, + next: None, + previous: None, + }, + keymaps: KeymapSet::from_settings(settings), + search: SearchState { + input, + filter_mode: default_filter_mode, + context: initial_context.clone(), + custom_context: None, + }, + engine: engines::engine(search_mode, settings), + results_len: 0, + accept: false, + keymap_mode: match settings.keymap_mode { + KeymapMode::Auto => KeymapMode::Emacs, + value => value, + }, + current_cursor: None, + now: if settings.prefers_reduced_motion { + let now = OffsetDateTime::now_utc(); + Box::new(move || now) + } else { + Box::new(OffsetDateTime::now_utc) + }, + prefix: false, + pending_vim_key: None, + original_input_empty: original_query.is_empty(), + }; + + app.initialize_keymap_cursor(settings); + + let mut results = app.query_results(&mut db, settings.smart_sort).await?; + + if inline_height > 0 && !popup_mode { + terminal.clear()?; + } + + let mut stats: Option<HistoryStats> = None; + let mut inspecting: Option<History> = None; + let accept; + let result = 'render: loop { + terminal.draw(|f| { + app.draw( + f, + &results, + stats.clone(), + inspecting.as_ref(), + settings, + theme, + popup_mode, + ); + })?; + + let initial_input = app.search.input.as_str().to_owned(); + let initial_filter_mode = app.search.filter_mode; + let initial_search_mode = app.search_mode; + let initial_custom_context = app.search.custom_context.clone(); + + let event_ready = tokio::task::spawn_blocking(|| event::poll(Duration::from_millis(250))); + + tokio::select! { + event_ready = event_ready => { + if event_ready?? { + loop { + match app.handle_input(settings, &event::read()?, &mut std::io::stdout())? { + InputAction::Continue => {}, + InputAction::Delete(index) => { + if results.is_empty() { + break; + } + app.results_len -= 1; + let selected = app.results_state.selected(); + if selected == app.results_len { + app.inspecting_state.reset(); + app.results_state.select(selected - 1); + } + + let entry = results.remove(index); + + if settings.sync.records { + let (id, _) = history_store.delete(entry.id).await?; + history_store.incremental_build(&db, &[id]).await?; + } else { + db.delete(entry.clone()).await?; + } + + app.tab_index = 0; + }, + InputAction::SwitchContext(index) => { + if let Some(index) = index && let Some(entry) = results.get(index) { + app.search.custom_context = Some(entry.id.clone()); + app.search.context = Context::from_history(entry); + app.search.filter_mode = FilterMode::Session; + app.search.input = Cursor::from(String::new()); + app.results_state = ListState::default(); + } else { + app.search.custom_context = None; + app.search.context = initial_context.clone(); + app.search.filter_mode = default_filter_mode; + } + }, + InputAction::Redraw => { + if !popup_mode { + terminal.clear()?; + } + terminal.draw(|f| { + app.draw(f, &results, stats.clone(), inspecting.as_ref(), settings, theme, popup_mode); + })?; + }, + r => { + accept = app.accept; + break 'render r; + }, + } + if !event::poll(Duration::ZERO)? { + break; + } + } + } + } + update_needed = &mut update_needed => { + // Don't fail interactive search if update check fails + // The update check is a nice-to-have feature, not critical + app.update_needed = update_needed.ok().flatten(); + } + } + + if initial_input != app.search.input.as_str() + || initial_filter_mode != app.search.filter_mode + || initial_search_mode != app.search_mode + || initial_custom_context != app.search.custom_context + { + results = app.query_results(&mut db, settings.smart_sort).await?; + } + + // In custom context mode, when no filter is applied, highlight the entry which was used + // to enter the context when changing modes. This helps to find your way around. + if app.search.custom_context.is_some() + && app.search.input.as_str().is_empty() + && (initial_custom_context != app.search.custom_context + || initial_filter_mode != app.search.filter_mode) + && let Some(history_id) = app.search.custom_context.clone() + && let Some(pos) = results.iter().position(|entry| entry.id == history_id) + { + app.results_state.select(pos); + } + + let inspecting_id = app.inspecting_state.clone().current; + // If inspecting ID is not the current inspecting History, update it. + match inspecting_id { + Some(inspecting_id) => { + if inspecting.is_none() || inspecting_id != inspecting.clone().unwrap().id { + inspecting = db.load(inspecting_id.0.as_str()).await?; + } + } + _ => { + inspecting = None; + } + } + + stats = if app.tab_index == 0 { + None + } else if !results.is_empty() { + // If we have stats, then we can indicate next available IDs. This avoids passing + // around a database object, or a full stats object. + let selected = match inspecting.clone() { + Some(insp) => insp, + None => results[app.results_state.selected()].clone(), + }; + let stats = db.stats(&selected).await?; + app.inspecting_state.current = Some(selected.id); + app.inspecting_state.previous = match stats.previous.clone() { + Some(p) => Some(p.id), + _ => None, + }; + app.inspecting_state.next = match stats.next.clone() { + Some(p) => Some(p.id), + _ => None, + }; + Some(stats) + } else { + None + }; + }; + + app.finalize_keymap_cursor(settings); + + if popup_mode { + // In popup mode, restore the screen area that was covered by the popup. + // This must happen before Stdout is dropped (which disables raw mode). + #[cfg(unix)] + if let Some(ref saved) = saved_screen { + restore_popup_area(saved, popup_rect, popup_scroll_offset); + } + } else if inline_height > 0 { + terminal.clear()?; + } + + let accept = accept + && matches!( + Shell::from_env(), + Shell::Zsh | Shell::Fish | Shell::Bash | Shell::Xonsh | Shell::Nu | Shell::Powershell + ); + + let accept_prefix = "__atuin_accept__:"; + + match result { + InputAction::AcceptInspecting => { + match inspecting { + Some(result) => { + let mut command = result.command; + + if accept { + command = String::from(accept_prefix) + &command; + } + + // index is in bounds so we return that entry + Ok(command) + } + None => Ok(String::new()), + } + } + InputAction::Accept(index) if index < results.len() => { + let mut command = results.swap_remove(index).command; + + if is_command_chaining { + command = format!("{} {}", original_query.trim_end(), command); + } else if accept { + command = String::from(accept_prefix) + &command; + } + + // index is in bounds so we return that entry + Ok(command) + } + InputAction::ReturnOriginal => Ok(String::new()), + InputAction::Copy(index) => { + let cmd = results.swap_remove(index).command; + set_clipboard(cmd); + Ok(String::new()) + } + InputAction::ReturnQuery | InputAction::Accept(_) => { + // Either: + // * index == RETURN_QUERY, in which case we should return the input + // * out of bounds -> usually implies no selected entry so we return the input + Ok(app.search.input.into_inner()) + } + InputAction::Continue + | InputAction::Redraw + | InputAction::Delete(_) + | InputAction::SwitchContext(_) => { + unreachable!("should have been handled!") + } + } +} + +// cli-clipboard only works on Windows, Mac, and Linux. + +#[cfg(all( + feature = "clipboard", + any(target_os = "windows", target_os = "macos", target_os = "linux") +))] +fn set_clipboard(s: String) { + let mut ctx = arboard::Clipboard::new().unwrap(); + ctx.set_text(s).unwrap(); + // Use the clipboard context to make sure it is saved + ctx.get_text().unwrap(); +} + +#[cfg(not(all( + feature = "clipboard", + any(target_os = "windows", target_os = "macos", target_os = "linux") +)))] +fn set_clipboard(_s: String) {} + +#[cfg(test)] +mod tests { + use atuin_client::database::Context; + use atuin_client::history::History; + use atuin_client::settings::{ + FilterMode, KeymapMode, Preview, PreviewStrategy, SearchMode, Settings, + }; + use time::OffsetDateTime; + + use crate::command::client::search::engines::{self, SearchState}; + use crate::command::client::search::history_list::ListState; + + use super::{Compactness, InspectingState, KeymapSet, State}; + + #[test] + #[allow(clippy::too_many_lines)] + fn calc_preview_height_test() { + let settings_preview_auto = Settings { + preview: Preview { + strategy: PreviewStrategy::Auto, + }, + show_preview: true, + ..Settings::utc() + }; + + let settings_preview_auto_h2 = Settings { + preview: Preview { + strategy: PreviewStrategy::Auto, + }, + show_preview: true, + max_preview_height: 2, + ..Settings::utc() + }; + + let settings_preview_h4 = Settings { + preview: Preview { + strategy: PreviewStrategy::Static, + }, + show_preview: true, + max_preview_height: 4, + ..Settings::utc() + }; + + let settings_preview_fixed = Settings { + preview: Preview { + strategy: PreviewStrategy::Fixed, + }, + show_preview: true, + max_preview_height: 15, + ..Settings::utc() + }; + + let cmd_60: History = History::capture() + .timestamp(time::OffsetDateTime::now_utc()) + .command("for i in $(seq -w 10); do echo \"item number $i - abcd\"; done") + .cwd("/") + .build() + .into(); + + let cmd_124: History = History::capture() + .timestamp(time::OffsetDateTime::now_utc()) + .command("echo 'Aurea prima sata est aetas, quae vindice nullo, sponte sua, sine lege fidem rectumque colebat. Poena metusque aberant'") + .cwd("/") + .build() + .into(); + + let cmd_200: History = History::capture() + .timestamp(time::OffsetDateTime::now_utc()) + .command("CREATE USER atuin WITH ENCRYPTED PASSWORD 'supersecretpassword'; CREATE DATABASE atuin WITH OWNER = atuin; \\c atuin; REVOKE ALL PRIVILEGES ON SCHEMA public FROM PUBLIC; echo 'All done. 200 characters'") + .cwd("/") + .build() + .into(); + + let results: Vec<History> = vec![cmd_60, cmd_124, cmd_200]; + + // the selected command does not require a preview + let no_preview = State::calc_preview_height( + &settings_preview_auto, + &results, + 0_usize, + 0_usize, + Compactness::Full, + 1, + 80, + ); + // the selected command requires 2 lines + let preview_h2 = State::calc_preview_height( + &settings_preview_auto, + &results, + 1_usize, + 0_usize, + Compactness::Full, + 1, + 80, + ); + // the selected command requires 3 lines + let preview_h3 = State::calc_preview_height( + &settings_preview_auto, + &results, + 2_usize, + 0_usize, + Compactness::Full, + 1, + 80, + ); + // the selected command requires a preview of 1 line (happens when the command is between preview_width-19 and preview_width) + let preview_one_line = State::calc_preview_height( + &settings_preview_auto, + &results, + 0_usize, + 0_usize, + Compactness::Full, + 1, + 66, + ); + // the selected command requires 3 lines, but we have a max preview height limit of 2 + let preview_limit_at_2 = State::calc_preview_height( + &settings_preview_auto_h2, + &results, + 2_usize, + 0_usize, + Compactness::Full, + 1, + 80, + ); + // the longest command requires 3 lines + let preview_static_h3 = State::calc_preview_height( + &settings_preview_h4, + &results, + 1_usize, + 0_usize, + Compactness::Full, + 1, + 80, + ); + // the longest command requires 10 lines, but we have a max preview height limit of 4 + let preview_static_limit_at_4 = State::calc_preview_height( + &settings_preview_h4, + &results, + 1_usize, + 0_usize, + Compactness::Full, + 1, + 20, + ); + // the longest command requires 10 lines, but we have a max preview height of 15 and a fixed preview strategy + let settings_preview_fixed = State::calc_preview_height( + &settings_preview_fixed, + &results, + 1_usize, + 0_usize, + Compactness::Full, + 1, + 20, + ); + + assert_eq!(no_preview, 1); + // 1 * 2 is the space for the border + let border_space = 2; + assert_eq!(preview_h2, 2 + border_space); + assert_eq!(preview_h3, 3 + border_space); + assert_eq!(preview_one_line, 1 + border_space); + assert_eq!(preview_limit_at_2, 2 + border_space); + assert_eq!(preview_static_h3, 3 + border_space); + assert_eq!(preview_static_limit_at_4, 4 + border_space); + assert_eq!(settings_preview_fixed, 15 + border_space); + } + + // Test when there's no results, scrolling up or down doesn't underflow + #[test] + fn state_scroll_up_underflow() { + let settings = Settings::utc(); + let mut state = State { + history_count: 0, + update_needed: None, + results_state: ListState::default(), + switched_search_mode: false, + search_mode: SearchMode::Fuzzy, + results_len: 0, + accept: false, + keymap_mode: KeymapMode::Auto, + prefix: false, + current_cursor: None, + tab_index: 0, + pending_vim_key: None, + original_input_empty: false, + inspecting_state: InspectingState { + current: None, + next: None, + previous: None, + }, + keymaps: KeymapSet::defaults(&settings), + search: SearchState { + input: String::new().into(), + filter_mode: FilterMode::Directory, + context: Context { + session: String::new(), + cwd: String::new(), + hostname: String::new(), + host_id: String::new(), + git_root: None, + }, + custom_context: None, + }, + engine: engines::engine(SearchMode::Fuzzy, &settings), + now: Box::new(OffsetDateTime::now_utc), + }; + + state.scroll_up(1); + state.scroll_down(1); + } + + #[test] + fn test_accept_keybindings() { + use atuin_client::settings::Keys; + use ratatui::crossterm::event::{KeyCode, KeyEvent, KeyModifiers}; + + let mut settings = Settings::utc(); + settings.keys = Keys { + scroll_exits: true, + exit_past_line_start: false, + accept_past_line_end: true, + accept_past_line_start: false, + accept_with_backspace: false, + prefix: "a".to_string(), + }; + + let mut state = State { + history_count: 1, + update_needed: None, + results_state: ListState::default(), + switched_search_mode: false, + search_mode: SearchMode::Fuzzy, + results_len: 1, + accept: false, + keymap_mode: KeymapMode::Emacs, + prefix: false, + current_cursor: None, + tab_index: 0, + pending_vim_key: None, + original_input_empty: false, + inspecting_state: InspectingState { + current: None, + next: None, + previous: None, + }, + keymaps: KeymapSet::defaults(&settings), + search: SearchState { + input: String::new().into(), + filter_mode: FilterMode::Global, + context: Context { + session: String::new(), + cwd: String::new(), + hostname: String::new(), + host_id: String::new(), + git_root: None, + }, + custom_context: None, + }, + engine: engines::engine(SearchMode::Fuzzy, &settings), + now: Box::new(OffsetDateTime::now_utc), + }; + + let tab_event = KeyEvent::new(KeyCode::Tab, KeyModifiers::NONE); + let result = state.handle_key_input(&settings, &tab_event); + assert!( + matches!(result, super::InputAction::Accept(_)), + "Tab should always accept" + ); + + // Test left arrow with accept_past_line_start disabled (should continue) + let left_event = KeyEvent::new(KeyCode::Left, KeyModifiers::NONE); + let result = state.handle_key_input(&settings, &left_event); + assert!( + matches!(result, super::InputAction::Continue), + "Left arrow should continue when disabled" + ); + + // Test left arrow with accept_past_line_start enabled (should accept at start of line) + settings.keys.accept_past_line_start = true; + state.keymaps = KeymapSet::defaults(&settings); + let result = state.handle_key_input(&settings, &left_event); + assert!( + matches!(result, super::InputAction::Accept(_)), + "Left arrow should accept at start of line when enabled" + ); + settings.keys.accept_past_line_start = false; + state.keymaps = KeymapSet::defaults(&settings); + + let backspace_event = KeyEvent::new(KeyCode::Backspace, KeyModifiers::NONE); + let result = state.handle_key_input(&settings, &backspace_event); + assert!( + matches!(result, super::InputAction::Continue), + "Backspace should continue when disabled" + ); + + settings.keys.accept_with_backspace = true; + state.keymaps = KeymapSet::defaults(&settings); + let result = state.handle_key_input(&settings, &backspace_event); + assert!( + matches!(result, super::InputAction::Accept(_)), + "Backspace should accept at start of line when enabled" + ); + + state.search.input.insert('t'); + state.search.input.insert('e'); + state.search.input.insert('s'); + state.search.input.insert('t'); + state.search.input.end(); + + let right_event = KeyEvent::new(KeyCode::Right, KeyModifiers::NONE); + let result = state.handle_key_input(&settings, &right_event); + assert!( + matches!(result, super::InputAction::Accept(_)), + "Right arrow should accept at end of line when enabled" + ); + + settings.keys.accept_past_line_start = true; + state.keymaps = KeymapSet::defaults(&settings); + let left_event = KeyEvent::new(KeyCode::Left, KeyModifiers::NONE); + let result = state.handle_key_input(&settings, &left_event); + assert!( + matches!(result, super::InputAction::Continue), + "Left arrow should continue and end of line, even when enabled" + ); + settings.keys.accept_past_line_start = false; + state.keymaps = KeymapSet::defaults(&settings); + + settings.keys.accept_with_backspace = true; + state.keymaps = KeymapSet::defaults(&settings); + let backspace_event = KeyEvent::new(KeyCode::Backspace, KeyModifiers::NONE); + let result = state.handle_key_input(&settings, &backspace_event); + assert!( + matches!(result, super::InputAction::Continue), + "Backspace should continue at end of line, even when enabled" + ); + settings.keys.accept_with_backspace = false; + state.keymaps = KeymapSet::defaults(&settings); + } + + #[test] + fn test_vim_gg_multikey_sequence() { + use ratatui::crossterm::event::{KeyCode, KeyEvent, KeyModifiers}; + + let settings = Settings::utc(); + + let mut state = State { + history_count: 100, + update_needed: None, + results_state: ListState::default(), + switched_search_mode: false, + search_mode: SearchMode::Fuzzy, + results_len: 100, + accept: false, + keymap_mode: KeymapMode::VimNormal, + prefix: false, + current_cursor: None, + tab_index: 0, + pending_vim_key: None, + original_input_empty: false, + inspecting_state: InspectingState { + current: None, + next: None, + previous: None, + }, + keymaps: KeymapSet::defaults(&settings), + search: SearchState { + input: String::new().into(), + filter_mode: FilterMode::Global, + context: Context { + session: String::new(), + cwd: String::new(), + hostname: String::new(), + host_id: String::new(), + git_root: None, + }, + custom_context: None, + }, + engine: engines::engine(SearchMode::Fuzzy, &settings), + now: Box::new(OffsetDateTime::now_utc), + }; + + // Start in the middle of the list + state.results_state.select(50); + + // First 'g' should set pending state + let g_event = KeyEvent::new(KeyCode::Char('g'), KeyModifiers::NONE); + let result = state.handle_key_input(&settings, &g_event); + assert!(matches!(result, super::InputAction::Continue)); + assert_eq!(state.pending_vim_key, Some('g')); + assert_eq!(state.results_state.selected(), 50); // Position unchanged + + // Second 'g' should jump to end (visual top in non-inverted mode) + let result = state.handle_key_input(&settings, &g_event); + assert!(matches!(result, super::InputAction::Continue)); + assert_eq!(state.pending_vim_key, None); + assert_eq!(state.results_state.selected(), 99); // Jumped to last index (visual top) + } + + #[test] + fn test_vim_g_key_clears_on_other_input() { + use ratatui::crossterm::event::{KeyCode, KeyEvent, KeyModifiers}; + + let settings = Settings::utc(); + + let mut state = State { + history_count: 100, + update_needed: None, + results_state: ListState::default(), + switched_search_mode: false, + search_mode: SearchMode::Fuzzy, + results_len: 100, + accept: false, + keymap_mode: KeymapMode::VimNormal, + prefix: false, + current_cursor: None, + tab_index: 0, + pending_vim_key: None, + original_input_empty: false, + inspecting_state: InspectingState { + current: None, + next: None, + previous: None, + }, + keymaps: KeymapSet::defaults(&settings), + search: SearchState { + input: String::new().into(), + filter_mode: FilterMode::Global, + context: Context { + session: String::new(), + cwd: String::new(), + hostname: String::new(), + host_id: String::new(), + git_root: None, + }, + custom_context: None, + }, + engine: engines::engine(SearchMode::Fuzzy, &settings), + now: Box::new(OffsetDateTime::now_utc), + }; + + state.results_state.select(50); + + // Press 'g' to set pending state + let g_event = KeyEvent::new(KeyCode::Char('g'), KeyModifiers::NONE); + state.handle_key_input(&settings, &g_event); + assert_eq!(state.pending_vim_key, Some('g')); + + // Press 'j' - should clear pending state + let j_event = KeyEvent::new(KeyCode::Char('j'), KeyModifiers::NONE); + state.handle_key_input(&settings, &j_event); + assert_eq!(state.pending_vim_key, None); + } + + #[test] + fn test_vim_big_g_jump_to_bottom() { + use ratatui::crossterm::event::{KeyCode, KeyEvent, KeyModifiers}; + + let settings = Settings::utc(); + + let mut state = State { + history_count: 100, + update_needed: None, + results_state: ListState::default(), + switched_search_mode: false, + search_mode: SearchMode::Fuzzy, + results_len: 100, + accept: false, + keymap_mode: KeymapMode::VimNormal, + prefix: false, + current_cursor: None, + tab_index: 0, + pending_vim_key: None, + original_input_empty: false, + inspecting_state: InspectingState { + current: None, + next: None, + previous: None, + }, + keymaps: KeymapSet::defaults(&settings), + search: SearchState { + input: String::new().into(), + filter_mode: FilterMode::Global, + context: Context { + session: String::new(), + cwd: String::new(), + hostname: String::new(), + host_id: String::new(), + git_root: None, + }, + custom_context: None, + }, + engine: engines::engine(SearchMode::Fuzzy, &settings), + now: Box::new(OffsetDateTime::now_utc), + }; + + state.results_state.select(50); + + // 'G' should jump to visual bottom (index 0 in non-inverted mode) + let big_g_event = KeyEvent::new(KeyCode::Char('G'), KeyModifiers::NONE); + let result = state.handle_key_input(&settings, &big_g_event); + assert!(matches!(result, super::InputAction::Continue)); + assert_eq!(state.results_state.selected(), 0); + } + + #[test] + fn test_vim_ctrl_u_d_half_page_scroll() { + use ratatui::crossterm::event::{KeyCode, KeyEvent, KeyModifiers}; + + let settings = Settings::utc(); + + let mut state = State { + history_count: 100, + update_needed: None, + results_state: ListState::default(), + switched_search_mode: false, + search_mode: SearchMode::Fuzzy, + results_len: 100, + accept: false, + keymap_mode: KeymapMode::VimNormal, + prefix: false, + current_cursor: None, + tab_index: 0, + pending_vim_key: None, + original_input_empty: false, + inspecting_state: InspectingState { + current: None, + next: None, + previous: None, + }, + keymaps: KeymapSet::defaults(&settings), + search: SearchState { + input: String::new().into(), + filter_mode: FilterMode::Global, + context: Context { + session: String::new(), + cwd: String::new(), + hostname: String::new(), + host_id: String::new(), + git_root: None, + }, + custom_context: None, + }, + engine: engines::engine(SearchMode::Fuzzy, &settings), + now: Box::new(OffsetDateTime::now_utc), + }; + + state.results_state.select(50); + + // Ctrl+d should return Continue and clear pending key + // (scroll amount depends on max_entries which is 0 in tests) + state.pending_vim_key = Some('g'); + let ctrl_d_event = KeyEvent::new(KeyCode::Char('d'), KeyModifiers::CONTROL); + let result = state.handle_key_input(&settings, &ctrl_d_event); + assert!(matches!(result, super::InputAction::Continue)); + assert_eq!(state.pending_vim_key, None); + + // Ctrl+u should return Continue and clear pending key + state.pending_vim_key = Some('g'); + let ctrl_u_event = KeyEvent::new(KeyCode::Char('u'), KeyModifiers::CONTROL); + let result = state.handle_key_input(&settings, &ctrl_u_event); + assert!(matches!(result, super::InputAction::Continue)); + assert_eq!(state.pending_vim_key, None); + } + + #[test] + fn test_vim_ctrl_f_b_full_page_scroll() { + use ratatui::crossterm::event::{KeyCode, KeyEvent, KeyModifiers}; + + let settings = Settings::utc(); + + let mut state = State { + history_count: 100, + update_needed: None, + results_state: ListState::default(), + switched_search_mode: false, + search_mode: SearchMode::Fuzzy, + results_len: 100, + accept: false, + keymap_mode: KeymapMode::VimNormal, + prefix: false, + current_cursor: None, + tab_index: 0, + pending_vim_key: None, + original_input_empty: false, + inspecting_state: InspectingState { + current: None, + next: None, + previous: None, + }, + keymaps: KeymapSet::defaults(&settings), + search: SearchState { + input: String::new().into(), + filter_mode: FilterMode::Global, + context: Context { + session: String::new(), + cwd: String::new(), + hostname: String::new(), + host_id: String::new(), + git_root: None, + }, + custom_context: None, + }, + engine: engines::engine(SearchMode::Fuzzy, &settings), + now: Box::new(OffsetDateTime::now_utc), + }; + + state.results_state.select(50); + + // Ctrl+f should return Continue and clear pending key + // (scroll amount depends on max_entries which is 0 in tests) + state.pending_vim_key = Some('g'); + let ctrl_f_event = KeyEvent::new(KeyCode::Char('f'), KeyModifiers::CONTROL); + let result = state.handle_key_input(&settings, &ctrl_f_event); + assert!(matches!(result, super::InputAction::Continue)); + assert_eq!(state.pending_vim_key, None); + + // Ctrl+b should return Continue and clear pending key + state.pending_vim_key = Some('g'); + let ctrl_b_event = KeyEvent::new(KeyCode::Char('b'), KeyModifiers::CONTROL); + let result = state.handle_key_input(&settings, &ctrl_b_event); + assert!(matches!(result, super::InputAction::Continue)); + assert_eq!(state.pending_vim_key, None); + } + + // ----------------------------------------------------------------------- + // Executor tests (execute_action) + // ----------------------------------------------------------------------- + + /// Helper to build a State for executor tests. + fn make_executor_state(results_len: usize, selected: usize) -> State { + let settings = Settings::utc(); + let mut state = State { + history_count: results_len as i64, + update_needed: None, + results_state: ListState::default(), + switched_search_mode: false, + search_mode: SearchMode::Fuzzy, + results_len, + accept: false, + keymap_mode: KeymapMode::Emacs, + prefix: false, + current_cursor: None, + tab_index: 0, + pending_vim_key: None, + original_input_empty: false, + inspecting_state: InspectingState { + current: None, + next: None, + previous: None, + }, + keymaps: KeymapSet::defaults(&settings), + search: SearchState { + input: String::new().into(), + filter_mode: FilterMode::Global, + context: Context { + session: String::new(), + cwd: String::new(), + hostname: String::new(), + host_id: String::new(), + git_root: None, + }, + custom_context: None, + }, + engine: engines::engine(SearchMode::Fuzzy, &settings), + now: Box::new(OffsetDateTime::now_utc), + }; + state.results_state.select(selected); + state + } + + #[test] + fn execute_select_next_no_invert() { + use crate::command::client::search::keybindings::Action; + + let mut state = make_executor_state(100, 50); + let settings = Settings::utc(); + let result = state.execute_action(&Action::SelectNext, &settings); + assert!(matches!(result, super::InputAction::Continue)); + // Non-inverted: SelectNext = scroll_down = selected - 1 + assert_eq!(state.results_state.selected(), 49); + } + + #[test] + fn execute_select_next_with_invert() { + use crate::command::client::search::keybindings::Action; + + let mut state = make_executor_state(100, 50); + let mut settings = Settings::utc(); + settings.invert = true; + let result = state.execute_action(&Action::SelectNext, &settings); + assert!(matches!(result, super::InputAction::Continue)); + // Inverted: SelectNext = scroll_up = selected + 1 + assert_eq!(state.results_state.selected(), 51); + } + + #[test] + fn execute_select_previous_no_invert() { + use crate::command::client::search::keybindings::Action; + + let mut state = make_executor_state(100, 50); + let settings = Settings::utc(); + let result = state.execute_action(&Action::SelectPrevious, &settings); + assert!(matches!(result, super::InputAction::Continue)); + // Non-inverted: SelectPrevious = scroll_up = selected + 1 + assert_eq!(state.results_state.selected(), 51); + } + + #[test] + fn execute_vim_enter_normal() { + use crate::command::client::search::keybindings::Action; + + let mut state = make_executor_state(100, 0); + let settings = Settings::utc(); + let result = state.execute_action(&Action::VimEnterNormal, &settings); + assert!(matches!(result, super::InputAction::Continue)); + assert_eq!(state.keymap_mode, KeymapMode::VimNormal); + } + + #[test] + fn execute_vim_enter_insert() { + use crate::command::client::search::keybindings::Action; + + let mut state = make_executor_state(100, 0); + state.keymap_mode = KeymapMode::VimNormal; + let settings = Settings::utc(); + let result = state.execute_action(&Action::VimEnterInsert, &settings); + assert!(matches!(result, super::InputAction::Continue)); + assert_eq!(state.keymap_mode, KeymapMode::VimInsert); + } + + #[test] + fn execute_accept_sets_accept_flag() { + use crate::command::client::search::keybindings::Action; + + let mut state = make_executor_state(100, 5); + let mut settings = Settings::utc(); + settings.enter_accept = true; + let result = state.execute_action(&Action::Accept, &settings); + assert!(matches!(result, super::InputAction::Accept(5))); + assert!(state.accept); + } + + #[test] + fn execute_return_selection_does_not_set_accept() { + use crate::command::client::search::keybindings::Action; + + let mut state = make_executor_state(100, 5); + let settings = Settings::utc(); + let result = state.execute_action(&Action::ReturnSelection, &settings); + assert!(matches!(result, super::InputAction::Accept(5))); + assert!(!state.accept); + } + + #[test] + fn execute_accept_nth() { + use crate::command::client::search::keybindings::Action; + + let mut state = make_executor_state(100, 5); + let settings = Settings::utc(); + let result = state.execute_action(&Action::AcceptNth(3), &settings); + assert!(matches!(result, super::InputAction::Accept(8))); + } + + #[test] + fn execute_scroll_to_top_no_invert() { + use crate::command::client::search::keybindings::Action; + + let mut state = make_executor_state(100, 50); + let settings = Settings::utc(); + let result = state.execute_action(&Action::ScrollToTop, &settings); + assert!(matches!(result, super::InputAction::Continue)); + // Non-inverted: visual top = highest index + assert_eq!(state.results_state.selected(), 99); + } + + #[test] + fn execute_scroll_to_top_with_invert() { + use crate::command::client::search::keybindings::Action; + + let mut state = make_executor_state(100, 50); + let mut settings = Settings::utc(); + settings.invert = true; + let result = state.execute_action(&Action::ScrollToTop, &settings); + assert!(matches!(result, super::InputAction::Continue)); + // Inverted: visual top = index 0 + assert_eq!(state.results_state.selected(), 0); + } + + #[test] + fn execute_scroll_to_bottom_no_invert() { + use crate::command::client::search::keybindings::Action; + + let mut state = make_executor_state(100, 50); + let settings = Settings::utc(); + let result = state.execute_action(&Action::ScrollToBottom, &settings); + assert!(matches!(result, super::InputAction::Continue)); + // Non-inverted: visual bottom = index 0 + assert_eq!(state.results_state.selected(), 0); + } + + #[test] + fn execute_toggle_tab() { + use crate::command::client::search::keybindings::Action; + + let mut state = make_executor_state(100, 0); + let settings = Settings::utc(); + assert_eq!(state.tab_index, 0); + state.execute_action(&Action::ToggleTab, &settings); + assert_eq!(state.tab_index, 1); + state.execute_action(&Action::ToggleTab, &settings); + assert_eq!(state.tab_index, 0); + } + + #[test] + fn execute_enter_prefix_mode() { + use crate::command::client::search::keybindings::Action; + + let mut state = make_executor_state(100, 0); + let settings = Settings::utc(); + assert!(!state.prefix); + state.execute_action(&Action::EnterPrefixMode, &settings); + assert!(state.prefix); + } + + #[test] + fn execute_exit_returns_based_on_exit_mode() { + use crate::command::client::search::keybindings::Action; + use atuin_client::settings::ExitMode; + + let mut state = make_executor_state(100, 0); + let mut settings = Settings::utc(); + + settings.exit_mode = ExitMode::ReturnOriginal; + let result = state.execute_action(&Action::Exit, &settings); + assert!(matches!(result, super::InputAction::ReturnOriginal)); + + settings.exit_mode = ExitMode::ReturnQuery; + let result = state.execute_action(&Action::Exit, &settings); + assert!(matches!(result, super::InputAction::ReturnQuery)); + } + + #[test] + fn execute_return_original() { + use crate::command::client::search::keybindings::Action; + + let mut state = make_executor_state(100, 0); + let settings = Settings::utc(); + let result = state.execute_action(&Action::ReturnOriginal, &settings); + assert!(matches!(result, super::InputAction::ReturnOriginal)); + } + + #[test] + fn execute_copy() { + use crate::command::client::search::keybindings::Action; + + let mut state = make_executor_state(100, 7); + let settings = Settings::utc(); + let result = state.execute_action(&Action::Copy, &settings); + assert!(matches!(result, super::InputAction::Copy(7))); + } + + #[test] + fn execute_delete() { + use crate::command::client::search::keybindings::Action; + + let mut state = make_executor_state(100, 7); + let settings = Settings::utc(); + let result = state.execute_action(&Action::Delete, &settings); + assert!(matches!(result, super::InputAction::Delete(7))); + } + + #[test] + fn execute_switch_context() { + use crate::command::client::search::keybindings::Action; + + let mut state = make_executor_state(100, 7); + let settings = Settings::utc(); + let result = state.execute_action(&Action::SwitchContext, &settings); + assert!(matches!(result, super::InputAction::SwitchContext(Some(7)))); + } + + #[test] + fn execute_clear_context() { + use crate::command::client::search::keybindings::Action; + + let mut state = make_executor_state(100, 7); + let settings = Settings::utc(); + let result = state.execute_action(&Action::ClearContext, &settings); + assert!(matches!(result, super::InputAction::SwitchContext(None))); + } + + #[test] + fn execute_noop() { + use crate::command::client::search::keybindings::Action; + + let mut state = make_executor_state(100, 50); + let settings = Settings::utc(); + let result = state.execute_action(&Action::Noop, &settings); + assert!(matches!(result, super::InputAction::Continue)); + assert_eq!(state.results_state.selected(), 50); + } + + #[test] + fn execute_accept_in_inspector_tab() { + use crate::command::client::search::keybindings::Action; + + let mut state = make_executor_state(100, 5); + state.tab_index = 1; + let settings = Settings::utc(); + let result = state.execute_action(&Action::Accept, &settings); + assert!(matches!(result, super::InputAction::AcceptInspecting)); + } + + #[test] + fn execute_cycle_search_mode() { + use crate::command::client::search::keybindings::Action; + + let mut state = make_executor_state(100, 0); + let settings = Settings::utc(); + let original_mode = state.search_mode; + let result = state.execute_action(&Action::CycleSearchMode, &settings); + assert!(matches!(result, super::InputAction::Continue)); + assert!(state.switched_search_mode); + assert_ne!(state.search_mode, original_mode); + } + + #[test] + fn execute_vim_search_insert() { + use crate::command::client::search::keybindings::Action; + + let mut state = make_executor_state(100, 0); + state.search.input.insert('h'); + state.search.input.insert('i'); + state.keymap_mode = KeymapMode::VimNormal; + let settings = Settings::utc(); + let result = state.execute_action(&Action::VimSearchInsert, &settings); + assert!(matches!(result, super::InputAction::Continue)); + // Should clear input and switch to insert mode + assert_eq!(state.search.input.as_str(), ""); + assert_eq!(state.keymap_mode, KeymapMode::VimInsert); + } + + #[test] + fn execute_cursor_movement() { + use crate::command::client::search::keybindings::Action; + + let mut state = make_executor_state(100, 0); + let settings = Settings::utc(); + + // Insert some text + state.search.input.insert('h'); + state.search.input.insert('e'); + state.search.input.insert('l'); + state.search.input.insert('l'); + state.search.input.insert('o'); + // cursor is at end (position 5) + + // CursorLeft + state.execute_action(&Action::CursorLeft, &settings); + assert_eq!(state.search.input.position(), 4); + + // CursorStart + state.execute_action(&Action::CursorStart, &settings); + assert_eq!(state.search.input.position(), 0); + + // CursorEnd + state.execute_action(&Action::CursorEnd, &settings); + assert_eq!(state.search.input.position(), 5); + + // CursorRight at end does nothing + state.execute_action(&Action::CursorRight, &settings); + assert_eq!(state.search.input.position(), 5); + } + + #[test] + fn execute_editing() { + use crate::command::client::search::keybindings::Action; + + let mut state = make_executor_state(100, 0); + let settings = Settings::utc(); + + // Insert "hello" + state.search.input.insert('h'); + state.search.input.insert('e'); + state.search.input.insert('l'); + state.search.input.insert('l'); + state.search.input.insert('o'); + + // DeleteCharBefore (backspace) + state.execute_action(&Action::DeleteCharBefore, &settings); + assert_eq!(state.search.input.as_str(), "hell"); + + // ClearLine + state.execute_action(&Action::ClearLine, &settings); + assert_eq!(state.search.input.as_str(), ""); + } + + #[test] + fn keymap_config_return_query() { + use atuin_client::settings::KeyBindingConfig; + use ratatui::crossterm::event::{KeyCode, KeyEvent, KeyModifiers}; + use std::collections::HashMap; + + let mut settings = Settings::utc(); + // Configure tab to return-query + settings.keymap.emacs = HashMap::from([( + "tab".to_string(), + KeyBindingConfig::Simple("return-query".to_string()), + )]); + + let mut state = State { + history_count: 100, + update_needed: None, + results_state: ListState::default(), + switched_search_mode: false, + search_mode: SearchMode::Fuzzy, + results_len: 100, + accept: false, + keymap_mode: KeymapMode::Emacs, + prefix: false, + current_cursor: None, + tab_index: 0, + pending_vim_key: None, + original_input_empty: false, + inspecting_state: InspectingState { + current: None, + next: None, + previous: None, + }, + keymaps: KeymapSet::from_settings(&settings), + search: SearchState { + input: "test query".to_string().into(), + filter_mode: FilterMode::Global, + context: Context { + session: String::new(), + cwd: String::new(), + hostname: String::new(), + host_id: String::new(), + git_root: None, + }, + custom_context: None, + }, + engine: engines::engine(SearchMode::Fuzzy, &settings), + now: Box::new(OffsetDateTime::now_utc), + }; + + let tab_event = KeyEvent::new(KeyCode::Tab, KeyModifiers::NONE); + let result = state.handle_key_input(&settings, &tab_event); + assert!( + matches!(result, super::InputAction::ReturnQuery), + "Tab configured as return-query should return InputAction::ReturnQuery" + ); + } +} diff --git a/crates/atuin/src/command/client/search/keybindings/actions.rs b/crates/atuin/src/command/client/search/keybindings/actions.rs new file mode 100644 index 00000000..66e2709e --- /dev/null +++ b/crates/atuin/src/command/client/search/keybindings/actions.rs @@ -0,0 +1,319 @@ +use std::fmt; + +use serde::{Deserialize, Deserializer, Serialize, Serializer}; + +/// All possible actions that can be triggered by a keybinding. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum Action { + // Cursor movement + CursorLeft, + CursorRight, + CursorWordLeft, + CursorWordRight, + CursorWordEnd, + CursorStart, + CursorEnd, + + // Editing + DeleteCharBefore, + DeleteCharAfter, + DeleteWordBefore, + DeleteWordAfter, + DeleteToWordBoundary, + ClearLine, + ClearToStart, + ClearToEnd, + + // List navigation + SelectNext, + SelectPrevious, + ScrollHalfPageUp, + ScrollHalfPageDown, + ScrollPageUp, + ScrollPageDown, + ScrollToTop, + ScrollToBottom, + ScrollToScreenTop, + ScrollToScreenMiddle, + ScrollToScreenBottom, + + // Commands — accept selection and execute immediately + Accept, + AcceptNth(u8), + // Commands — return selection to command line without executing + ReturnSelection, + ReturnSelectionNth(u8), + // Commands — other + Copy, + Delete, + ReturnOriginal, + ReturnQuery, + Exit, + Redraw, + CycleFilterMode, + CycleSearchMode, + SwitchContext, + ClearContext, + ToggleTab, + + // Mode changes + VimEnterNormal, + VimEnterInsert, + VimEnterInsertAfter, + VimEnterInsertAtStart, + VimEnterInsertAtEnd, + VimSearchInsert, + VimChangeToEnd, + EnterPrefixMode, + + // Inspector + InspectPrevious, + InspectNext, + + // Special + Noop, +} + +impl Action { + /// Convert from a kebab-case string. + pub fn from_str(s: &str) -> Result<Self, String> { + // Handle accept-N and return-selection-N patterns + if let Some(rest) = s.strip_prefix("accept-") + && let Ok(n) = rest.parse::<u8>() + && (1..=9).contains(&n) + { + return Ok(Action::AcceptNth(n)); + } + if let Some(rest) = s.strip_prefix("return-selection-") + && let Ok(n) = rest.parse::<u8>() + && (1..=9).contains(&n) + { + return Ok(Action::ReturnSelectionNth(n)); + } + + match s { + "cursor-left" => Ok(Action::CursorLeft), + "cursor-right" => Ok(Action::CursorRight), + "cursor-word-left" => Ok(Action::CursorWordLeft), + "cursor-word-right" => Ok(Action::CursorWordRight), + "cursor-word-end" => Ok(Action::CursorWordEnd), + "cursor-start" => Ok(Action::CursorStart), + "cursor-end" => Ok(Action::CursorEnd), + + "delete-char-before" => Ok(Action::DeleteCharBefore), + "delete-char-after" => Ok(Action::DeleteCharAfter), + "delete-word-before" => Ok(Action::DeleteWordBefore), + "delete-word-after" => Ok(Action::DeleteWordAfter), + "delete-to-word-boundary" => Ok(Action::DeleteToWordBoundary), + "clear-line" => Ok(Action::ClearLine), + "clear-to-start" => Ok(Action::ClearToStart), + "clear-to-end" => Ok(Action::ClearToEnd), + + "select-next" => Ok(Action::SelectNext), + "select-previous" => Ok(Action::SelectPrevious), + "scroll-half-page-up" => Ok(Action::ScrollHalfPageUp), + "scroll-half-page-down" => Ok(Action::ScrollHalfPageDown), + "scroll-page-up" => Ok(Action::ScrollPageUp), + "scroll-page-down" => Ok(Action::ScrollPageDown), + "scroll-to-top" => Ok(Action::ScrollToTop), + "scroll-to-bottom" => Ok(Action::ScrollToBottom), + "scroll-to-screen-top" => Ok(Action::ScrollToScreenTop), + "scroll-to-screen-middle" => Ok(Action::ScrollToScreenMiddle), + "scroll-to-screen-bottom" => Ok(Action::ScrollToScreenBottom), + + "accept" => Ok(Action::Accept), + "return-selection" => Ok(Action::ReturnSelection), + "copy" => Ok(Action::Copy), + "delete" => Ok(Action::Delete), + "return-original" => Ok(Action::ReturnOriginal), + "return-query" => Ok(Action::ReturnQuery), + "exit" => Ok(Action::Exit), + "redraw" => Ok(Action::Redraw), + "cycle-filter-mode" => Ok(Action::CycleFilterMode), + "cycle-search-mode" => Ok(Action::CycleSearchMode), + "switch-context" => Ok(Action::SwitchContext), + "clear-context" => Ok(Action::ClearContext), + "toggle-tab" => Ok(Action::ToggleTab), + + "vim-enter-normal" => Ok(Action::VimEnterNormal), + "vim-enter-insert" => Ok(Action::VimEnterInsert), + "vim-enter-insert-after" => Ok(Action::VimEnterInsertAfter), + "vim-enter-insert-at-start" => Ok(Action::VimEnterInsertAtStart), + "vim-enter-insert-at-end" => Ok(Action::VimEnterInsertAtEnd), + "vim-search-insert" => Ok(Action::VimSearchInsert), + "vim-change-to-end" => Ok(Action::VimChangeToEnd), + "enter-prefix-mode" => Ok(Action::EnterPrefixMode), + + "inspect-previous" => Ok(Action::InspectPrevious), + "inspect-next" => Ok(Action::InspectNext), + + "noop" => Ok(Action::Noop), + + _ => Err(format!("unknown action: {s}")), + } + } + + /// Convert to a kebab-case string. + pub fn as_str(&self) -> String { + match self { + Action::CursorLeft => "cursor-left".to_string(), + Action::CursorRight => "cursor-right".to_string(), + Action::CursorWordLeft => "cursor-word-left".to_string(), + Action::CursorWordRight => "cursor-word-right".to_string(), + Action::CursorWordEnd => "cursor-word-end".to_string(), + Action::CursorStart => "cursor-start".to_string(), + Action::CursorEnd => "cursor-end".to_string(), + + Action::DeleteCharBefore => "delete-char-before".to_string(), + Action::DeleteCharAfter => "delete-char-after".to_string(), + Action::DeleteWordBefore => "delete-word-before".to_string(), + Action::DeleteWordAfter => "delete-word-after".to_string(), + Action::DeleteToWordBoundary => "delete-to-word-boundary".to_string(), + Action::ClearLine => "clear-line".to_string(), + Action::ClearToStart => "clear-to-start".to_string(), + Action::ClearToEnd => "clear-to-end".to_string(), + + Action::SelectNext => "select-next".to_string(), + Action::SelectPrevious => "select-previous".to_string(), + Action::ScrollHalfPageUp => "scroll-half-page-up".to_string(), + Action::ScrollHalfPageDown => "scroll-half-page-down".to_string(), + Action::ScrollPageUp => "scroll-page-up".to_string(), + Action::ScrollPageDown => "scroll-page-down".to_string(), + Action::ScrollToTop => "scroll-to-top".to_string(), + Action::ScrollToBottom => "scroll-to-bottom".to_string(), + Action::ScrollToScreenTop => "scroll-to-screen-top".to_string(), + Action::ScrollToScreenMiddle => "scroll-to-screen-middle".to_string(), + Action::ScrollToScreenBottom => "scroll-to-screen-bottom".to_string(), + + Action::Accept => "accept".to_string(), + Action::AcceptNth(n) => format!("accept-{n}"), + Action::ReturnSelection => "return-selection".to_string(), + Action::ReturnSelectionNth(n) => format!("return-selection-{n}"), + Action::Copy => "copy".to_string(), + Action::Delete => "delete".to_string(), + Action::ReturnOriginal => "return-original".to_string(), + Action::ReturnQuery => "return-query".to_string(), + Action::Exit => "exit".to_string(), + Action::Redraw => "redraw".to_string(), + Action::CycleFilterMode => "cycle-filter-mode".to_string(), + Action::CycleSearchMode => "cycle-search-mode".to_string(), + Action::SwitchContext => "switch-context".to_string(), + Action::ClearContext => "clear-context".to_string(), + Action::ToggleTab => "toggle-tab".to_string(), + + Action::VimEnterNormal => "vim-enter-normal".to_string(), + Action::VimEnterInsert => "vim-enter-insert".to_string(), + Action::VimEnterInsertAfter => "vim-enter-insert-after".to_string(), + Action::VimEnterInsertAtStart => "vim-enter-insert-at-start".to_string(), + Action::VimEnterInsertAtEnd => "vim-enter-insert-at-end".to_string(), + Action::VimSearchInsert => "vim-search-insert".to_string(), + Action::VimChangeToEnd => "vim-change-to-end".to_string(), + Action::EnterPrefixMode => "enter-prefix-mode".to_string(), + + Action::InspectPrevious => "inspect-previous".to_string(), + Action::InspectNext => "inspect-next".to_string(), + + Action::Noop => "noop".to_string(), + } + } +} + +impl fmt::Display for Action { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{}", self.as_str()) + } +} + +impl Serialize for Action { + fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> { + serializer.serialize_str(&self.as_str()) + } +} + +impl<'de> Deserialize<'de> for Action { + fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> { + let s = String::deserialize(deserializer)?; + Action::from_str(&s).map_err(serde::de::Error::custom) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parse_basic_actions() { + assert_eq!(Action::from_str("cursor-left").unwrap(), Action::CursorLeft); + assert_eq!(Action::from_str("accept").unwrap(), Action::Accept); + assert_eq!(Action::from_str("exit").unwrap(), Action::Exit); + assert_eq!(Action::from_str("noop").unwrap(), Action::Noop); + assert_eq!( + Action::from_str("vim-enter-normal").unwrap(), + Action::VimEnterNormal + ); + } + + #[test] + fn parse_accept_nth() { + assert_eq!(Action::from_str("accept-1").unwrap(), Action::AcceptNth(1)); + assert_eq!(Action::from_str("accept-9").unwrap(), Action::AcceptNth(9)); + } + + #[test] + fn parse_return_selection() { + assert_eq!( + Action::from_str("return-selection").unwrap(), + Action::ReturnSelection + ); + assert_eq!( + Action::from_str("return-selection-1").unwrap(), + Action::ReturnSelectionNth(1) + ); + assert_eq!( + Action::from_str("return-selection-9").unwrap(), + Action::ReturnSelectionNth(9) + ); + } + + #[test] + fn parse_unknown_action() { + assert!(Action::from_str("unknown-action").is_err()); + assert!(Action::from_str("accept-0").is_err()); + assert!(Action::from_str("accept-10").is_err()); + assert!(Action::from_str("return-selection-0").is_err()); + assert!(Action::from_str("return-selection-10").is_err()); + } + + #[test] + fn round_trip() { + let actions = vec![ + Action::CursorLeft, + Action::Accept, + Action::AcceptNth(5), + Action::ReturnSelection, + Action::ReturnSelectionNth(3), + Action::VimSearchInsert, + Action::ScrollToScreenMiddle, + ]; + for action in actions { + let s = action.as_str(); + let parsed = Action::from_str(&s).unwrap(); + assert_eq!(action, parsed); + } + } + + #[test] + fn serde_round_trip() { + let action = Action::CursorLeft; + let json = serde_json::to_string(&action).unwrap(); + assert_eq!(json, "\"cursor-left\""); + let parsed: Action = serde_json::from_str(&json).unwrap(); + assert_eq!(parsed, Action::CursorLeft); + + let action = Action::AcceptNth(3); + let json = serde_json::to_string(&action).unwrap(); + assert_eq!(json, "\"accept-3\""); + let parsed: Action = serde_json::from_str(&json).unwrap(); + assert_eq!(parsed, Action::AcceptNth(3)); + } +} diff --git a/crates/atuin/src/command/client/search/keybindings/conditions.rs b/crates/atuin/src/command/client/search/keybindings/conditions.rs new file mode 100644 index 00000000..d460d7d4 --- /dev/null +++ b/crates/atuin/src/command/client/search/keybindings/conditions.rs @@ -0,0 +1,801 @@ +use std::fmt; + +use serde::{Deserialize, Deserializer, Serialize, Serializer}; + +/// Atomic (leaf) conditions that can be evaluated against state. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ConditionAtom { + CursorAtStart, + CursorAtEnd, + InputEmpty, + OriginalInputEmpty, + ListAtEnd, + ListAtStart, + NoResults, + HasResults, + HasContext, +} + +/// Boolean expression tree over condition atoms. +/// +/// Supports negation, conjunction, and disjunction with standard precedence: +/// `!` binds tightest, then `&&`, then `||`. +/// +/// Examples of valid expression strings: +/// - `"cursor-at-start"` (bare atom) +/// - `"!no-results"` (negation) +/// - `"cursor-at-start && input-empty"` (conjunction) +/// - `"list-at-start || no-results"` (disjunction) +/// - `"(cursor-at-start && !input-empty) || no-results"` (grouping) +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ConditionExpr { + Atom(ConditionAtom), + Not(Box<ConditionExpr>), + And(Box<ConditionExpr>, Box<ConditionExpr>), + Or(Box<ConditionExpr>, Box<ConditionExpr>), +} + +/// Context needed to evaluate conditions. This is a pure snapshot of state — +/// no references to mutable data. +pub struct EvalContext { + /// Current cursor position (unicode width units). + pub cursor_position: usize, + /// Width of the input string in unicode width units. + pub input_width: usize, + /// Byte length of the input string. + pub input_byte_len: usize, + /// Currently selected index in the results list. + pub selected_index: usize, + /// Total number of results. + pub results_len: usize, + /// Whether the original input (query passed to the TUI) was empty. + pub original_input_empty: bool, + /// Whether we use a search context of a command from the history. + pub has_context: bool, +} + +// --------------------------------------------------------------------------- +// ConditionAtom +// --------------------------------------------------------------------------- + +impl ConditionAtom { + /// Evaluate this atom against the given context. + pub fn evaluate(&self, ctx: &EvalContext) -> bool { + match self { + ConditionAtom::CursorAtStart => ctx.cursor_position == 0, + ConditionAtom::CursorAtEnd => ctx.cursor_position == ctx.input_width, + ConditionAtom::InputEmpty => ctx.input_byte_len == 0, + ConditionAtom::OriginalInputEmpty => ctx.original_input_empty, + ConditionAtom::ListAtEnd => { + ctx.results_len == 0 || ctx.selected_index >= ctx.results_len.saturating_sub(1) + } + ConditionAtom::ListAtStart => ctx.results_len == 0 || ctx.selected_index == 0, + ConditionAtom::NoResults => ctx.results_len == 0, + ConditionAtom::HasResults => ctx.results_len > 0, + ConditionAtom::HasContext => ctx.has_context, + } + } + + /// Parse from a kebab-case string. + pub fn from_str(s: &str) -> Result<Self, String> { + match s { + "cursor-at-start" => Ok(ConditionAtom::CursorAtStart), + "cursor-at-end" => Ok(ConditionAtom::CursorAtEnd), + "input-empty" => Ok(ConditionAtom::InputEmpty), + "original-input-empty" => Ok(ConditionAtom::OriginalInputEmpty), + "list-at-end" => Ok(ConditionAtom::ListAtEnd), + "list-at-start" => Ok(ConditionAtom::ListAtStart), + "no-results" => Ok(ConditionAtom::NoResults), + "has-results" => Ok(ConditionAtom::HasResults), + "has-context" => Ok(ConditionAtom::HasContext), + _ => Err(format!("unknown condition: {s}")), + } + } + + /// Convert to a kebab-case string. + pub fn as_str(&self) -> &'static str { + match self { + ConditionAtom::CursorAtStart => "cursor-at-start", + ConditionAtom::CursorAtEnd => "cursor-at-end", + ConditionAtom::InputEmpty => "input-empty", + ConditionAtom::OriginalInputEmpty => "original-input-empty", + ConditionAtom::ListAtEnd => "list-at-end", + ConditionAtom::ListAtStart => "list-at-start", + ConditionAtom::NoResults => "no-results", + ConditionAtom::HasResults => "has-results", + ConditionAtom::HasContext => "has-context", + } + } +} + +impl fmt::Display for ConditionAtom { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{}", self.as_str()) + } +} + +// --------------------------------------------------------------------------- +// ConditionExpr — evaluation +// --------------------------------------------------------------------------- + +impl ConditionExpr { + /// Evaluate this expression against the given context. + pub fn evaluate(&self, ctx: &EvalContext) -> bool { + match self { + ConditionExpr::Atom(atom) => atom.evaluate(ctx), + ConditionExpr::Not(inner) => !inner.evaluate(ctx), + ConditionExpr::And(lhs, rhs) => lhs.evaluate(ctx) && rhs.evaluate(ctx), + ConditionExpr::Or(lhs, rhs) => lhs.evaluate(ctx) || rhs.evaluate(ctx), + } + } +} + +// --------------------------------------------------------------------------- +// ConditionExpr — ergonomic builders +// --------------------------------------------------------------------------- + +impl From<ConditionAtom> for ConditionExpr { + fn from(atom: ConditionAtom) -> Self { + ConditionExpr::Atom(atom) + } +} + +#[allow(dead_code)] +impl ConditionExpr { + /// Negate this expression: `!self`. + pub fn not(self) -> Self { + ConditionExpr::Not(Box::new(self)) + } + + /// Conjoin with another expression: `self && other`. + pub fn and(self, other: ConditionExpr) -> Self { + ConditionExpr::And(Box::new(self), Box::new(other)) + } + + /// Disjoin with another expression: `self || other`. + pub fn or(self, other: ConditionExpr) -> Self { + ConditionExpr::Or(Box::new(self), Box::new(other)) + } +} + +// --------------------------------------------------------------------------- +// ConditionExpr — parser +// --------------------------------------------------------------------------- + +/// Recursive descent parser for boolean condition expressions. +/// +/// Grammar (standard boolean precedence): +/// ```text +/// expr = or_expr +/// or_expr = and_expr ("||" and_expr)* +/// and_expr = unary ("&&" unary)* +/// unary = "!" unary | primary +/// primary = atom | "(" expr ")" +/// atom = [a-z][a-z0-9-]* +/// ``` +struct ExprParser<'a> { + input: &'a str, + pos: usize, +} + +impl<'a> ExprParser<'a> { + fn new(input: &'a str) -> Self { + Self { input, pos: 0 } + } + + fn skip_whitespace(&mut self) { + while self.pos < self.input.len() && self.input.as_bytes()[self.pos].is_ascii_whitespace() { + self.pos += 1; + } + } + + fn starts_with(&mut self, s: &str) -> bool { + self.skip_whitespace(); + self.input[self.pos..].starts_with(s) + } + + fn consume(&mut self, s: &str) -> bool { + self.skip_whitespace(); + if self.input[self.pos..].starts_with(s) { + self.pos += s.len(); + true + } else { + false + } + } + + /// Parse a full expression, expecting to consume all input. + fn parse(mut self) -> Result<ConditionExpr, String> { + let expr = self.parse_or()?; + self.skip_whitespace(); + if self.pos < self.input.len() { + return Err(format!( + "unexpected input at position {}: {:?}", + self.pos, + &self.input[self.pos..] + )); + } + Ok(expr) + } + + /// `or_expr` = `and_expr` ("||" `and_expr`)* + fn parse_or(&mut self) -> Result<ConditionExpr, String> { + let mut left = self.parse_and()?; + while self.starts_with("||") { + self.consume("||"); + let right = self.parse_and()?; + left = ConditionExpr::Or(Box::new(left), Box::new(right)); + } + Ok(left) + } + + /// `and_expr` = unary ("&&" unary)* + fn parse_and(&mut self) -> Result<ConditionExpr, String> { + let mut left = self.parse_unary()?; + while self.starts_with("&&") { + self.consume("&&"); + let right = self.parse_unary()?; + left = ConditionExpr::And(Box::new(left), Box::new(right)); + } + Ok(left) + } + + /// unary = "!" unary | primary + fn parse_unary(&mut self) -> Result<ConditionExpr, String> { + if self.consume("!") { + let inner = self.parse_unary()?; + Ok(ConditionExpr::Not(Box::new(inner))) + } else { + self.parse_primary() + } + } + + /// primary = "(" expr ")" | atom + fn parse_primary(&mut self) -> Result<ConditionExpr, String> { + if self.consume("(") { + let expr = self.parse_or()?; + if !self.consume(")") { + return Err(format!("expected ')' at position {}", self.pos)); + } + Ok(expr) + } else { + self.parse_atom() + } + } + + /// atom = [a-z][a-z0-9-]* + fn parse_atom(&mut self) -> Result<ConditionExpr, String> { + self.skip_whitespace(); + let start = self.pos; + while self.pos < self.input.len() { + let b = self.input.as_bytes()[self.pos]; + if b.is_ascii_lowercase() || b.is_ascii_digit() || b == b'-' { + self.pos += 1; + } else { + break; + } + } + if self.pos == start { + return Err(format!("expected condition name at position {}", self.pos)); + } + let name = &self.input[start..self.pos]; + let atom = ConditionAtom::from_str(name)?; + Ok(ConditionExpr::Atom(atom)) + } +} + +impl ConditionExpr { + /// Parse a condition expression from a string. + pub fn parse(s: &str) -> Result<Self, String> { + let parser = ExprParser::new(s); + parser.parse() + } +} + +// --------------------------------------------------------------------------- +// ConditionExpr — Display +// --------------------------------------------------------------------------- + +/// Precedence levels for minimal-parentheses display. +#[derive(PartialEq, Eq, PartialOrd, Ord, Clone, Copy)] +enum Prec { + Or = 0, + And = 1, + Not = 2, + Atom = 3, +} + +impl ConditionExpr { + fn prec(&self) -> Prec { + match self { + ConditionExpr::Or(..) => Prec::Or, + ConditionExpr::And(..) => Prec::And, + ConditionExpr::Not(..) => Prec::Not, + ConditionExpr::Atom(..) => Prec::Atom, + } + } + + fn fmt_with_prec(&self, f: &mut fmt::Formatter<'_>, parent_prec: Prec) -> fmt::Result { + let needs_parens = self.prec() < parent_prec; + if needs_parens { + write!(f, "(")?; + } + match self { + ConditionExpr::Atom(atom) => write!(f, "{atom}")?, + ConditionExpr::Not(inner) => { + write!(f, "!")?; + inner.fmt_with_prec(f, Prec::Not)?; + } + ConditionExpr::And(lhs, rhs) => { + lhs.fmt_with_prec(f, Prec::And)?; + write!(f, " && ")?; + rhs.fmt_with_prec(f, Prec::And)?; + } + ConditionExpr::Or(lhs, rhs) => { + lhs.fmt_with_prec(f, Prec::Or)?; + write!(f, " || ")?; + rhs.fmt_with_prec(f, Prec::Or)?; + } + } + if needs_parens { + write!(f, ")")?; + } + Ok(()) + } +} + +impl fmt::Display for ConditionExpr { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + self.fmt_with_prec(f, Prec::Or) + } +} + +// --------------------------------------------------------------------------- +// Serde +// --------------------------------------------------------------------------- + +impl Serialize for ConditionExpr { + fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> { + serializer.serialize_str(&self.to_string()) + } +} + +impl<'de> Deserialize<'de> for ConditionExpr { + fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> { + let s = String::deserialize(deserializer)?; + ConditionExpr::parse(&s).map_err(serde::de::Error::custom) + } +} + +// --------------------------------------------------------------------------- +// Tests +// --------------------------------------------------------------------------- + +#[cfg(test)] +mod tests { + use super::*; + + fn ctx( + cursor: usize, + width: usize, + byte_len: usize, + selected: usize, + len: usize, + ) -> EvalContext { + ctx_with_original(cursor, width, byte_len, selected, len, false) + } + + fn ctx_with_original( + cursor: usize, + width: usize, + byte_len: usize, + selected: usize, + len: usize, + original_input_empty: bool, + ) -> EvalContext { + EvalContext { + cursor_position: cursor, + input_width: width, + input_byte_len: byte_len, + selected_index: selected, + results_len: len, + original_input_empty, + has_context: false, + } + } + + // -- Atom evaluation (carried over from Phase 0) -- + + #[test] + fn atom_cursor_at_start() { + assert!(ConditionAtom::CursorAtStart.evaluate(&ctx(0, 5, 5, 0, 10))); + assert!(!ConditionAtom::CursorAtStart.evaluate(&ctx(3, 5, 5, 0, 10))); + } + + #[test] + fn atom_cursor_at_end() { + assert!(ConditionAtom::CursorAtEnd.evaluate(&ctx(5, 5, 5, 0, 10))); + assert!(!ConditionAtom::CursorAtEnd.evaluate(&ctx(3, 5, 5, 0, 10))); + assert!(ConditionAtom::CursorAtEnd.evaluate(&ctx(0, 0, 0, 0, 10))); + } + + #[test] + fn atom_input_empty() { + assert!(ConditionAtom::InputEmpty.evaluate(&ctx(0, 0, 0, 0, 10))); + assert!(!ConditionAtom::InputEmpty.evaluate(&ctx(0, 5, 5, 0, 10))); + } + + #[test] + fn atom_original_input_empty() { + // original_input_empty = true + assert!( + ConditionAtom::OriginalInputEmpty.evaluate(&ctx_with_original(0, 0, 0, 0, 10, true)) + ); + // original_input_empty = false + assert!( + !ConditionAtom::OriginalInputEmpty.evaluate(&ctx_with_original(0, 0, 0, 0, 10, false)) + ); + // original_input_empty is independent of current input state + assert!( + ConditionAtom::OriginalInputEmpty.evaluate(&ctx_with_original(0, 5, 5, 0, 10, true)) + ); + } + + #[test] + fn atom_list_at_end() { + assert!(ConditionAtom::ListAtEnd.evaluate(&ctx(0, 0, 0, 99, 100))); + assert!(!ConditionAtom::ListAtEnd.evaluate(&ctx(0, 0, 0, 50, 100))); + assert!(ConditionAtom::ListAtEnd.evaluate(&ctx(0, 0, 0, 0, 0))); + } + + #[test] + fn atom_list_at_start() { + assert!(ConditionAtom::ListAtStart.evaluate(&ctx(0, 0, 0, 0, 100))); + assert!(!ConditionAtom::ListAtStart.evaluate(&ctx(0, 0, 0, 50, 100))); + assert!(ConditionAtom::ListAtStart.evaluate(&ctx(0, 0, 0, 0, 0))); + } + + #[test] + fn atom_no_results_and_has_results() { + assert!(ConditionAtom::NoResults.evaluate(&ctx(0, 0, 0, 0, 0))); + assert!(!ConditionAtom::NoResults.evaluate(&ctx(0, 0, 0, 0, 5))); + assert!(ConditionAtom::HasResults.evaluate(&ctx(0, 0, 0, 0, 5))); + assert!(!ConditionAtom::HasResults.evaluate(&ctx(0, 0, 0, 0, 0))); + } + + #[test] + fn atom_has_context() { + let mut context = ctx(0, 0, 0, 0, 0); + assert!(!ConditionAtom::HasContext.evaluate(&context)); + context.has_context = true; + assert!(ConditionAtom::HasContext.evaluate(&context)); + } + + #[test] + fn atom_parse_round_trip() { + let conditions = [ + "cursor-at-start", + "cursor-at-end", + "input-empty", + "original-input-empty", + "list-at-end", + "list-at-start", + "no-results", + "has-results", + ]; + for s in conditions { + let c = ConditionAtom::from_str(s).unwrap(); + assert_eq!(c.as_str(), s); + } + } + + #[test] + fn atom_parse_unknown() { + assert!(ConditionAtom::from_str("unknown-condition").is_err()); + } + + // -- Parser tests -- + + #[test] + fn parse_bare_atom() { + let expr = ConditionExpr::parse("cursor-at-start").unwrap(); + assert_eq!(expr, ConditionExpr::Atom(ConditionAtom::CursorAtStart)); + } + + #[test] + fn parse_negation() { + let expr = ConditionExpr::parse("!no-results").unwrap(); + assert_eq!( + expr, + ConditionExpr::Not(Box::new(ConditionExpr::Atom(ConditionAtom::NoResults))) + ); + } + + #[test] + fn parse_double_negation() { + let expr = ConditionExpr::parse("!!no-results").unwrap(); + assert_eq!( + expr, + ConditionExpr::Not(Box::new(ConditionExpr::Not(Box::new(ConditionExpr::Atom( + ConditionAtom::NoResults + ))))) + ); + } + + #[test] + fn parse_and() { + let expr = ConditionExpr::parse("cursor-at-start && input-empty").unwrap(); + assert_eq!( + expr, + ConditionExpr::And( + Box::new(ConditionExpr::Atom(ConditionAtom::CursorAtStart)), + Box::new(ConditionExpr::Atom(ConditionAtom::InputEmpty)), + ) + ); + } + + #[test] + fn parse_or() { + let expr = ConditionExpr::parse("list-at-start || no-results").unwrap(); + assert_eq!( + expr, + ConditionExpr::Or( + Box::new(ConditionExpr::Atom(ConditionAtom::ListAtStart)), + Box::new(ConditionExpr::Atom(ConditionAtom::NoResults)), + ) + ); + } + + #[test] + fn parse_precedence_and_binds_tighter_than_or() { + // "a || b && c" should parse as "a || (b && c)" + let expr = ConditionExpr::parse("cursor-at-start || input-empty && no-results").unwrap(); + assert_eq!( + expr, + ConditionExpr::Or( + Box::new(ConditionExpr::Atom(ConditionAtom::CursorAtStart)), + Box::new(ConditionExpr::And( + Box::new(ConditionExpr::Atom(ConditionAtom::InputEmpty)), + Box::new(ConditionExpr::Atom(ConditionAtom::NoResults)), + )), + ) + ); + } + + #[test] + fn parse_parens_override_precedence() { + // "(a || b) && c" + let expr = ConditionExpr::parse("(cursor-at-start || input-empty) && no-results").unwrap(); + assert_eq!( + expr, + ConditionExpr::And( + Box::new(ConditionExpr::Or( + Box::new(ConditionExpr::Atom(ConditionAtom::CursorAtStart)), + Box::new(ConditionExpr::Atom(ConditionAtom::InputEmpty)), + )), + Box::new(ConditionExpr::Atom(ConditionAtom::NoResults)), + ) + ); + } + + #[test] + fn parse_complex_nested() { + // "(a && !b) || c" + let expr = ConditionExpr::parse("(cursor-at-start && !input-empty) || no-results").unwrap(); + assert_eq!( + expr, + ConditionExpr::Or( + Box::new(ConditionExpr::And( + Box::new(ConditionExpr::Atom(ConditionAtom::CursorAtStart)), + Box::new(ConditionExpr::Not(Box::new(ConditionExpr::Atom( + ConditionAtom::InputEmpty + )))), + )), + Box::new(ConditionExpr::Atom(ConditionAtom::NoResults)), + ) + ); + } + + #[test] + fn parse_whitespace_tolerance() { + let a = ConditionExpr::parse("cursor-at-start||input-empty").unwrap(); + let b = ConditionExpr::parse("cursor-at-start || input-empty").unwrap(); + let c = ConditionExpr::parse(" cursor-at-start || input-empty ").unwrap(); + assert_eq!(a, b); + assert_eq!(b, c); + } + + #[test] + fn parse_error_unknown_atom() { + assert!(ConditionExpr::parse("unknown-thing").is_err()); + } + + #[test] + fn parse_error_trailing_input() { + assert!(ConditionExpr::parse("cursor-at-start blah").is_err()); + } + + #[test] + fn parse_error_unmatched_paren() { + assert!(ConditionExpr::parse("(cursor-at-start").is_err()); + } + + #[test] + fn parse_error_empty() { + assert!(ConditionExpr::parse("").is_err()); + } + + // -- Expression evaluation -- + + #[test] + fn eval_not() { + let expr = ConditionExpr::parse("!no-results").unwrap(); + // Has results → !no-results is true + assert!(expr.evaluate(&ctx(0, 0, 0, 0, 5))); + // No results → !no-results is false + assert!(!expr.evaluate(&ctx(0, 0, 0, 0, 0))); + } + + #[test] + fn eval_and() { + let expr = ConditionExpr::parse("cursor-at-start && input-empty").unwrap(); + // Both true + assert!(expr.evaluate(&ctx(0, 0, 0, 0, 10))); + // First true, second false (non-empty input) + assert!(!expr.evaluate(&ctx(0, 5, 5, 0, 10))); + // First false (cursor not at start) + assert!(!expr.evaluate(&ctx(3, 5, 5, 0, 10))); + } + + #[test] + fn eval_or() { + let expr = ConditionExpr::parse("list-at-start || no-results").unwrap(); + // list at bottom (selected=0) + assert!(expr.evaluate(&ctx(0, 0, 0, 0, 10))); + // no results + assert!(expr.evaluate(&ctx(0, 0, 0, 0, 0))); + // neither + assert!(!expr.evaluate(&ctx(0, 0, 0, 5, 10))); + } + + #[test] + fn eval_complex_nested() { + // (cursor-at-start && !input-empty) || no-results + let expr = ConditionExpr::parse("(cursor-at-start && !input-empty) || no-results").unwrap(); + + // cursor at start, input not empty → true (left branch) + assert!(expr.evaluate(&ctx(0, 5, 5, 0, 10))); + // no results → true (right branch) + assert!(expr.evaluate(&ctx(3, 5, 5, 0, 0))); + // cursor not at start, has results → false + assert!(!expr.evaluate(&ctx(3, 5, 5, 0, 10))); + // cursor at start, input empty → false (left: && fails; right: has results) + assert!(!expr.evaluate(&ctx(0, 0, 0, 0, 10))); + } + + // -- Display -- + + #[test] + fn display_atom() { + let expr = ConditionExpr::Atom(ConditionAtom::CursorAtStart); + assert_eq!(expr.to_string(), "cursor-at-start"); + } + + #[test] + fn display_not() { + let expr = ConditionExpr::Atom(ConditionAtom::NoResults).not(); + assert_eq!(expr.to_string(), "!no-results"); + } + + #[test] + fn display_and() { + let expr = ConditionExpr::Atom(ConditionAtom::CursorAtStart) + .and(ConditionExpr::Atom(ConditionAtom::InputEmpty)); + assert_eq!(expr.to_string(), "cursor-at-start && input-empty"); + } + + #[test] + fn display_or() { + let expr = ConditionExpr::Atom(ConditionAtom::ListAtStart) + .or(ConditionExpr::Atom(ConditionAtom::NoResults)); + assert_eq!(expr.to_string(), "list-at-start || no-results"); + } + + #[test] + fn display_parens_when_needed() { + // (a || b) && c — the Or inside And needs parens + let expr = ConditionExpr::Atom(ConditionAtom::CursorAtStart) + .or(ConditionExpr::Atom(ConditionAtom::InputEmpty)) + .and(ConditionExpr::Atom(ConditionAtom::NoResults)); + assert_eq!( + expr.to_string(), + "(cursor-at-start || input-empty) && no-results" + ); + } + + #[test] + fn display_no_parens_when_not_needed() { + // a || b && c — no parens needed (and binds tighter) + let inner_and = ConditionExpr::Atom(ConditionAtom::InputEmpty) + .and(ConditionExpr::Atom(ConditionAtom::NoResults)); + let expr = ConditionExpr::Atom(ConditionAtom::CursorAtStart).or(inner_and); + assert_eq!( + expr.to_string(), + "cursor-at-start || input-empty && no-results" + ); + } + + // -- Display round-trip -- + + #[test] + fn display_round_trip() { + let cases = [ + "cursor-at-start", + "!no-results", + "cursor-at-start && input-empty", + "list-at-start || no-results", + "(cursor-at-start || input-empty) && no-results", + "(cursor-at-start && !input-empty) || no-results", + ]; + for s in cases { + let expr = ConditionExpr::parse(s).unwrap(); + let displayed = expr.to_string(); + let reparsed = ConditionExpr::parse(&displayed).unwrap(); + assert_eq!(expr, reparsed, "round-trip failed for: {s}"); + } + } + + // -- Serde -- + + #[test] + fn serde_simple_atom() { + let expr = ConditionExpr::Atom(ConditionAtom::CursorAtStart); + let json = serde_json::to_string(&expr).unwrap(); + assert_eq!(json, "\"cursor-at-start\""); + let parsed: ConditionExpr = serde_json::from_str(&json).unwrap(); + assert_eq!(parsed, expr); + } + + #[test] + fn serde_compound_expression() { + let json = "\"cursor-at-start && !input-empty\""; + let parsed: ConditionExpr = serde_json::from_str(json).unwrap(); + let expected = ConditionExpr::And( + Box::new(ConditionExpr::Atom(ConditionAtom::CursorAtStart)), + Box::new(ConditionExpr::Not(Box::new(ConditionExpr::Atom( + ConditionAtom::InputEmpty, + )))), + ); + assert_eq!(parsed, expected); + } + + #[test] + fn serde_round_trip() { + let expr = ConditionExpr::parse("(cursor-at-start && !input-empty) || no-results").unwrap(); + let json = serde_json::to_string(&expr).unwrap(); + let parsed: ConditionExpr = serde_json::from_str(&json).unwrap(); + assert_eq!(expr, parsed); + } + + // -- From<ConditionAtom> -- + + #[test] + fn from_atom_into_expr() { + let expr: ConditionExpr = ConditionAtom::CursorAtStart.into(); + assert_eq!(expr, ConditionExpr::Atom(ConditionAtom::CursorAtStart)); + } + + // -- Builder helpers -- + + #[test] + fn builder_chain() { + let expr = ConditionExpr::from(ConditionAtom::CursorAtStart) + .and(ConditionExpr::from(ConditionAtom::InputEmpty).not()) + .or(ConditionExpr::from(ConditionAtom::NoResults)); + // And binds tighter than Or, so no parens needed around the And + assert_eq!( + expr.to_string(), + "cursor-at-start && !input-empty || no-results" + ); + } +} diff --git a/crates/atuin/src/command/client/search/keybindings/defaults.rs b/crates/atuin/src/command/client/search/keybindings/defaults.rs new file mode 100644 index 00000000..f19bf377 --- /dev/null +++ b/crates/atuin/src/command/client/search/keybindings/defaults.rs @@ -0,0 +1,1285 @@ +use std::collections::HashMap; + +use atuin_client::settings::{KeyBindingConfig, Settings}; +use tracing::warn; + +use super::actions::Action; +use super::conditions::{ConditionAtom, ConditionExpr}; +use super::key::KeyInput; +use super::keymap::{KeyBinding, KeyRule, Keymap}; + +/// Helper to bind a scroll key with optional exit behavior. +/// +/// When `scroll_exits` is true AND the key scrolls toward index 0 (the newest +/// entry), we add a conditional rule: at `ListAtStart` → `Exit`, otherwise → +/// the scroll action. +/// +/// Whether a key scrolls toward index 0 depends on the `invert` setting: +/// - Non-inverted: "down" / "j" move toward index 0, "up" / "k" move away +/// - Inverted: "up" / "k" move toward index 0, "down" / "j" move away +/// +/// If `toward_index_zero` is false, or `scroll_exits` is false, we just bind +/// the key to the plain scroll action (no exit). +fn bind_scroll_key( + km: &mut Keymap, + key_str: &str, + action: Action, + toward_index_zero: bool, + scroll_exits: bool, +) { + let k = key(key_str); + if scroll_exits && toward_index_zero { + km.bind_conditional( + k, + vec![ + KeyRule::when(ConditionAtom::ListAtStart, Action::Exit), + KeyRule::always(action), + ], + ); + } else { + km.bind(k, action); + } +} + +/// Helper to parse a key string, panicking on invalid keys (these are all +/// compile-time-known strings). +fn key(s: &str) -> KeyInput { + KeyInput::parse(s).unwrap_or_else(|e| panic!("invalid default key {s:?}: {e}")) +} + +/// All five keymaps bundled together. +#[derive(Debug, Clone)] +pub struct KeymapSet { + pub emacs: Keymap, + pub vim_normal: Keymap, + pub vim_insert: Keymap, + pub inspector: Keymap, + pub prefix: Keymap, +} + +// --------------------------------------------------------------------------- +// Common bindings shared across search-tab keymaps +// --------------------------------------------------------------------------- + +/// Add the bindings that are common to all search-tab keymaps: +/// ctrl-c, ctrl-g, ctrl-o, and tab. +/// +/// Note: `esc`/`ctrl-[` are NOT included here because their behavior differs +/// between emacs (exit), vim-normal (exit), and vim-insert (enter normal mode). +fn add_common_bindings(km: &mut Keymap) { + km.bind(key("ctrl-c"), Action::ReturnOriginal); + km.bind(key("ctrl-g"), Action::ReturnOriginal); + km.bind(key("ctrl-o"), Action::ToggleTab); + + // Tab: always returns selection without executing (unlike Enter which respects enter_accept) + km.bind(key("tab"), Action::ReturnSelection); +} + +/// Returns `Accept` or `ReturnSelection` based on the `enter_accept` setting. +fn accept_action(settings: &Settings) -> Action { + if settings.enter_accept { + Action::Accept + } else { + Action::ReturnSelection + } +} + +// --------------------------------------------------------------------------- +// Emacs keymap (also base for vim-insert) +// --------------------------------------------------------------------------- + +/// Build the default emacs keymap. This encodes the behavior from +/// `handle_key_input` common section + `handle_search_input` shared section. +/// +/// The `settings` parameter is used for: +/// - `keys.prefix` — which ctrl-key enters prefix mode +/// - `keys.scroll_exits`, `invert` — scroll-at-boundary exit behavior +/// - `keys.accept_past_line_end` — right arrow at end of line accepts +/// - `keys.exit_past_line_start` — left arrow at start of line exits +/// - `keys.accept_past_line_start` — left arrow at start accepts (overrides exit) +/// - `keys.accept_with_backspace` — backspace at start of line accepts +/// - `ctrl_n_shortcuts` — whether alt or ctrl is used for numeric shortcuts +// Keymap builder that enumerates every default binding; not worth splitting. +#[allow(clippy::too_many_lines)] +pub fn default_emacs_keymap(settings: &Settings) -> Keymap { + let mut km = Keymap::new(); + add_common_bindings(&mut km); + + let accept = accept_action(settings); + + // esc / ctrl-[ → exit + km.bind(key("esc"), Action::Exit); + km.bind(key("ctrl-["), Action::Exit); + + // Prefix key: ctrl-<prefix_char> → enter prefix mode + let prefix_char = settings.keys.prefix.chars().next().unwrap_or('a'); + km.bind(key(&format!("ctrl-{prefix_char}")), Action::EnterPrefixMode); + + // --- Accept / navigation edge behaviors (from [keys] settings) --- + + // right: behavior at end of line + if settings.keys.accept_past_line_end { + km.bind_conditional( + key("right"), + vec![ + KeyRule::when(ConditionAtom::CursorAtEnd, Action::ReturnSelection), + KeyRule::always(Action::CursorRight), + ], + ); + } else { + km.bind(key("right"), Action::CursorRight); + } + + // left: behavior at start of line + // accept_past_line_start takes precedence over exit_past_line_start + if settings.keys.accept_past_line_start { + km.bind_conditional( + key("left"), + vec![ + KeyRule::when(ConditionAtom::CursorAtStart, Action::ReturnSelection), + KeyRule::always(Action::CursorLeft), + ], + ); + } else if settings.keys.exit_past_line_start { + km.bind_conditional( + key("left"), + vec![ + KeyRule::when(ConditionAtom::CursorAtStart, Action::Exit), + KeyRule::always(Action::CursorLeft), + ], + ); + } else { + km.bind(key("left"), Action::CursorLeft); + } + + // down/up: scroll with optional exit at boundary. + // Non-inverted: down moves toward index 0 (can exit); up moves away (no exit). + // Inverted: up moves toward index 0 (can exit); down moves away (no exit). + let scroll_exits = settings.keys.scroll_exits; + let invert = settings.invert; + bind_scroll_key(&mut km, "down", Action::SelectNext, !invert, scroll_exits); + bind_scroll_key(&mut km, "up", Action::SelectPrevious, invert, scroll_exits); + + // backspace: behavior at start of line + if settings.keys.accept_with_backspace { + km.bind_conditional( + key("backspace"), + vec![ + KeyRule::when(ConditionAtom::CursorAtStart, Action::ReturnSelection), + KeyRule::always(Action::DeleteCharBefore), + ], + ); + } else { + km.bind(key("backspace"), Action::DeleteCharBefore); + } + + // --- Accept --- + km.bind(key("enter"), accept.clone()); + km.bind(key("ctrl-m"), accept); + + // --- Copy --- + km.bind(key("ctrl-y"), Action::Copy); + + // --- Numeric shortcuts (alt-1..9 by default, ctrl-1..9 if ctrl_n_shortcuts) --- + // These return the selection without executing, regardless of enter_accept. + let num_mod = if settings.ctrl_n_shortcuts { + "ctrl" + } else { + "alt" + }; + for n in 1..=9u8 { + km.bind( + key(&format!("{num_mod}-{n}")), + Action::ReturnSelectionNth(n), + ); + } + + // --- Cursor movement --- + km.bind(key("ctrl-left"), Action::CursorWordLeft); + km.bind(key("alt-b"), Action::CursorWordLeft); + km.bind(key("ctrl-b"), Action::CursorLeft); + km.bind(key("ctrl-right"), Action::CursorWordRight); + km.bind(key("alt-f"), Action::CursorWordRight); + km.bind(key("ctrl-f"), Action::CursorRight); + km.bind(key("home"), Action::CursorStart); + // ctrl-a → CursorStart only if prefix char is NOT 'a' + // (otherwise ctrl-a is already bound to EnterPrefixMode above) + if prefix_char != 'a' { + km.bind(key("ctrl-a"), Action::CursorStart); + } + km.bind(key("ctrl-e"), Action::CursorEnd); + km.bind(key("end"), Action::CursorEnd); + + // --- Editing --- + km.bind(key("ctrl-backspace"), Action::DeleteWordBefore); + km.bind(key("ctrl-h"), Action::DeleteCharBefore); + km.bind(key("ctrl-?"), Action::DeleteCharBefore); + km.bind(key("ctrl-delete"), Action::DeleteWordAfter); + km.bind(key("delete"), Action::DeleteCharAfter); + // ctrl-d: if input empty → return original, otherwise delete char + km.bind_conditional( + key("ctrl-d"), + vec![ + KeyRule::when(ConditionAtom::InputEmpty, Action::ReturnOriginal), + KeyRule::always(Action::DeleteCharAfter), + ], + ); + km.bind(key("ctrl-w"), Action::DeleteToWordBoundary); + km.bind(key("ctrl-u"), Action::ClearLine); + + // --- Search mode --- + km.bind(key("ctrl-r"), Action::CycleFilterMode); + km.bind(key("ctrl-s"), Action::CycleSearchMode); + + // --- Scroll (no exit) --- + km.bind(key("ctrl-n"), Action::SelectNext); + km.bind(key("ctrl-j"), Action::SelectNext); + km.bind(key("ctrl-p"), Action::SelectPrevious); + km.bind(key("ctrl-k"), Action::SelectPrevious); + + // --- Redraw --- + km.bind(key("ctrl-l"), Action::Redraw); + + // --- Page scroll --- + km.bind(key("pagedown"), Action::ScrollPageDown); + km.bind(key("pageup"), Action::ScrollPageUp); + + km +} + +// --------------------------------------------------------------------------- +// Vim Normal keymap +// --------------------------------------------------------------------------- + +/// Build the default vim-normal keymap. +pub fn default_vim_normal_keymap(settings: &Settings) -> Keymap { + let mut km = Keymap::new(); + add_common_bindings(&mut km); + + // esc / ctrl-[ → exit (vim-normal exits, unlike vim-insert) + km.bind(key("esc"), Action::Exit); + km.bind(key("ctrl-["), Action::Exit); + + // Prefix key + let prefix_char = settings.keys.prefix.chars().next().unwrap_or('a'); + km.bind(key(&format!("ctrl-{prefix_char}")), Action::EnterPrefixMode); + + // --- Vim navigation --- + // j/k: scroll with optional exit at boundary. + let scroll_exits = settings.keys.scroll_exits; + let invert = settings.invert; + bind_scroll_key(&mut km, "j", Action::SelectNext, !invert, scroll_exits); + bind_scroll_key(&mut km, "k", Action::SelectPrevious, invert, scroll_exits); + km.bind(key("h"), Action::CursorLeft); + km.bind(key("l"), Action::CursorRight); + + // --- Vim cursor movement --- + km.bind(key("0"), Action::CursorStart); + km.bind(key("$"), Action::CursorEnd); + km.bind(key("w"), Action::CursorWordRight); + km.bind(key("b"), Action::CursorWordLeft); + km.bind(key("e"), Action::CursorWordEnd); + + // --- Vim editing --- + km.bind(key("x"), Action::DeleteCharAfter); + km.bind(key("d d"), Action::ClearLine); + km.bind(key("D"), Action::ClearToEnd); + km.bind(key("C"), Action::VimChangeToEnd); + + // --- Mode switching --- + km.bind(key("?"), Action::VimSearchInsert); + km.bind(key("/"), Action::VimSearchInsert); + km.bind(key("a"), Action::VimEnterInsertAfter); + km.bind(key("A"), Action::VimEnterInsertAtEnd); + km.bind(key("i"), Action::VimEnterInsert); + km.bind(key("I"), Action::VimEnterInsertAtStart); + + // --- Numeric shortcuts (return selection without executing) --- + for n in 1..=9u8 { + km.bind(key(&n.to_string()), Action::ReturnSelectionNth(n)); + } + + // --- Half/full page scroll --- + km.bind(key("ctrl-u"), Action::ScrollHalfPageUp); + km.bind(key("ctrl-d"), Action::ScrollHalfPageDown); + km.bind(key("ctrl-b"), Action::ScrollPageUp); + km.bind(key("ctrl-f"), Action::ScrollPageDown); + + // --- Jump --- + km.bind(key("G"), Action::ScrollToBottom); + km.bind(key("g g"), Action::ScrollToTop); + km.bind(key("H"), Action::ScrollToScreenTop); + km.bind(key("M"), Action::ScrollToScreenMiddle); + km.bind(key("L"), Action::ScrollToScreenBottom); + + // --- Arrow keys (same as emacs for convenience) --- + bind_scroll_key(&mut km, "down", Action::SelectNext, !invert, scroll_exits); + bind_scroll_key(&mut km, "up", Action::SelectPrevious, invert, scroll_exits); + + // --- Page scroll --- + km.bind(key("pagedown"), Action::ScrollPageDown); + km.bind(key("pageup"), Action::ScrollPageUp); + + // --- Accept --- + let accept = accept_action(settings); + km.bind(key("enter"), accept); + + km +} + +// --------------------------------------------------------------------------- +// Vim Insert keymap +// --------------------------------------------------------------------------- + +/// Build the default vim-insert keymap. This clones the emacs keymap and +/// overlays vim-insert-specific bindings (esc → enter normal mode). +pub fn default_vim_insert_keymap(settings: &Settings) -> Keymap { + let mut km = default_emacs_keymap(settings); + + // Override esc and ctrl-[ to enter normal mode instead of exiting + km.bind(key("esc"), Action::VimEnterNormal); + km.bind(key("ctrl-["), Action::VimEnterNormal); + + km +} + +// --------------------------------------------------------------------------- +// Inspector keymap +// --------------------------------------------------------------------------- + +/// Build the default inspector keymap (tab index 1). +/// +/// The inspector shows details about the selected history item and has no +/// text input, so we build a minimal keymap with only inspector-relevant +/// bindings. We respect the user's `keymap_mode` to provide vim-style j/k +/// navigation for vim users. +pub fn default_inspector_keymap(settings: &Settings) -> Keymap { + use atuin_client::settings::KeymapMode; + + let mut km = Keymap::new(); + + // Common bindings (same as search tab) + km.bind(key("ctrl-c"), Action::ReturnOriginal); + km.bind(key("ctrl-g"), Action::ReturnOriginal); + km.bind(key("esc"), Action::Exit); + km.bind(key("ctrl-["), Action::Exit); + km.bind(key("tab"), Action::ReturnSelection); + km.bind(key("ctrl-o"), Action::ToggleTab); + + // Accept behavior respects enter_accept setting + let accept = if settings.enter_accept { + Action::Accept + } else { + Action::ReturnSelection + }; + km.bind(key("enter"), accept); + + // Inspector-specific: delete history entry + km.bind(key("ctrl-d"), Action::Delete); + + // Inspector navigation + km.bind(key("up"), Action::InspectPrevious); + km.bind(key("down"), Action::InspectNext); + km.bind(key("pageup"), Action::InspectPrevious); + km.bind(key("pagedown"), Action::InspectNext); + + // For vim users, add j/k navigation + if matches!( + settings.keymap_mode, + KeymapMode::VimNormal | KeymapMode::VimInsert + ) { + km.bind(key("j"), Action::InspectNext); + km.bind(key("k"), Action::InspectPrevious); + } + + km +} + +// --------------------------------------------------------------------------- +// Prefix keymap +// --------------------------------------------------------------------------- + +/// Build the default prefix keymap (active after ctrl-a prefix). +pub fn default_prefix_keymap() -> Keymap { + let mut km = Keymap::new(); + + km.bind(key("d"), Action::Delete); + km.bind(key("a"), Action::CursorStart); + km.bind_conditional( + key("c"), + vec![ + KeyRule::when(ConditionAtom::HasContext, Action::ClearContext), + KeyRule::always(Action::SwitchContext), + ], + ); + + km +} + +// --------------------------------------------------------------------------- +// KeymapSet construction +// --------------------------------------------------------------------------- + +// --------------------------------------------------------------------------- +// Config → Keymap conversion +// --------------------------------------------------------------------------- + +/// Convert a `KeyBindingConfig` (from TOML) into a `KeyBinding`. +/// Returns `Err` if an action name or condition expression is invalid. +fn parse_binding_config(config: &KeyBindingConfig) -> Result<KeyBinding, String> { + match config { + KeyBindingConfig::Simple(action_str) => { + let action = Action::from_str(action_str)?; + Ok(KeyBinding::simple(action)) + } + KeyBindingConfig::Rules(rules) => { + let mut parsed_rules = Vec::with_capacity(rules.len()); + for rule_cfg in rules { + let action = Action::from_str(&rule_cfg.action)?; + let rule = match &rule_cfg.when { + None => KeyRule::always(action), + Some(cond_str) => { + let cond = ConditionExpr::parse(cond_str)?; + KeyRule::when(cond, action) + } + }; + parsed_rules.push(rule); + } + Ok(KeyBinding::conditional(parsed_rules)) + } + } +} + +/// Apply a map of key-string → binding-config overrides to a keymap. +/// Per-key override replaces the entire rule list for that key. +/// Invalid keys or action names are logged and skipped. +fn apply_config_to_keymap(keymap: &mut Keymap, overrides: &HashMap<String, KeyBindingConfig>) { + for (key_str, binding_cfg) in overrides { + let key = match KeyInput::parse(key_str) { + Ok(k) => k, + Err(e) => { + warn!("invalid key in keymap config: {key_str:?}: {e}"); + continue; + } + }; + match parse_binding_config(binding_cfg) { + Ok(binding) => { + keymap.bindings.insert(key, binding); + } + Err(e) => { + warn!("invalid binding for {key_str:?} in keymap config: {e}"); + } + } + } +} + +impl KeymapSet { + /// Build the complete set of default keymaps from settings. + pub fn defaults(settings: &Settings) -> Self { + KeymapSet { + emacs: default_emacs_keymap(settings), + vim_normal: default_vim_normal_keymap(settings), + vim_insert: default_vim_insert_keymap(settings), + inspector: default_inspector_keymap(settings), + prefix: default_prefix_keymap(), + } + } + + /// Build keymaps from settings, applying any user `[keymap]` overrides. + /// + /// Precedence rules: + /// - If `[keymap]` has any entries, `[keys]` is **ignored entirely**. + /// Defaults are built with standard `[keys]` values, then `[keymap]` + /// overrides are applied per-key. + /// - If `[keymap]` is empty/absent, `[keys]` customizes the defaults + /// (current behavior for backward compatibility). + pub fn from_settings(settings: &Settings) -> Self { + use atuin_client::settings::Keys; + + if settings.keymap.is_empty() { + // No [keymap] section → use [keys] to customize defaults + Self::defaults(settings) + } else { + // [keymap] present → ignore [keys], use standard defaults as base + let mut base_settings = settings.clone(); + base_settings.keys = Keys::standard_defaults(); + let mut set = Self::defaults(&base_settings); + set.apply_config(settings); + set + } + } + + /// Apply user keymap config overrides to all modes. + fn apply_config(&mut self, settings: &Settings) { + let config = &settings.keymap; + apply_config_to_keymap(&mut self.emacs, &config.emacs); + apply_config_to_keymap(&mut self.vim_normal, &config.vim_normal); + apply_config_to_keymap(&mut self.vim_insert, &config.vim_insert); + apply_config_to_keymap(&mut self.inspector, &config.inspector); + apply_config_to_keymap(&mut self.prefix, &config.prefix); + } +} + +// --------------------------------------------------------------------------- +// Tests +// --------------------------------------------------------------------------- + +#[cfg(test)] +mod tests { + use super::*; + use crate::command::client::search::keybindings::conditions::EvalContext; + + fn make_ctx(cursor: usize, width: usize, selected: usize, len: usize) -> EvalContext { + EvalContext { + cursor_position: cursor, + input_width: width, + input_byte_len: width, + selected_index: selected, + results_len: len, + original_input_empty: false, + has_context: false, + } + } + + fn default_settings() -> Settings { + Settings::utc() + } + + // -- Emacs keymap tests -- + + #[test] + fn emacs_ctrl_c_returns_original() { + let km = default_emacs_keymap(&default_settings()); + let ctx = make_ctx(0, 0, 0, 10); + assert_eq!( + km.resolve(&key("ctrl-c"), &ctx), + Some(Action::ReturnOriginal) + ); + } + + #[test] + fn emacs_esc_exits() { + let km = default_emacs_keymap(&default_settings()); + let ctx = make_ctx(0, 0, 0, 10); + assert_eq!(km.resolve(&key("esc"), &ctx), Some(Action::Exit)); + } + + #[test] + fn emacs_tab_returns_selection() { + // enter_accept=false in test defaults → ReturnSelection + let km = default_emacs_keymap(&default_settings()); + let ctx = make_ctx(0, 0, 0, 10); + assert_eq!(km.resolve(&key("tab"), &ctx), Some(Action::ReturnSelection)); + } + + #[test] + fn emacs_enter_returns_selection() { + // enter_accept=false in test defaults → ReturnSelection + let km = default_emacs_keymap(&default_settings()); + let ctx = make_ctx(0, 0, 0, 10); + assert_eq!( + km.resolve(&key("enter"), &ctx), + Some(Action::ReturnSelection) + ); + } + + #[test] + fn emacs_enter_accept_true_uses_accept() { + let mut settings = default_settings(); + settings.enter_accept = true; + let km = default_emacs_keymap(&settings); + let ctx = make_ctx(0, 0, 0, 10); + assert_eq!(km.resolve(&key("enter"), &ctx), Some(Action::Accept)); + assert_eq!(km.resolve(&key("tab"), &ctx), Some(Action::ReturnSelection)); + } + + #[test] + fn emacs_right_at_end_returns_selection() { + let km = default_emacs_keymap(&default_settings()); + // cursor at end of "hello" (width 5) + let ctx = make_ctx(5, 5, 0, 10); + assert_eq!( + km.resolve(&key("right"), &ctx), + Some(Action::ReturnSelection) + ); + } + + #[test] + fn emacs_right_not_at_end_moves() { + let km = default_emacs_keymap(&default_settings()); + let ctx = make_ctx(2, 5, 0, 10); + assert_eq!(km.resolve(&key("right"), &ctx), Some(Action::CursorRight)); + } + + #[test] + fn emacs_left_at_start_exits() { + let km = default_emacs_keymap(&default_settings()); + let ctx = make_ctx(0, 5, 0, 10); + assert_eq!(km.resolve(&key("left"), &ctx), Some(Action::Exit)); + } + + #[test] + fn emacs_left_not_at_start_moves() { + let km = default_emacs_keymap(&default_settings()); + let ctx = make_ctx(3, 5, 0, 10); + assert_eq!(km.resolve(&key("left"), &ctx), Some(Action::CursorLeft)); + } + + #[test] + fn emacs_down_at_start_exits() { + let km = default_emacs_keymap(&default_settings()); + // selected=0 → ListAtStart → Exit + let ctx = make_ctx(0, 0, 0, 10); + assert_eq!(km.resolve(&key("down"), &ctx), Some(Action::Exit)); + } + + #[test] + fn emacs_down_not_at_start_selects_next() { + let km = default_emacs_keymap(&default_settings()); + // selected=5 → not at start → SelectNext + let ctx = make_ctx(0, 0, 5, 10); + assert_eq!(km.resolve(&key("down"), &ctx), Some(Action::SelectNext)); + } + + #[test] + fn emacs_up_selects_previous() { + let km = default_emacs_keymap(&default_settings()); + // Non-inverted: up never exits (moves away from index 0) + let ctx = make_ctx(0, 0, 5, 10); + assert_eq!(km.resolve(&key("up"), &ctx), Some(Action::SelectPrevious)); + } + + #[test] + fn emacs_ctrl_d_empty_returns_original() { + let km = default_emacs_keymap(&default_settings()); + // input empty (byte_len = 0) + let ctx = make_ctx(0, 0, 0, 10); + assert_eq!( + km.resolve(&key("ctrl-d"), &ctx), + Some(Action::ReturnOriginal) + ); + } + + #[test] + fn emacs_ctrl_d_nonempty_deletes() { + let km = default_emacs_keymap(&default_settings()); + let ctx = make_ctx(2, 5, 0, 10); + assert_eq!( + km.resolve(&key("ctrl-d"), &ctx), + Some(Action::DeleteCharAfter) + ); + } + + #[test] + fn emacs_ctrl_n_selects_next_no_exit_condition() { + let km = default_emacs_keymap(&default_settings()); + // at start, but ctrl-n should NOT exit (no exit condition bound) + let ctx = make_ctx(0, 0, 0, 10); + assert_eq!(km.resolve(&key("ctrl-n"), &ctx), Some(Action::SelectNext)); + } + + #[test] + fn emacs_prefix_key_enters_prefix() { + let km = default_emacs_keymap(&default_settings()); + let ctx = make_ctx(0, 0, 0, 10); + assert_eq!( + km.resolve(&key("ctrl-a"), &ctx), + Some(Action::EnterPrefixMode) + ); + } + + #[test] + fn emacs_home_cursor_start() { + let km = default_emacs_keymap(&default_settings()); + let ctx = make_ctx(5, 10, 0, 10); + assert_eq!(km.resolve(&key("home"), &ctx), Some(Action::CursorStart)); + } + + // -- Vim Normal keymap tests -- + + #[test] + fn vim_normal_j_at_start_exits() { + let km = default_vim_normal_keymap(&default_settings()); + // selected=0 → ListAtStart → Exit (non-inverted: j moves toward index 0) + let ctx = make_ctx(0, 0, 0, 10); + assert_eq!(km.resolve(&key("j"), &ctx), Some(Action::Exit)); + } + + #[test] + fn vim_normal_j_not_at_start_selects_next() { + let km = default_vim_normal_keymap(&default_settings()); + let ctx = make_ctx(0, 0, 5, 10); + assert_eq!(km.resolve(&key("j"), &ctx), Some(Action::SelectNext)); + } + + #[test] + fn vim_normal_k_selects_previous() { + let km = default_vim_normal_keymap(&default_settings()); + // Non-inverted: k never exits (moves away from index 0) + let ctx = make_ctx(0, 0, 5, 10); + assert_eq!(km.resolve(&key("k"), &ctx), Some(Action::SelectPrevious)); + } + + #[test] + fn vim_normal_i_enters_insert() { + let km = default_vim_normal_keymap(&default_settings()); + let ctx = make_ctx(0, 0, 0, 10); + assert_eq!(km.resolve(&key("i"), &ctx), Some(Action::VimEnterInsert)); + } + + #[test] + fn vim_normal_slash_search_insert() { + let km = default_vim_normal_keymap(&default_settings()); + let ctx = make_ctx(0, 0, 0, 10); + assert_eq!(km.resolve(&key("/"), &ctx), Some(Action::VimSearchInsert)); + } + + #[test] + fn vim_normal_gg_scroll_to_top() { + let km = default_vim_normal_keymap(&default_settings()); + let ctx = make_ctx(0, 0, 50, 100); + assert_eq!(km.resolve(&key("g g"), &ctx), Some(Action::ScrollToTop)); + } + + #[test] + fn vim_normal_big_g_scroll_to_bottom() { + let km = default_vim_normal_keymap(&default_settings()); + let ctx = make_ctx(0, 0, 50, 100); + assert_eq!(km.resolve(&key("G"), &ctx), Some(Action::ScrollToBottom)); + } + + #[test] + fn vim_normal_numeric_returns_selection() { + let km = default_vim_normal_keymap(&default_settings()); + let ctx = make_ctx(0, 0, 0, 10); + assert_eq!( + km.resolve(&key("3"), &ctx), + Some(Action::ReturnSelectionNth(3)) + ); + } + + #[test] + fn vim_normal_ctrl_u_half_page_up() { + let km = default_vim_normal_keymap(&default_settings()); + let ctx = make_ctx(0, 0, 50, 100); + assert_eq!( + km.resolve(&key("ctrl-u"), &ctx), + Some(Action::ScrollHalfPageUp) + ); + } + + #[test] + fn vim_normal_screen_jumps() { + let km = default_vim_normal_keymap(&default_settings()); + let ctx = make_ctx(0, 0, 50, 100); + assert_eq!(km.resolve(&key("H"), &ctx), Some(Action::ScrollToScreenTop)); + assert_eq!( + km.resolve(&key("M"), &ctx), + Some(Action::ScrollToScreenMiddle) + ); + assert_eq!( + km.resolve(&key("L"), &ctx), + Some(Action::ScrollToScreenBottom) + ); + } + + #[test] + fn vim_normal_enter_returns_selection() { + // enter_accept=false in test defaults → ReturnSelection + let km = default_vim_normal_keymap(&default_settings()); + let ctx = make_ctx(0, 0, 0, 10); + assert_eq!( + km.resolve(&key("enter"), &ctx), + Some(Action::ReturnSelection) + ); + } + + #[test] + fn vim_normal_enter_accept_true_uses_accept() { + let mut settings = default_settings(); + settings.enter_accept = true; + let km = default_vim_normal_keymap(&settings); + let ctx = make_ctx(0, 0, 0, 10); + assert_eq!(km.resolve(&key("enter"), &ctx), Some(Action::Accept)); + } + + // -- Vim Insert keymap tests -- + + #[test] + fn vim_insert_inherits_emacs_enter() { + let km = default_vim_insert_keymap(&default_settings()); + let ctx = make_ctx(0, 0, 0, 10); + // enter_accept=false → ReturnSelection + assert_eq!( + km.resolve(&key("enter"), &ctx), + Some(Action::ReturnSelection) + ); + } + + #[test] + fn vim_insert_esc_enters_normal() { + let km = default_vim_insert_keymap(&default_settings()); + let ctx = make_ctx(0, 0, 0, 10); + assert_eq!(km.resolve(&key("esc"), &ctx), Some(Action::VimEnterNormal)); + } + + #[test] + fn vim_insert_ctrl_bracket_enters_normal() { + let km = default_vim_insert_keymap(&default_settings()); + let ctx = make_ctx(0, 0, 0, 10); + assert_eq!( + km.resolve(&key("ctrl-["), &ctx), + Some(Action::VimEnterNormal) + ); + } + + #[test] + fn vim_insert_inherits_emacs_ctrl_d() { + let km = default_vim_insert_keymap(&default_settings()); + let ctx = make_ctx(0, 0, 0, 10); + // input empty → return original + assert_eq!( + km.resolve(&key("ctrl-d"), &ctx), + Some(Action::ReturnOriginal) + ); + } + + // -- Inspector keymap tests -- + + #[test] + fn inspector_ctrl_d_deletes() { + let km = default_inspector_keymap(&default_settings()); + let ctx = make_ctx(0, 0, 0, 10); + assert_eq!(km.resolve(&key("ctrl-d"), &ctx), Some(Action::Delete)); + } + + #[test] + fn inspector_up_inspects_previous() { + let km = default_inspector_keymap(&default_settings()); + let ctx = make_ctx(0, 0, 0, 10); + assert_eq!(km.resolve(&key("up"), &ctx), Some(Action::InspectPrevious)); + } + + #[test] + fn inspector_down_inspects_next() { + let km = default_inspector_keymap(&default_settings()); + let ctx = make_ctx(0, 0, 0, 10); + assert_eq!(km.resolve(&key("down"), &ctx), Some(Action::InspectNext)); + } + + #[test] + fn inspector_esc_exits() { + let km = default_inspector_keymap(&default_settings()); + let ctx = make_ctx(0, 0, 0, 10); + assert_eq!(km.resolve(&key("esc"), &ctx), Some(Action::Exit)); + } + + #[test] + fn inspector_tab_returns_selection() { + // enter_accept=false → ReturnSelection + let km = default_inspector_keymap(&default_settings()); + let ctx = make_ctx(0, 0, 0, 10); + assert_eq!(km.resolve(&key("tab"), &ctx), Some(Action::ReturnSelection)); + } + + // -- Prefix keymap tests -- + + #[test] + fn prefix_d_deletes() { + let km = default_prefix_keymap(); + let ctx = make_ctx(0, 0, 0, 10); + assert_eq!(km.resolve(&key("d"), &ctx), Some(Action::Delete)); + } + + #[test] + fn prefix_a_cursor_start() { + let km = default_prefix_keymap(); + let ctx = make_ctx(0, 0, 0, 10); + assert_eq!(km.resolve(&key("a"), &ctx), Some(Action::CursorStart)); + } + + #[test] + fn prefix_unknown_key_returns_none() { + let km = default_prefix_keymap(); + let ctx = make_ctx(0, 0, 0, 10); + assert_eq!(km.resolve(&key("x"), &ctx), None); + } + + // -- KeymapSet tests -- + + #[test] + fn keymap_set_defaults_builds() { + let settings = default_settings(); + let set = KeymapSet::defaults(&settings); + let ctx = make_ctx(0, 0, 0, 10); + + // Sanity check each keymap has bindings + assert!(set.emacs.resolve(&key("ctrl-c"), &ctx).is_some()); + assert!(set.vim_normal.resolve(&key("ctrl-c"), &ctx).is_some()); + assert!(set.vim_insert.resolve(&key("ctrl-c"), &ctx).is_some()); + assert!(set.inspector.resolve(&key("ctrl-c"), &ctx).is_some()); + assert!(set.prefix.resolve(&key("d"), &ctx).is_some()); + } + + // -- Settings-dependent behavior -- + + #[test] + fn custom_prefix_char() { + let mut settings = default_settings(); + settings.keys.prefix = "x".to_string(); + let km = default_emacs_keymap(&settings); + let ctx = make_ctx(0, 0, 0, 10); + + // ctrl-x should be prefix mode + assert_eq!( + km.resolve(&key("ctrl-x"), &ctx), + Some(Action::EnterPrefixMode) + ); + // ctrl-a should now be CursorStart (not prefix) + assert_eq!(km.resolve(&key("ctrl-a"), &ctx), Some(Action::CursorStart)); + } + + #[test] + fn ctrl_n_shortcuts_changes_numeric_modifier() { + let mut settings = default_settings(); + settings.ctrl_n_shortcuts = true; + let km = default_emacs_keymap(&settings); + let ctx = make_ctx(0, 0, 0, 10); + + // ctrl-1 should work + assert_eq!( + km.resolve(&key("ctrl-1"), &ctx), + Some(Action::ReturnSelectionNth(1)) + ); + // alt-1 should NOT be bound + assert_eq!(km.resolve(&key("alt-1"), &ctx), None); + } + + #[test] + fn default_alt_numeric_shortcuts() { + let settings = default_settings(); + let km = default_emacs_keymap(&settings); + let ctx = make_ctx(0, 0, 0, 10); + + // alt-1 should work by default + assert_eq!( + km.resolve(&key("alt-1"), &ctx), + Some(Action::ReturnSelectionNth(1)) + ); + } + + // ----------------------------------------------------------------------- + // Config parsing and merging tests + // ----------------------------------------------------------------------- + + #[test] + fn parse_simple_binding_config() { + use atuin_client::settings::KeyBindingConfig; + let cfg = KeyBindingConfig::Simple("accept".to_string()); + let binding = super::parse_binding_config(&cfg).unwrap(); + assert_eq!(binding.rules.len(), 1); + assert!(binding.rules[0].condition.is_none()); + assert_eq!(binding.rules[0].action, Action::Accept); + } + + #[test] + fn parse_conditional_binding_config() { + use atuin_client::settings::{KeyBindingConfig, KeyRuleConfig}; + let cfg = KeyBindingConfig::Rules(vec![ + KeyRuleConfig { + when: Some("cursor-at-start".to_string()), + action: "exit".to_string(), + }, + KeyRuleConfig { + when: None, + action: "cursor-left".to_string(), + }, + ]); + let binding = super::parse_binding_config(&cfg).unwrap(); + assert_eq!(binding.rules.len(), 2); + assert!(binding.rules[0].condition.is_some()); + assert_eq!(binding.rules[0].action, Action::Exit); + assert!(binding.rules[1].condition.is_none()); + assert_eq!(binding.rules[1].action, Action::CursorLeft); + } + + #[test] + fn parse_binding_config_invalid_action() { + use atuin_client::settings::KeyBindingConfig; + let cfg = KeyBindingConfig::Simple("not-a-real-action".to_string()); + assert!(super::parse_binding_config(&cfg).is_err()); + } + + #[test] + fn parse_binding_config_invalid_condition() { + use atuin_client::settings::{KeyBindingConfig, KeyRuleConfig}; + let cfg = KeyBindingConfig::Rules(vec![KeyRuleConfig { + when: Some("not-a-real-condition".to_string()), + action: "exit".to_string(), + }]); + assert!(super::parse_binding_config(&cfg).is_err()); + } + + #[test] + fn config_override_replaces_key() { + use atuin_client::settings::KeyBindingConfig; + use std::collections::HashMap; + + let mut settings = default_settings(); + let set = KeymapSet::defaults(&settings); + + // Default: ctrl-c → ReturnOriginal + let ctx = make_ctx(0, 0, 0, 10); + assert_eq!( + set.emacs.resolve(&key("ctrl-c"), &ctx), + Some(Action::ReturnOriginal) + ); + + // Override ctrl-c → Exit via config + settings.keymap.emacs = HashMap::from([( + "ctrl-c".to_string(), + KeyBindingConfig::Simple("exit".to_string()), + )]); + + let set = KeymapSet::from_settings(&settings); + assert_eq!(set.emacs.resolve(&key("ctrl-c"), &ctx), Some(Action::Exit)); + } + + #[test] + fn config_override_preserves_unoverridden_keys() { + use atuin_client::settings::KeyBindingConfig; + use std::collections::HashMap; + + let mut settings = default_settings(); + // Override only ctrl-c; enter should keep its default + settings.keymap.emacs = HashMap::from([( + "ctrl-c".to_string(), + KeyBindingConfig::Simple("exit".to_string()), + )]); + + let set = KeymapSet::from_settings(&settings); + let ctx = make_ctx(0, 0, 0, 10); + + // ctrl-c overridden + assert_eq!(set.emacs.resolve(&key("ctrl-c"), &ctx), Some(Action::Exit)); + // enter still has default (enter_accept=false → ReturnSelection) + assert_eq!( + set.emacs.resolve(&key("enter"), &ctx), + Some(Action::ReturnSelection) + ); + } + + #[test] + fn config_conditional_override() { + use atuin_client::settings::{KeyBindingConfig, KeyRuleConfig}; + use std::collections::HashMap; + + let mut settings = default_settings(); + // Override "up" with a custom conditional + settings.keymap.emacs = HashMap::from([( + "up".to_string(), + KeyBindingConfig::Rules(vec![ + KeyRuleConfig { + when: Some("no-results".to_string()), + action: "exit".to_string(), + }, + KeyRuleConfig { + when: None, + action: "select-previous".to_string(), + }, + ]), + )]); + + let set = KeymapSet::from_settings(&settings); + + // With no results → exit + let ctx = make_ctx(0, 0, 0, 0); + assert_eq!(set.emacs.resolve(&key("up"), &ctx), Some(Action::Exit)); + + // With results → select-previous + let ctx = make_ctx(0, 0, 0, 10); + assert_eq!( + set.emacs.resolve(&key("up"), &ctx), + Some(Action::SelectPrevious) + ); + } + + #[test] + fn from_settings_with_empty_config_equals_defaults() { + let settings = default_settings(); + let defaults = KeymapSet::defaults(&settings); + let from_settings = KeymapSet::from_settings(&settings); + + // Verify a sample of keys produce the same results + let ctx = make_ctx(0, 0, 0, 10); + let test_keys = [ + "ctrl-c", "enter", "esc", "tab", "up", "down", "left", "right", + ]; + for k in &test_keys { + assert_eq!( + defaults.emacs.resolve(&key(k), &ctx), + from_settings.emacs.resolve(&key(k), &ctx), + "mismatch for emacs key {k}" + ); + } + } + + // ----------------------------------------------------------------------- + // Phase 5: [keys] vs [keymap] backward compatibility + // ----------------------------------------------------------------------- + + #[test] + fn keymap_overrides_ignore_keys_section() { + use atuin_client::settings::KeyBindingConfig; + + // Set up: [keys] disables scroll_exits, but [keymap] is present + let mut settings = default_settings(); + settings.keys.scroll_exits = false; + + // Without [keymap], scroll_exits=false means no exit condition on down + let set_legacy = KeymapSet::defaults(&settings); + // At list-at-start (selected=0), down should still be SelectNext (no exit) + let ctx_at_boundary = make_ctx(0, 0, 0, 10); + assert_eq!( + set_legacy.emacs.resolve(&key("down"), &ctx_at_boundary), + Some(Action::SelectNext), + "legacy: down at boundary should be SelectNext with scroll_exits=false" + ); + + // With [keymap] present (even just one override), [keys] is ignored + // so the standard defaults (scroll_exits=true) apply + settings.keymap.emacs = HashMap::from([( + "ctrl-c".to_string(), + KeyBindingConfig::Simple("exit".to_string()), + )]); + let set_keymap = KeymapSet::from_settings(&settings); + + // Not at boundary (selected=5): should SelectNext normally + let ctx_not_at_boundary = make_ctx(0, 0, 5, 10); + assert_eq!( + set_keymap.emacs.resolve(&key("down"), &ctx_not_at_boundary), + Some(Action::SelectNext), + "keymap: down not at boundary should SelectNext" + ); + // At list-at-start (selected=0): should Exit (standard scroll_exits=true) + assert_eq!( + set_keymap.emacs.resolve(&key("down"), &ctx_at_boundary), + Some(Action::Exit), + "keymap: down at boundary should Exit (standard defaults restored)" + ); + } + + #[test] + fn keymap_present_resets_to_standard_keys_defaults() { + use atuin_client::settings::KeyBindingConfig; + + let mut settings = default_settings(); + // Disable all [keys] behaviors + settings.keys.exit_past_line_start = false; + settings.keys.accept_past_line_end = false; + + // Without [keymap], left should be plain CursorLeft + let set_legacy = KeymapSet::defaults(&settings); + let ctx_at_start = make_ctx(0, 5, 0, 10); + assert_eq!( + set_legacy.emacs.resolve(&key("left"), &ctx_at_start), + Some(Action::CursorLeft), + "legacy: left should be plain CursorLeft without exit_past_line_start" + ); + + // Add a [keymap] entry (for a different key) + settings.keymap.emacs = HashMap::from([( + "ctrl-c".to_string(), + KeyBindingConfig::Simple("exit".to_string()), + )]); + let set_keymap = KeymapSet::from_settings(&settings); + + // Now left should use standard defaults (exit_past_line_start=true) + // At cursor start → Exit + assert_eq!( + set_keymap.emacs.resolve(&key("left"), &ctx_at_start), + Some(Action::Exit), + "keymap: left at cursor start should exit (standard defaults)" + ); + + // Right at cursor end should return selection (standard defaults: accept_past_line_end=true, enter_accept=false) + let ctx_at_end = make_ctx(5, 5, 0, 10); + assert_eq!( + set_keymap.emacs.resolve(&key("right"), &ctx_at_end), + Some(Action::ReturnSelection), + "keymap: right at cursor end should return selection (standard defaults)" + ); + } + + #[test] + fn keys_has_non_default_values_detection() { + use atuin_client::settings::Keys; + + let standard = Keys::standard_defaults(); + assert!(!standard.has_non_default_values()); + + let mut modified = Keys::standard_defaults(); + modified.scroll_exits = false; + assert!(modified.has_non_default_values()); + + let mut modified = Keys::standard_defaults(); + modified.prefix = "x".to_string(); + assert!(modified.has_non_default_values()); + } + + #[test] + fn original_input_empty_condition_in_config() { + use atuin_client::settings::{KeyBindingConfig, KeyRuleConfig}; + use std::collections::HashMap; + + let mut settings = default_settings(); + // Configure esc to: if original-input-empty -> return-query, else return-original + settings.keymap.emacs = HashMap::from([( + "esc".to_string(), + KeyBindingConfig::Rules(vec![ + KeyRuleConfig { + when: Some("original-input-empty".to_string()), + action: "return-query".to_string(), + }, + KeyRuleConfig { + when: None, + action: "return-original".to_string(), + }, + ]), + )]); + + let set = KeymapSet::from_settings(&settings); + + // When original input was empty, should return-query + let ctx_original_empty = EvalContext { + cursor_position: 0, + input_width: 5, + input_byte_len: 5, + selected_index: 0, + results_len: 10, + original_input_empty: true, + has_context: false, + }; + assert_eq!( + set.emacs.resolve(&key("esc"), &ctx_original_empty), + Some(Action::ReturnQuery), + "esc with original_input_empty=true should return-query" + ); + + // When original input was not empty, should return-original + let ctx_original_not_empty = EvalContext { + cursor_position: 0, + input_width: 5, + input_byte_len: 5, + selected_index: 0, + results_len: 10, + original_input_empty: false, + has_context: false, + }; + assert_eq!( + set.emacs.resolve(&key("esc"), &ctx_original_not_empty), + Some(Action::ReturnOriginal), + "esc with original_input_empty=false should return-original" + ); + } +} diff --git a/crates/atuin/src/command/client/search/keybindings/key.rs b/crates/atuin/src/command/client/search/keybindings/key.rs new file mode 100644 index 00000000..717b406d --- /dev/null +++ b/crates/atuin/src/command/client/search/keybindings/key.rs @@ -0,0 +1,629 @@ +use std::fmt; + +use ratatui::crossterm::event::{KeyCode, KeyEvent, KeyModifiers, MediaKeyCode}; +use serde::{Deserialize, Deserializer, Serialize, Serializer}; + +/// A single key press with modifiers (e.g. `ctrl-c`, `alt-f`, `enter`). +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +#[allow(clippy::struct_excessive_bools)] +pub struct SingleKey { + pub code: KeyCodeValue, + pub ctrl: bool, + pub alt: bool, + pub shift: bool, + pub super_key: bool, +} + +/// The key code portion of a key press. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum KeyCodeValue { + Char(char), + Enter, + Esc, + Tab, + Backspace, + Delete, + Insert, + Up, + Down, + Left, + Right, + Home, + End, + PageUp, + PageDown, + Space, + F(u8), + Media(MediaKeyCode), +} + +/// A key input that may be a single key or a multi-key sequence (e.g. `g g`). +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub enum KeyInput { + Single(SingleKey), + Sequence(Vec<SingleKey>), +} + +impl SingleKey { + /// Convert a crossterm `KeyEvent` into a `SingleKey`. + pub fn from_event(event: &KeyEvent) -> Option<Self> { + let ctrl = event.modifiers.contains(KeyModifiers::CONTROL); + let alt = event.modifiers.contains(KeyModifiers::ALT); + let shift = event.modifiers.contains(KeyModifiers::SHIFT); + let super_key = event.modifiers.contains(KeyModifiers::SUPER); + + let code = match event.code { + KeyCode::Char(' ') => KeyCodeValue::Space, + KeyCode::Char(c) => { + // If shift is the only modifier and it's an uppercase letter, + // we store the uppercase char directly and clear the shift flag + // since the case already encodes it. + if shift && !ctrl && !alt && !super_key && c.is_ascii_uppercase() { + return Some(SingleKey { + code: KeyCodeValue::Char(c), + ctrl: false, + alt: false, + shift: false, + super_key: false, + }); + } + KeyCodeValue::Char(c) + } + KeyCode::Enter => KeyCodeValue::Enter, + KeyCode::Esc => KeyCodeValue::Esc, + KeyCode::Tab => KeyCodeValue::Tab, + // BackTab is sent by many terminals for Shift+Tab + KeyCode::BackTab => { + return Some(SingleKey { + code: KeyCodeValue::Tab, + ctrl, + alt, + shift: true, + super_key, + }); + } + KeyCode::Backspace => KeyCodeValue::Backspace, + KeyCode::Delete => KeyCodeValue::Delete, + KeyCode::Insert => KeyCodeValue::Insert, + KeyCode::Up => KeyCodeValue::Up, + KeyCode::Down => KeyCodeValue::Down, + KeyCode::Left => KeyCodeValue::Left, + KeyCode::Right => KeyCodeValue::Right, + KeyCode::Home => KeyCodeValue::Home, + KeyCode::End => KeyCodeValue::End, + KeyCode::PageUp => KeyCodeValue::PageUp, + KeyCode::PageDown => KeyCodeValue::PageDown, + KeyCode::F(n) => KeyCodeValue::F(n), + KeyCode::Media(m) => KeyCodeValue::Media(m), + _ => return None, + }; + + Some(SingleKey { + code, + ctrl, + alt, + shift: if matches!(code, KeyCodeValue::Char(_)) { + false + } else { + shift + }, + super_key, + }) + } + + /// Parse a key string like `"ctrl-c"`, `"alt-f"`, `"enter"`, `"G"`. + pub fn parse(s: &str) -> Result<Self, String> { + let s = s.trim(); + let parts: Vec<&str> = s.split('-').collect(); + + let mut ctrl = false; + let mut alt = false; + let mut shift = false; + let mut super_key = false; + + // All parts except the last are modifiers + for &part in &parts[..parts.len() - 1] { + match part.to_lowercase().as_str() { + "ctrl" => ctrl = true, + "alt" => alt = true, + "shift" => shift = true, + "super" | "cmd" | "win" => super_key = true, + _ => return Err(format!("unknown modifier: {part}")), + } + } + + let key_part = parts[parts.len() - 1]; + let code = match key_part.to_lowercase().as_str() { + "enter" | "return" => KeyCodeValue::Enter, + "esc" | "escape" => KeyCodeValue::Esc, + "tab" => KeyCodeValue::Tab, + "backspace" => KeyCodeValue::Backspace, + "delete" | "del" => KeyCodeValue::Delete, + "insert" | "ins" => KeyCodeValue::Insert, + "up" => KeyCodeValue::Up, + "down" => KeyCodeValue::Down, + "left" => KeyCodeValue::Left, + "right" => KeyCodeValue::Right, + "home" => KeyCodeValue::Home, + "end" => KeyCodeValue::End, + "pageup" => KeyCodeValue::PageUp, + "pagedown" => KeyCodeValue::PageDown, + "space" => KeyCodeValue::Space, + s if s.starts_with('f') && s.len() > 1 => { + // Parse function keys like "f1", "f12" + if let Ok(n) = s[1..].parse::<u8>() { + if (1..=24).contains(&n) { + KeyCodeValue::F(n) + } else { + return Err(format!("function key out of range: {key_part}")); + } + } else { + return Err(format!("unknown key: {key_part}")); + } + } + "[" => KeyCodeValue::Char('['), + "]" => KeyCodeValue::Char(']'), + "?" => KeyCodeValue::Char('?'), + "/" => KeyCodeValue::Char('/'), + "$" => KeyCodeValue::Char('$'), + // Media keys (no dashes - the parser splits on dash for modifiers) + "play" => KeyCodeValue::Media(MediaKeyCode::Play), + "pause" => KeyCodeValue::Media(MediaKeyCode::Pause), + "playpause" => KeyCodeValue::Media(MediaKeyCode::PlayPause), + "stop" => KeyCodeValue::Media(MediaKeyCode::Stop), + "fastforward" => KeyCodeValue::Media(MediaKeyCode::FastForward), + "rewind" => KeyCodeValue::Media(MediaKeyCode::Rewind), + "tracknext" => KeyCodeValue::Media(MediaKeyCode::TrackNext), + "trackprevious" => KeyCodeValue::Media(MediaKeyCode::TrackPrevious), + "record" => KeyCodeValue::Media(MediaKeyCode::Record), + "lowervolume" => KeyCodeValue::Media(MediaKeyCode::LowerVolume), + "raisevolume" => KeyCodeValue::Media(MediaKeyCode::RaiseVolume), + "mutevolume" | "mute" => KeyCodeValue::Media(MediaKeyCode::MuteVolume), + _ => { + let chars: Vec<char> = key_part.chars().collect(); + if chars.len() == 1 { + let c = chars[0]; + // An uppercase letter implies shift (unless shift already specified) + if c.is_ascii_uppercase() && !ctrl && !alt && !super_key { + return Ok(SingleKey { + code: KeyCodeValue::Char(c), + ctrl: false, + alt: false, + shift: false, + super_key: false, + }); + } + KeyCodeValue::Char(c) + } else { + return Err(format!("unknown key: {key_part}")); + } + } + }; + + Ok(SingleKey { + code, + ctrl, + alt, + shift, + super_key, + }) + } +} + +impl fmt::Display for SingleKey { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + if self.super_key { + write!(f, "super-")?; + } + if self.ctrl { + write!(f, "ctrl-")?; + } + if self.alt { + write!(f, "alt-")?; + } + if self.shift { + write!(f, "shift-")?; + } + match &self.code { + KeyCodeValue::Char(c) => write!(f, "{c}"), + KeyCodeValue::Enter => write!(f, "enter"), + KeyCodeValue::Esc => write!(f, "esc"), + KeyCodeValue::Tab => write!(f, "tab"), + KeyCodeValue::Backspace => write!(f, "backspace"), + KeyCodeValue::Delete => write!(f, "delete"), + KeyCodeValue::Insert => write!(f, "insert"), + KeyCodeValue::Up => write!(f, "up"), + KeyCodeValue::Down => write!(f, "down"), + KeyCodeValue::Left => write!(f, "left"), + KeyCodeValue::Right => write!(f, "right"), + KeyCodeValue::Home => write!(f, "home"), + KeyCodeValue::End => write!(f, "end"), + KeyCodeValue::PageUp => write!(f, "pageup"), + KeyCodeValue::PageDown => write!(f, "pagedown"), + KeyCodeValue::Space => write!(f, "space"), + KeyCodeValue::F(n) => write!(f, "f{n}"), + KeyCodeValue::Media(m) => match m { + MediaKeyCode::Play => write!(f, "play"), + MediaKeyCode::Pause => write!(f, "media-pause"), + MediaKeyCode::PlayPause => write!(f, "playpause"), + MediaKeyCode::Stop => write!(f, "stop"), + MediaKeyCode::FastForward => write!(f, "fastforward"), + MediaKeyCode::Rewind => write!(f, "rewind"), + MediaKeyCode::TrackNext => write!(f, "tracknext"), + MediaKeyCode::TrackPrevious => write!(f, "trackprevious"), + MediaKeyCode::Record => write!(f, "record"), + MediaKeyCode::LowerVolume => write!(f, "lowervolume"), + MediaKeyCode::RaiseVolume => write!(f, "raisevolume"), + MediaKeyCode::MuteVolume => write!(f, "mutevolume"), + MediaKeyCode::Reverse => write!(f, "reverse"), + }, + } + } +} + +impl KeyInput { + /// Parse a key input string. Supports multi-key sequences separated by spaces + /// (e.g. `"g g"`). + pub fn parse(s: &str) -> Result<Self, String> { + let s = s.trim(); + // Check for space-separated multi-key sequences + // But don't split "space" or modifier combos like "ctrl-a" + let parts: Vec<&str> = s.split_whitespace().collect(); + if parts.len() > 1 { + let keys: Result<Vec<SingleKey>, String> = + parts.iter().map(|p| SingleKey::parse(p)).collect(); + Ok(KeyInput::Sequence(keys?)) + } else { + Ok(KeyInput::Single(SingleKey::parse(s)?)) + } + } +} + +impl fmt::Display for KeyInput { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + KeyInput::Single(k) => write!(f, "{k}"), + KeyInput::Sequence(keys) => { + for (i, k) in keys.iter().enumerate() { + if i > 0 { + write!(f, " ")?; + } + write!(f, "{k}")?; + } + Ok(()) + } + } + } +} + +impl Serialize for KeyInput { + fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> { + serializer.serialize_str(&self.to_string()) + } +} + +impl<'de> Deserialize<'de> for KeyInput { + fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> { + let s = String::deserialize(deserializer)?; + KeyInput::parse(&s).map_err(serde::de::Error::custom) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use ratatui::crossterm::event::{KeyCode, KeyEvent, KeyModifiers}; + + #[test] + fn parse_simple_keys() { + let k = SingleKey::parse("a").unwrap(); + assert_eq!(k.code, KeyCodeValue::Char('a')); + assert!(!k.ctrl && !k.alt && !k.shift); + + let k = SingleKey::parse("enter").unwrap(); + assert_eq!(k.code, KeyCodeValue::Enter); + + let k = SingleKey::parse("esc").unwrap(); + assert_eq!(k.code, KeyCodeValue::Esc); + + let k = SingleKey::parse("tab").unwrap(); + assert_eq!(k.code, KeyCodeValue::Tab); + + let k = SingleKey::parse("space").unwrap(); + assert_eq!(k.code, KeyCodeValue::Space); + } + + #[test] + fn parse_modifiers() { + let k = SingleKey::parse("ctrl-c").unwrap(); + assert_eq!(k.code, KeyCodeValue::Char('c')); + assert!(k.ctrl); + assert!(!k.alt); + + let k = SingleKey::parse("alt-f").unwrap(); + assert_eq!(k.code, KeyCodeValue::Char('f')); + assert!(k.alt); + assert!(!k.ctrl); + + let k = SingleKey::parse("ctrl-alt-x").unwrap(); + assert_eq!(k.code, KeyCodeValue::Char('x')); + assert!(k.ctrl && k.alt); + } + + #[test] + fn parse_uppercase_implies_no_shift_flag() { + let k = SingleKey::parse("G").unwrap(); + assert_eq!(k.code, KeyCodeValue::Char('G')); + assert!(!k.shift); + assert!(!k.ctrl); + } + + #[test] + fn parse_special_chars() { + let k = SingleKey::parse("ctrl-[").unwrap(); + assert_eq!(k.code, KeyCodeValue::Char('[')); + assert!(k.ctrl); + + let k = SingleKey::parse("?").unwrap(); + assert_eq!(k.code, KeyCodeValue::Char('?')); + + let k = SingleKey::parse("/").unwrap(); + assert_eq!(k.code, KeyCodeValue::Char('/')); + } + + #[test] + fn parse_multi_key_sequence() { + let ki = KeyInput::parse("g g").unwrap(); + match ki { + KeyInput::Sequence(keys) => { + assert_eq!(keys.len(), 2); + assert_eq!(keys[0].code, KeyCodeValue::Char('g')); + assert_eq!(keys[1].code, KeyCodeValue::Char('g')); + } + _ => panic!("expected sequence"), + } + } + + #[test] + fn display_round_trip() { + let cases = ["ctrl-c", "alt-f", "enter", "G", "tab", "pageup"]; + for s in cases { + let k = KeyInput::parse(s).unwrap(); + let display = k.to_string(); + let k2 = KeyInput::parse(&display).unwrap(); + assert_eq!(k, k2, "round-trip failed for {s}"); + } + + let ki = KeyInput::parse("g g").unwrap(); + assert_eq!(ki.to_string(), "g g"); + } + + #[test] + fn from_event_basic() { + let event = KeyEvent::new(KeyCode::Char('c'), KeyModifiers::CONTROL); + let k = SingleKey::from_event(&event).unwrap(); + assert_eq!(k.code, KeyCodeValue::Char('c')); + assert!(k.ctrl); + assert!(!k.alt); + + let event = KeyEvent::new(KeyCode::Enter, KeyModifiers::NONE); + let k = SingleKey::from_event(&event).unwrap(); + assert_eq!(k.code, KeyCodeValue::Enter); + } + + #[test] + fn from_event_uppercase() { + // Crossterm sends uppercase chars with SHIFT modifier + let event = KeyEvent::new(KeyCode::Char('G'), KeyModifiers::SHIFT); + let k = SingleKey::from_event(&event).unwrap(); + assert_eq!(k.code, KeyCodeValue::Char('G')); + // shift flag should be cleared since the case encodes it + assert!(!k.shift); + } + + #[test] + fn from_event_matches_parsed() { + // Verify that from_event and parse produce the same SingleKey + let event = KeyEvent::new(KeyCode::Char('c'), KeyModifiers::CONTROL); + let from_event = SingleKey::from_event(&event).unwrap(); + let parsed = SingleKey::parse("ctrl-c").unwrap(); + assert_eq!(from_event, parsed); + + let event = KeyEvent::new(KeyCode::Char('G'), KeyModifiers::SHIFT); + let from_event = SingleKey::from_event(&event).unwrap(); + let parsed = SingleKey::parse("G").unwrap(); + assert_eq!(from_event, parsed); + } + + #[test] + fn parse_super_modifier() { + let k = SingleKey::parse("super-a").unwrap(); + assert_eq!(k.code, KeyCodeValue::Char('a')); + assert!(k.super_key); + assert!(!k.ctrl && !k.alt && !k.shift); + + // "cmd" is an alias for "super" + let k2 = SingleKey::parse("cmd-a").unwrap(); + assert_eq!(k, k2); + + // "win" is an alias for "super" + let k3 = SingleKey::parse("win-a").unwrap(); + assert_eq!(k, k3); + } + + #[test] + fn parse_super_with_other_modifiers() { + let k = SingleKey::parse("super-ctrl-c").unwrap(); + assert_eq!(k.code, KeyCodeValue::Char('c')); + assert!(k.super_key && k.ctrl); + assert!(!k.alt && !k.shift); + } + + #[test] + fn display_super_modifier() { + let k = SingleKey::parse("super-a").unwrap(); + assert_eq!(k.to_string(), "super-a"); + + let k = SingleKey::parse("super-ctrl-x").unwrap(); + assert_eq!(k.to_string(), "super-ctrl-x"); + } + + #[test] + fn display_round_trip_super() { + let k = KeyInput::parse("super-a").unwrap(); + let display = k.to_string(); + let k2 = KeyInput::parse(&display).unwrap(); + assert_eq!(k, k2, "round-trip failed for super-a"); + } + + #[test] + fn from_event_super() { + let event = KeyEvent::new(KeyCode::Char('a'), KeyModifiers::SUPER); + let k = SingleKey::from_event(&event).unwrap(); + assert_eq!(k.code, KeyCodeValue::Char('a')); + assert!(k.super_key); + assert!(!k.ctrl && !k.alt && !k.shift); + } + + #[test] + fn from_event_super_matches_parsed() { + let event = KeyEvent::new(KeyCode::Char('a'), KeyModifiers::SUPER); + let from_event = SingleKey::from_event(&event).unwrap(); + let parsed = SingleKey::parse("super-a").unwrap(); + assert_eq!(from_event, parsed); + } + + #[test] + fn super_uppercase_preserves_super() { + // super-G should keep the super flag (unlike bare "G" which clears shift) + let k = SingleKey::parse("super-G").unwrap(); + assert_eq!(k.code, KeyCodeValue::Char('G')); + assert!(k.super_key); + } + + #[test] + fn parse_errors() { + assert!(SingleKey::parse("ctrl-alt-shift-xxx").is_err()); + assert!(SingleKey::parse("foobar-a").is_err()); + } + + #[test] + fn parse_function_keys() { + let k = SingleKey::parse("f1").unwrap(); + assert_eq!(k.code, KeyCodeValue::F(1)); + assert!(!k.ctrl && !k.alt && !k.shift); + + let k = SingleKey::parse("F12").unwrap(); + assert_eq!(k.code, KeyCodeValue::F(12)); + + let k = SingleKey::parse("ctrl-f5").unwrap(); + assert_eq!(k.code, KeyCodeValue::F(5)); + assert!(k.ctrl); + + // F24 is valid (some keyboards have extended function keys) + let k = SingleKey::parse("f24").unwrap(); + assert_eq!(k.code, KeyCodeValue::F(24)); + + // F0 and F25+ are invalid + assert!(SingleKey::parse("f0").is_err()); + assert!(SingleKey::parse("f25").is_err()); + } + + #[test] + fn from_event_function_keys() { + let event = KeyEvent::new(KeyCode::F(1), KeyModifiers::NONE); + let k = SingleKey::from_event(&event).unwrap(); + assert_eq!(k.code, KeyCodeValue::F(1)); + + let event = KeyEvent::new(KeyCode::F(12), KeyModifiers::CONTROL); + let k = SingleKey::from_event(&event).unwrap(); + assert_eq!(k.code, KeyCodeValue::F(12)); + assert!(k.ctrl); + } + + #[test] + fn display_function_keys() { + let k = SingleKey::parse("f1").unwrap(); + assert_eq!(k.to_string(), "f1"); + + let k = SingleKey::parse("ctrl-f12").unwrap(); + assert_eq!(k.to_string(), "ctrl-f12"); + } + + #[test] + fn function_key_round_trip() { + let cases = ["f1", "f12", "ctrl-f5", "alt-f10"]; + for s in cases { + let k = KeyInput::parse(s).unwrap(); + let display = k.to_string(); + let k2 = KeyInput::parse(&display).unwrap(); + assert_eq!(k, k2, "round-trip failed for {s}"); + } + } + + #[test] + fn from_event_function_key_matches_parsed() { + let event = KeyEvent::new(KeyCode::F(12), KeyModifiers::NONE); + let from_event = SingleKey::from_event(&event).unwrap(); + let parsed = SingleKey::parse("f12").unwrap(); + assert_eq!(from_event, parsed); + } + + #[test] + fn from_event_backtab_becomes_shift_tab() { + // Many terminals send BackTab for Shift+Tab + let event = KeyEvent::new(KeyCode::BackTab, KeyModifiers::NONE); + let k = SingleKey::from_event(&event).unwrap(); + assert_eq!(k.code, KeyCodeValue::Tab); + assert!(k.shift); + assert!(!k.ctrl && !k.alt); + } + + #[test] + fn from_event_backtab_matches_parsed_shift_tab() { + let event = KeyEvent::new(KeyCode::BackTab, KeyModifiers::NONE); + let from_event = SingleKey::from_event(&event).unwrap(); + let parsed = SingleKey::parse("shift-tab").unwrap(); + assert_eq!(from_event, parsed); + } + + #[test] + fn from_event_backtab_with_ctrl() { + // BackTab with ctrl modifier + let event = KeyEvent::new(KeyCode::BackTab, KeyModifiers::CONTROL); + let k = SingleKey::from_event(&event).unwrap(); + assert_eq!(k.code, KeyCodeValue::Tab); + assert!(k.shift); + assert!(k.ctrl); + } + + #[test] + fn parse_insert_key() { + let k = SingleKey::parse("insert").unwrap(); + assert_eq!(k.code, KeyCodeValue::Insert); + assert!(!k.ctrl && !k.alt && !k.shift); + + let k = SingleKey::parse("ins").unwrap(); + assert_eq!(k.code, KeyCodeValue::Insert); + + let k = SingleKey::parse("ctrl-insert").unwrap(); + assert_eq!(k.code, KeyCodeValue::Insert); + assert!(k.ctrl); + } + + #[test] + fn from_event_insert_key() { + let event = KeyEvent::new(KeyCode::Insert, KeyModifiers::NONE); + let k = SingleKey::from_event(&event).unwrap(); + assert_eq!(k.code, KeyCodeValue::Insert); + } + + #[test] + fn insert_key_round_trip() { + let k = KeyInput::parse("insert").unwrap(); + let display = k.to_string(); + assert_eq!(display, "insert"); + let k2 = KeyInput::parse(&display).unwrap(); + assert_eq!(k, k2); + } +} diff --git a/crates/atuin/src/command/client/search/keybindings/keymap.rs b/crates/atuin/src/command/client/search/keybindings/keymap.rs new file mode 100644 index 00000000..8c7fcfa8 --- /dev/null +++ b/crates/atuin/src/command/client/search/keybindings/keymap.rs @@ -0,0 +1,233 @@ +use std::collections::HashMap; + +use super::actions::Action; +use super::conditions::{ConditionExpr, EvalContext}; +use super::key::{KeyInput, SingleKey}; + +/// A single rule within a keybinding: an optional condition and an action. +/// If the condition is `None`, the rule always matches. +#[derive(Debug, Clone)] +pub struct KeyRule { + pub condition: Option<ConditionExpr>, + pub action: Action, +} + +/// A keybinding is an ordered list of rules. The first rule whose condition +/// matches (or has no condition) wins. +#[derive(Debug, Clone)] +pub struct KeyBinding { + pub rules: Vec<KeyRule>, +} + +/// A keymap is a collection of keybindings indexed by key input. +#[derive(Debug, Clone)] +pub struct Keymap { + pub bindings: HashMap<KeyInput, KeyBinding>, +} + +impl KeyRule { + /// Create an unconditional rule. + pub fn always(action: Action) -> Self { + KeyRule { + condition: None, + action, + } + } + + /// Create a conditional rule. Accepts any type convertible to `ConditionExpr`, + /// including bare `ConditionAtom` values. + pub fn when(condition: impl Into<ConditionExpr>, action: Action) -> Self { + KeyRule { + condition: Some(condition.into()), + action, + } + } +} + +impl KeyBinding { + /// Create a simple (unconditional) binding. + pub fn simple(action: Action) -> Self { + KeyBinding { + rules: vec![KeyRule::always(action)], + } + } + + /// Create a conditional binding from a list of rules. + pub fn conditional(rules: Vec<KeyRule>) -> Self { + KeyBinding { rules } + } +} + +impl Keymap { + /// Create an empty keymap. + pub fn new() -> Self { + Keymap { + bindings: HashMap::new(), + } + } + + /// Bind a key input to a simple (unconditional) action. + pub fn bind(&mut self, key: KeyInput, action: Action) { + self.bindings.insert(key, KeyBinding::simple(action)); + } + + /// Bind a key input to a conditional set of rules. + pub fn bind_conditional(&mut self, key: KeyInput, rules: Vec<KeyRule>) { + self.bindings.insert(key, KeyBinding::conditional(rules)); + } + + /// Resolve a key input to an action given the current evaluation context. + /// Returns `None` if the key has no binding or no rule's condition matches. + pub fn resolve(&self, key: &KeyInput, ctx: &EvalContext) -> Option<Action> { + let binding = self.bindings.get(key)?; + for rule in &binding.rules { + match &rule.condition { + None => return Some(rule.action.clone()), + Some(cond) if cond.evaluate(ctx) => return Some(rule.action.clone()), + Some(_) => {} + } + } + None + } + + /// Check if any binding starts with the given single key as the first key + /// of a multi-key sequence. Used to detect pending multi-key sequences. + pub fn has_sequence_starting_with(&self, prefix: &SingleKey) -> bool { + self.bindings.keys().any(|ki| match ki { + KeyInput::Sequence(keys) => keys.first() == Some(prefix), + KeyInput::Single(_) => false, + }) + } + + /// Merge another keymap into this one. Keys from `other` override keys in `self`. + #[allow(dead_code)] + pub fn merge(&mut self, other: &Keymap) { + for (key, binding) in &other.bindings { + self.bindings.insert(key.clone(), binding.clone()); + } + } +} + +impl Default for Keymap { + fn default() -> Self { + Self::new() + } +} + +#[cfg(test)] +mod tests { + use super::super::conditions::ConditionAtom; + use super::*; + + fn make_ctx(cursor: usize, width: usize, selected: usize, len: usize) -> EvalContext { + EvalContext { + cursor_position: cursor, + input_width: width, + input_byte_len: width, + selected_index: selected, + results_len: len, + original_input_empty: false, + has_context: false, + } + } + + #[test] + fn simple_binding_resolves() { + let mut keymap = Keymap::new(); + let key = KeyInput::parse("ctrl-c").unwrap(); + keymap.bind(key.clone(), Action::ReturnOriginal); + + let ctx = make_ctx(0, 0, 0, 10); + assert_eq!(keymap.resolve(&key, &ctx), Some(Action::ReturnOriginal)); + } + + #[test] + fn conditional_first_match_wins() { + let mut keymap = Keymap::new(); + let key = KeyInput::parse("left").unwrap(); + keymap.bind_conditional( + key.clone(), + vec![ + KeyRule::when(ConditionAtom::CursorAtStart, Action::Exit), + KeyRule::always(Action::CursorLeft), + ], + ); + + // Cursor at start → Exit + let ctx = make_ctx(0, 5, 0, 10); + assert_eq!(keymap.resolve(&key, &ctx), Some(Action::Exit)); + + // Cursor not at start → CursorLeft + let ctx = make_ctx(3, 5, 0, 10); + assert_eq!(keymap.resolve(&key, &ctx), Some(Action::CursorLeft)); + } + + #[test] + fn no_match_returns_none() { + let keymap = Keymap::new(); + let key = KeyInput::parse("ctrl-c").unwrap(); + let ctx = make_ctx(0, 0, 0, 0); + assert_eq!(keymap.resolve(&key, &ctx), None); + } + + #[test] + fn conditional_no_condition_matches_returns_none() { + let mut keymap = Keymap::new(); + let key = KeyInput::parse("left").unwrap(); + // Only one rule with a condition that won't match + keymap.bind_conditional( + key.clone(), + vec![KeyRule::when(ConditionAtom::CursorAtStart, Action::Exit)], + ); + + // Cursor not at start → no match + let ctx = make_ctx(3, 5, 0, 10); + assert_eq!(keymap.resolve(&key, &ctx), None); + } + + #[test] + fn has_sequence_starting_with() { + let mut keymap = Keymap::new(); + let seq = KeyInput::parse("g g").unwrap(); + keymap.bind(seq, Action::ScrollToTop); + + let g = SingleKey::parse("g").unwrap(); + assert!(keymap.has_sequence_starting_with(&g)); + + let h = SingleKey::parse("h").unwrap(); + assert!(!keymap.has_sequence_starting_with(&h)); + } + + #[test] + fn merge_overrides() { + let mut base = Keymap::new(); + let key = KeyInput::parse("ctrl-c").unwrap(); + base.bind(key.clone(), Action::ReturnOriginal); + + let mut overlay = Keymap::new(); + overlay.bind(key.clone(), Action::Exit); + + base.merge(&overlay); + + let ctx = make_ctx(0, 0, 0, 0); + assert_eq!(base.resolve(&key, &ctx), Some(Action::Exit)); + } + + #[test] + fn merge_preserves_unoverridden() { + let mut base = Keymap::new(); + let key1 = KeyInput::parse("ctrl-c").unwrap(); + let key2 = KeyInput::parse("ctrl-d").unwrap(); + base.bind(key1.clone(), Action::ReturnOriginal); + base.bind(key2.clone(), Action::DeleteCharAfter); + + let mut overlay = Keymap::new(); + overlay.bind(key1.clone(), Action::Exit); + + base.merge(&overlay); + + let ctx = make_ctx(0, 0, 0, 0); + assert_eq!(base.resolve(&key1, &ctx), Some(Action::Exit)); + assert_eq!(base.resolve(&key2, &ctx), Some(Action::DeleteCharAfter)); + } +} diff --git a/crates/atuin/src/command/client/search/keybindings/mod.rs b/crates/atuin/src/command/client/search/keybindings/mod.rs new file mode 100644 index 00000000..a9454b0d --- /dev/null +++ b/crates/atuin/src/command/client/search/keybindings/mod.rs @@ -0,0 +1,14 @@ +pub mod actions; +pub mod conditions; +pub mod defaults; +pub mod key; +pub mod keymap; + +pub use actions::Action; +#[allow(unused_imports)] +pub use conditions::{ConditionAtom, ConditionExpr, EvalContext}; +pub use defaults::KeymapSet; +#[allow(unused_imports)] +pub use key::{KeyCodeValue, KeyInput, SingleKey}; +#[allow(unused_imports)] +pub use keymap::{KeyBinding, KeyRule, Keymap}; diff --git a/crates/atuin/src/command/client/setup.rs b/crates/atuin/src/command/client/setup.rs new file mode 100644 index 00000000..acdf0cad --- /dev/null +++ b/crates/atuin/src/command/client/setup.rs @@ -0,0 +1,71 @@ +use atuin_client::settings::Settings; + +use colored::Colorize; +use eyre::Result; +use std::io::{self, Write}; +use toml_edit::{DocumentMut, value}; + +pub async fn run(_settings: &Settings) -> Result<()> { + let enable_ai = prompt( + "Atuin AI", + "This will enable command generation and other AI features via the question mark key", + )?; + + let enable_daemon = prompt( + "Atuin Daemon", + "This will enable improved search and history sync using a persistent background process", + )?; + + let config_file = Settings::get_config_path()?; + let config_str = tokio::fs::read_to_string(&config_file).await?; + let mut doc = config_str.parse::<DocumentMut>()?; + + let mut changed = false; + if enable_ai { + changed = true; + if !doc.contains_key("ai") { + doc["ai"] = toml_edit::table(); + } + doc["ai"]["enabled"] = value(true); + } + + if enable_daemon { + changed = true; + if !doc.contains_key("daemon") { + doc["daemon"] = toml_edit::table(); + } + doc["daemon"]["enabled"] = value(true); + doc["daemon"]["autostart"] = value(true); + doc["search_mode"] = value("daemon-fuzzy"); + } + + if changed { + tokio::fs::write(config_file, doc.to_string()).await?; + + println!( + "{check} Settings updated successfully", + check = "✓".bold().bright_green() + ); + } else { + println!( + "{check} No settings changed", + check = "✓".bold().bright_green() + ); + } + + Ok(()) +} + +pub fn prompt(feature: &str, description: &str) -> Result<bool> { + println!( + "> Enable {feature}?", + feature = feature.bold().bright_blue() + ); + print!(" {description} {q} ", q = "[Y/n]".bold()); + io::stdout().flush().ok(); + + let mut input = String::new(); + io::stdin().read_line(&mut input)?; + let answer = input.trim().to_lowercase(); + Ok(answer.is_empty() || answer == "y" || answer == "yes") +} diff --git a/crates/atuin/src/command/client/stats.rs b/crates/atuin/src/command/client/stats.rs new file mode 100644 index 00000000..a7fc00ac --- /dev/null +++ b/crates/atuin/src/command/client/stats.rs @@ -0,0 +1,85 @@ +use clap::Parser; +use eyre::Result; +use interim::parse_date_string; +use time::{Duration, OffsetDateTime, Time}; + +use atuin_client::{ + database::{Database, current_context}, + settings::Settings, + theme::Theme, +}; + +use atuin_history::stats::{compute, pretty_print}; + +fn parse_ngram_size(s: &str) -> Result<usize, String> { + let value = s + .parse::<usize>() + .map_err(|_| format!("'{s}' is not a valid window size"))?; + + if value == 0 { + return Err("ngram window size must be at least 1".to_string()); + } + + Ok(value) +} + +#[derive(Parser, Debug)] +#[command(infer_subcommands = true)] +pub struct Cmd { + /// Compute statistics for the specified period, leave blank for statistics since the beginning. See [this](https://docs.atuin.sh/reference/stats/) for more details. + period: Vec<String>, + + /// How many top commands to list + #[arg(long, short, default_value = "10")] + count: usize, + + /// The number of consecutive commands to consider + #[arg(long, short, default_value = "1", value_parser = parse_ngram_size)] + ngram_size: usize, +} + +impl Cmd { + pub async fn run(&self, db: &impl Database, settings: &Settings, theme: &Theme) -> Result<()> { + let context = current_context().await?; + let words = if self.period.is_empty() { + String::from("all") + } else { + self.period.join(" ") + }; + + let now = OffsetDateTime::now_utc().to_offset(settings.timezone.0); + let last_night = now.replace_time(Time::MIDNIGHT); + + let history = if words.as_str() == "all" { + db.list(&[], &context, None, false, false).await? + } else if words.trim() == "today" { + let start = last_night; + let end = start + Duration::days(1); + db.range(start, end).await? + } else if words.trim() == "month" { + let end = last_night; + let start = end - Duration::days(31); + db.range(start, end).await? + } else if words.trim() == "week" { + let end = last_night; + let start = end - Duration::days(7); + db.range(start, end).await? + } else if words.trim() == "year" { + let end = last_night; + let start = end - Duration::days(365); + db.range(start, end).await? + } else { + let start = parse_date_string(&words, now, settings.dialect.into())?; + let end = start + Duration::days(1); + db.range(start, end).await? + }; + + let stats = compute(settings, &history, self.count, self.ngram_size); + + if let Some(stats) = stats { + pretty_print(stats, self.ngram_size, theme); + } + + Ok(()) + } +} diff --git a/crates/atuin/src/command/client/store.rs b/crates/atuin/src/command/client/store.rs new file mode 100644 index 00000000..513c404a --- /dev/null +++ b/crates/atuin/src/command/client/store.rs @@ -0,0 +1,120 @@ +use clap::Subcommand; +use eyre::Result; + +use atuin_client::{ + database::Database, + record::{sqlite_store::SqliteStore, store::Store}, + settings::Settings, +}; +use itertools::Itertools; +use time::{OffsetDateTime, UtcOffset}; + +#[cfg(feature = "sync")] +mod push; + +#[cfg(feature = "sync")] +mod pull; + +mod purge; +mod rebuild; +mod rekey; +mod verify; + +#[derive(Subcommand, Debug)] +#[command(infer_subcommands = true)] +pub enum Cmd { + /// Print the current status of the record store + Status, + + /// Rebuild a store (eg atuin store rebuild history) + Rebuild(rebuild::Rebuild), + + /// Re-encrypt the store with a new key (potential for data loss!) + Rekey(rekey::Rekey), + + /// Delete all records in the store that cannot be decrypted with the current key + Purge(purge::Purge), + + /// Verify that all records in the store can be decrypted with the current key + Verify(verify::Verify), + + /// Push all records to the remote sync server (one way sync) + #[cfg(feature = "sync")] + Push(push::Push), + + /// Pull records from the remote sync server (one way sync) + #[cfg(feature = "sync")] + Pull(pull::Pull), +} + +impl Cmd { + pub async fn run( + &self, + settings: &Settings, + database: &dyn Database, + store: SqliteStore, + ) -> Result<()> { + match self { + Self::Status => self.status(store).await, + Self::Rebuild(rebuild) => rebuild.run(settings, store, database).await, + Self::Rekey(rekey) => rekey.run(settings, store).await, + Self::Verify(verify) => verify.run(settings, store).await, + Self::Purge(purge) => purge.run(settings, store).await, + + #[cfg(feature = "sync")] + Self::Push(push) => push.run(settings, store).await, + + #[cfg(feature = "sync")] + Self::Pull(pull) => pull.run(settings, store, database).await, + } + } + + pub async fn status(&self, store: SqliteStore) -> Result<()> { + let host_id = Settings::host_id().await?; + let offset = UtcOffset::current_local_offset().unwrap_or(UtcOffset::UTC); + + let status = store.status().await?; + + // TODO: should probs build some data structure and then pretty-print it or smth + for (host, st) in status.hosts.iter().sorted_by_key(|(h, _)| *h) { + let host_string = if host == &host_id { + format!("host: {} <- CURRENT HOST", host.0.as_hyphenated()) + } else { + format!("host: {}", host.0.as_hyphenated()) + }; + + println!("{host_string}"); + + for (tag, idx) in st.iter().sorted_by_key(|(tag, _)| *tag) { + println!("\tstore: {tag}"); + + let first = store.first(*host, tag).await?; + let last = store.last(*host, tag).await?; + + println!("\t\tidx: {idx}"); + + if let Some(first) = first { + println!("\t\tfirst: {}", first.id.0.as_hyphenated()); + + let time = + OffsetDateTime::from_unix_timestamp_nanos(i128::from(first.timestamp))? + .to_offset(offset); + println!("\t\t\tcreated: {time}"); + } + + if let Some(last) = last { + println!("\t\tlast: {}", last.id.0.as_hyphenated()); + + let time = + OffsetDateTime::from_unix_timestamp_nanos(i128::from(last.timestamp))? + .to_offset(offset); + println!("\t\t\tcreated: {time}"); + } + } + + println!(); + } + + Ok(()) + } +} diff --git a/crates/atuin/src/command/client/store/pull.rs b/crates/atuin/src/command/client/store/pull.rs new file mode 100644 index 00000000..af7ca00e --- /dev/null +++ b/crates/atuin/src/command/client/store/pull.rs @@ -0,0 +1,83 @@ +use clap::Args; +use eyre::Result; + +use atuin_client::{ + database::Database, + record::store::Store, + record::sync::Operation, + record::{sqlite_store::SqliteStore, sync}, + settings::Settings, +}; + +#[derive(Args, Debug)] +pub struct Pull { + /// The tag to push (eg, 'history'). Defaults to all tags + #[arg(long, short)] + pub tag: Option<String>, + + /// Force push records + /// This will first wipe the local store, and then download all records from the remote + #[arg(long, default_value = "false")] + pub force: bool, + + /// Page Size + /// How many records to download at once. Defaults to 100 + #[arg(long, default_value = "100")] + pub page: u64, +} + +impl Pull { + pub async fn run( + &self, + settings: &Settings, + store: SqliteStore, + db: &dyn Database, + ) -> Result<()> { + if self.force { + println!("Forcing local overwrite!"); + println!("Clearing local store"); + + store.delete_all().await?; + } + + // We can actually just use the existing diff/etc to push + // 1. Diff + // 2. Get operations + // 3. Filter operations by + // a) are they a download op? + // b) are they for the host/tag we are pushing here? + let (diff, _) = sync::diff(settings, &store).await?; + let operations = sync::operations(diff, &store).await?; + + let operations = operations + .into_iter() + .filter(|op| match op { + // No noops or downloads thx + Operation::Noop { .. } | Operation::Upload { .. } => false, + + // pull, so yes plz to downloads! + Operation::Download { tag, .. } => { + if self.force { + return true; + } + + if let Some(t) = self.tag.clone() + && t != *tag + { + return false; + } + + true + } + }) + .collect(); + + let (_, downloaded) = sync::sync_remote(operations, &store, settings, self.page).await?; + + println!("Downloaded {} records", downloaded.len()); + + crate::sync::build(settings, &store, db, Some(&downloaded)).await?; + + Ok(()) + } +} diff --git a/crates/atuin/src/command/client/store/purge.rs b/crates/atuin/src/command/client/store/purge.rs new file mode 100644 index 00000000..ad2369ce --- /dev/null +++ b/crates/atuin/src/command/client/store/purge.rs @@ -0,0 +1,26 @@ +use clap::Args; +use eyre::Result; + +use atuin_client::{ + encryption::load_key, + record::{sqlite_store::SqliteStore, store::Store}, + settings::Settings, +}; + +#[derive(Args, Debug)] +pub struct Purge {} + +impl Purge { + pub async fn run(&self, settings: &Settings, store: SqliteStore) -> Result<()> { + println!("Purging local records that cannot be decrypted"); + + let key = load_key(settings)?; + + match store.purge(&key.into()).await { + Ok(()) => println!("Local store purge completed OK"), + Err(e) => println!("Failed to purge local store: {e:?}"), + } + + Ok(()) + } +} diff --git a/crates/atuin/src/command/client/store/push.rs b/crates/atuin/src/command/client/store/push.rs new file mode 100644 index 00000000..23958d7f --- /dev/null +++ b/crates/atuin/src/command/client/store/push.rs @@ -0,0 +1,101 @@ +use atuin_common::record::HostId; +use clap::Args; +use eyre::Result; +use uuid::Uuid; + +use atuin_client::{ + api_client::Client, + record::sync::Operation, + record::{sqlite_store::SqliteStore, sync}, + settings::Settings, +}; + +#[derive(Args, Debug)] +pub struct Push { + /// The tag to push (eg, 'history'). Defaults to all tags + #[arg(long, short)] + pub tag: Option<String>, + + /// The host to push, in the form of a UUID host ID. Defaults to the current host. + #[arg(long)] + pub host: Option<Uuid>, + + /// Force push records + /// This will override both host and tag, to be all hosts and all tags. First clear the remote store, then upload all of the + /// local store + #[arg(long, default_value = "false")] + pub force: bool, + + /// Page Size + /// How many records to upload at once. Defaults to 100 + #[arg(long, default_value = "100")] + pub page: u64, +} + +impl Push { + pub async fn run(&self, settings: &Settings, store: SqliteStore) -> Result<()> { + let host_id = Settings::host_id().await?; + + if self.force { + println!("Forcing remote store overwrite!"); + println!("Clearing remote store"); + + let client = Client::new( + &settings.sync_address, + settings.sync_auth_token().await?, + settings.network_connect_timeout, + settings.network_timeout * 10, // we may be deleting a lot of data... so up the + // timeout + ) + .expect("failed to create client"); + + client.delete_store().await?; + } + + // We can actually just use the existing diff/etc to push + // 1. Diff + // 2. Get operations + // 3. Filter operations by + // a) are they an upload op? + // b) are they for the host/tag we are pushing here? + let (diff, _) = sync::diff(settings, &store).await?; + let operations = sync::operations(diff, &store).await?; + + let operations = operations + .into_iter() + .filter(|op| match op { + // No noops or downloads thx + Operation::Noop { .. } | Operation::Download { .. } => false, + + // push, so yes plz to uploads! + Operation::Upload { host, tag, .. } => { + if self.force { + return true; + } + + if let Some(h) = self.host { + if HostId(h) != *host { + return false; + } + } else if *host != host_id { + return false; + } + + if let Some(t) = self.tag.clone() + && t != *tag + { + return false; + } + + true + } + }) + .collect(); + + let (uploaded, _) = sync::sync_remote(operations, &store, settings, self.page).await?; + + println!("Uploaded {uploaded} records"); + + Ok(()) + } +} diff --git a/crates/atuin/src/command/client/store/rebuild.rs b/crates/atuin/src/command/client/store/rebuild.rs new file mode 100644 index 00000000..a98f8142 --- /dev/null +++ b/crates/atuin/src/command/client/store/rebuild.rs @@ -0,0 +1,94 @@ +use atuin_dotfiles::store::{AliasStore, var::VarStore}; +use atuin_scripts::store::ScriptStore; +use clap::Args; +use eyre::{Result, bail}; + +#[cfg(feature = "daemon")] +use atuin_daemon::emit_event; + +use atuin_client::{ + database::Database, encryption, history::store::HistoryStore, + record::sqlite_store::SqliteStore, settings::Settings, +}; + +#[derive(Args, Debug)] +pub struct Rebuild { + pub tag: String, +} + +impl Rebuild { + pub async fn run( + &self, + settings: &Settings, + store: SqliteStore, + database: &dyn Database, + ) -> Result<()> { + // keep it as a string and not an enum atm + // would be super cool to build this dynamically in the future + // eg register handles for rebuilding various tags without having to make this part of the + // binary big + match self.tag.as_str() { + "history" => { + self.rebuild_history(settings, store.clone(), database) + .await?; + } + + "dotfiles" => { + self.rebuild_dotfiles(settings, store.clone()).await?; + } + + "scripts" => { + self.rebuild_scripts(settings, store.clone()).await?; + } + + tag => bail!("unknown tag: {tag}"), + } + + Ok(()) + } + + async fn rebuild_history( + &self, + settings: &Settings, + store: SqliteStore, + database: &dyn Database, + ) -> Result<()> { + let encryption_key: [u8; 32] = encryption::load_key(settings)?.into(); + + let host_id = Settings::host_id().await?; + let history_store = HistoryStore::new(store, host_id, encryption_key); + + history_store.build(database).await?; + + #[cfg(feature = "daemon")] + let _ = emit_event(atuin_daemon::DaemonEvent::HistoryRebuilt).await; + + Ok(()) + } + + async fn rebuild_dotfiles(&self, settings: &Settings, store: SqliteStore) -> Result<()> { + let encryption_key: [u8; 32] = encryption::load_key(settings)?.into(); + + let host_id = Settings::host_id().await?; + + let alias_store = AliasStore::new(store.clone(), host_id, encryption_key); + let var_store = VarStore::new(store.clone(), host_id, encryption_key); + + alias_store.build().await?; + var_store.build().await?; + + Ok(()) + } + + async fn rebuild_scripts(&self, settings: &Settings, store: SqliteStore) -> Result<()> { + let encryption_key: [u8; 32] = encryption::load_key(settings)?.into(); + let host_id = Settings::host_id().await?; + let script_store = ScriptStore::new(store, host_id, encryption_key); + let database = + atuin_scripts::database::Database::new(settings.scripts.db_path.clone(), 1.0).await?; + + script_store.build(database).await?; + + Ok(()) + } +} diff --git a/crates/atuin/src/command/client/store/rekey.rs b/crates/atuin/src/command/client/store/rekey.rs new file mode 100644 index 00000000..4c537a48 --- /dev/null +++ b/crates/atuin/src/command/client/store/rekey.rs @@ -0,0 +1,57 @@ +use clap::Args; +use eyre::{Result, bail}; +use tokio::{fs::File, io::AsyncWriteExt}; + +use atuin_client::{ + encryption::{Key, decode_key, encode_key, generate_encoded_key, load_key}, + record::sqlite_store::SqliteStore, + record::store::Store, + settings::Settings, +}; + +#[derive(Args, Debug)] +pub struct Rekey { + /// The new key to use for encryption. Omit for a randomly-generated key + key: Option<String>, +} + +impl Rekey { + pub async fn run(&self, settings: &Settings, store: SqliteStore) -> Result<()> { + let key = if let Some(key) = self.key.clone() { + println!("Re-encrypting store with specified key"); + + match bip39::Mnemonic::from_phrase(&key, bip39::Language::English) { + Ok(mnemonic) => encode_key(Key::from_slice(mnemonic.entropy()))?, + Err(err) => { + match err { + // assume they copied in the base64 key + bip39::ErrorKind::InvalidWord(_) => key, + bip39::ErrorKind::InvalidChecksum => { + bail!("key mnemonic was not valid") + } + bip39::ErrorKind::InvalidKeysize(_) + | bip39::ErrorKind::InvalidWordLength(_) + | bip39::ErrorKind::InvalidEntropyLength(_, _) => { + bail!("key was not the correct length") + } + } + } + } + } else { + println!("Re-encrypting store with freshly-generated key"); + let (_, encoded) = generate_encoded_key()?; + encoded + }; + + let current_key: [u8; 32] = load_key(settings)?.into(); + let new_key: [u8; 32] = decode_key(key.clone())?.into(); + + store.re_encrypt(¤t_key, &new_key).await?; + + println!("Store rewritten. Saving new key"); + let mut file = File::create(settings.key_path.clone()).await?; + file.write_all(key.as_bytes()).await?; + + Ok(()) + } +} diff --git a/crates/atuin/src/command/client/store/verify.rs b/crates/atuin/src/command/client/store/verify.rs new file mode 100644 index 00000000..84bec96a --- /dev/null +++ b/crates/atuin/src/command/client/store/verify.rs @@ -0,0 +1,26 @@ +use clap::Args; +use eyre::Result; + +use atuin_client::{ + encryption::load_key, + record::{sqlite_store::SqliteStore, store::Store}, + settings::Settings, +}; + +#[derive(Args, Debug)] +pub struct Verify {} + +impl Verify { + pub async fn run(&self, settings: &Settings, store: SqliteStore) -> Result<()> { + println!("Verifying local store can be decrypted with the current key"); + + let key = load_key(settings)?; + + match store.verify(&key.into()).await { + Ok(()) => println!("Local store encryption verified OK"), + Err(e) => println!("Failed to verify local store encryption: {e:?}"), + } + + Ok(()) + } +} diff --git a/crates/atuin/src/command/client/sync.rs b/crates/atuin/src/command/client/sync.rs new file mode 100644 index 00000000..250e98aa --- /dev/null +++ b/crates/atuin/src/command/client/sync.rs @@ -0,0 +1,131 @@ +use clap::Subcommand; +use eyre::{Result, WrapErr}; + +use atuin_client::{ + database::Database, + encryption, + history::store::HistoryStore, + record::{sqlite_store::SqliteStore, store::Store, sync}, + settings::Settings, +}; + +mod status; + +use crate::command::client::account; + +#[derive(Subcommand, Debug)] +#[command(infer_subcommands = true)] +pub enum Cmd { + /// Sync with the configured server + Sync { + /// Force re-download everything + #[arg(long, short)] + force: bool, + }, + + /// Login to the configured server + Login(account::login::Cmd), + + /// Log out + Logout, + + /// Register with the configured server + Register(account::register::Cmd), + + /// Print the encryption key for transfer to another machine + Key { + /// Switch to base64 output of the key + #[arg(long)] + base64: bool, + }, + + /// Display the sync status + Status, +} + +impl Cmd { + pub async fn run( + self, + settings: Settings, + db: &impl Database, + store: SqliteStore, + ) -> Result<()> { + match self { + Self::Sync { force } => run(&settings, force, db, store).await, + Self::Login(l) => l.run(&settings, &store).await, + Self::Logout => account::logout::run().await, + Self::Register(r) => r.run(&settings, &store).await, + Self::Status => status::run(&settings, db).await, + Self::Key { base64 } => { + use atuin_client::encryption::{encode_key, load_key}; + let key = load_key(&settings).wrap_err("could not load encryption key")?; + + if base64 { + let encode = encode_key(&key).wrap_err("could not encode encryption key")?; + println!("{encode}"); + } else { + let mnemonic = bip39::Mnemonic::from_entropy(&key, bip39::Language::English) + .map_err(|_| eyre::eyre!("invalid key"))?; + println!("{mnemonic}"); + } + Ok(()) + } + } + } +} + +async fn run( + settings: &Settings, + force: bool, + db: &impl Database, + store: SqliteStore, +) -> Result<()> { + if settings.sync.records { + let encryption_key: [u8; 32] = encryption::load_key(settings) + .context("could not load encryption key")? + .into(); + + let host_id = Settings::host_id().await?; + let history_store = HistoryStore::new(store.clone(), host_id, encryption_key); + + let (uploaded, downloaded) = sync::sync(settings, &store).await?; + + crate::sync::build(settings, &store, db, Some(&downloaded)).await?; + + println!("{uploaded}/{} up/down to record store", downloaded.len()); + + let history_length = db.history_count(true).await?; + let store_history_length = store.len_tag("history").await?; + + #[allow(clippy::cast_sign_loss)] + if history_length as u64 > store_history_length { + println!( + "{history_length} in history index, but {store_history_length} in history store" + ); + println!("Running automatic history store init..."); + + // Internally we use the global filter mode, so this context is ignored. + // don't recurse or loop here. + history_store.init_store(db).await?; + + println!("Re-running sync due to new records locally"); + + // we'll want to run sync once more, as there will now be stuff to upload + let (uploaded, downloaded) = sync::sync(settings, &store).await?; + + crate::sync::build(settings, &store, db, Some(&downloaded)).await?; + + println!("{uploaded}/{} up/down to record store", downloaded.len()); + } + } else { + atuin_client::sync::sync(settings, force, db).await?; + } + + println!( + "Sync complete! {} items in history database, force: {}", + db.history_count(true).await?, + force + ); + + Ok(()) +} diff --git a/crates/atuin/src/command/client/sync/status.rs b/crates/atuin/src/command/client/sync/status.rs new file mode 100644 index 00000000..54911cc8 --- /dev/null +++ b/crates/atuin/src/command/client/sync/status.rs @@ -0,0 +1,45 @@ +use crate::{SHA, VERSION}; +use atuin_client::{api_client, database::Database, settings::Settings}; +use colored::Colorize; +use eyre::{Result, bail}; + +pub async fn run(settings: &Settings, db: &impl Database) -> Result<()> { + if !settings.logged_in().await? { + bail!("You are not logged in to a sync server - cannot show sync status"); + } + + let client = api_client::Client::new( + &settings.sync_address, + settings.sync_auth_token().await?, + settings.network_connect_timeout, + settings.network_timeout, + )?; + + let me = client.me().await?; + let last_sync = Settings::last_sync().await?; + + println!("Atuin v{VERSION} - Build rev {SHA}\n"); + + println!("{}", "[Local]".green()); + + if settings.auto_sync { + println!("Sync frequency: {}", settings.sync_frequency); + println!("Last sync: {}", last_sync.to_offset(settings.timezone.0)); + } + + if !settings.sync.records { + let local_count = db.history_count(false).await?; + let deleted_count = db.history_count(true).await? - local_count; + + println!("History count: {local_count}"); + println!("Deleted history count: {deleted_count}\n"); + } + + if settings.auto_sync { + println!("{}", "[Remote]".green()); + println!("Address: {}", settings.sync_address); + println!("Username: {}", me.username); + } + + Ok(()) +} diff --git a/crates/atuin/src/command/client/wrapped.rs b/crates/atuin/src/command/client/wrapped.rs new file mode 100644 index 00000000..20b79a2e --- /dev/null +++ b/crates/atuin/src/command/client/wrapped.rs @@ -0,0 +1,368 @@ +use crossterm::style::{ResetColor, SetAttribute}; +use eyre::Result; +use std::collections::{HashMap, HashSet}; +use time::{Date, Duration, Month, OffsetDateTime, Time}; + +use atuin_client::{ + database::Database, encryption, record::sqlite_store::SqliteStore, settings::Settings, + theme::Theme, +}; +use atuin_dotfiles::store::AliasStore; + +use atuin_history::stats::{Stats, compute}; + +#[derive(Debug)] +struct WrappedStats { + nav_commands: usize, + pkg_commands: usize, + error_rate: f64, + first_half_commands: Vec<(String, usize)>, + second_half_commands: Vec<(String, usize)>, + git_percentage: f64, + busiest_hour: Option<(String, usize)>, +} + +impl WrappedStats { + #[allow(clippy::too_many_lines, clippy::cast_precision_loss)] + fn new( + settings: &Settings, + stats: &Stats, + history: &[atuin_client::history::History], + alias_map: &HashMap<String, String>, + ) -> Self { + // Helper to expand alias to its first command word + let expand_alias = |cmd: &str| -> String { + alias_map.get(cmd).map_or_else( + || cmd.to_string(), + |expanded| { + expanded + .split_whitespace() + .next() + .unwrap_or(cmd) + .to_string() + }, + ) + }; + + let nav_commands = stats + .top + .iter() + .filter(|(cmd, _)| { + let cmd = &cmd[0]; + cmd == "cd" || cmd == "ls" || cmd == "pwd" || cmd == "pushd" || cmd == "popd" + }) + .map(|(_, count)| count) + .sum(); + + let pkg_managers = [ + "cargo", + "npm", + "pnpm", + "yarn", + "pip", + "pip3", + "pipenv", + "poetry", + "pipx", + "uv", + "brew", + "apt", + "apt-get", + "apk", + "pacman", + "yay", + "paru", + "yum", + "dnf", + "dnf5", + "rpm", + "rpm-ostree", + "zypper", + "pkg", + "chocolatey", + "choco", + "scoop", + "winget", + "gem", + "bundle", + "shards", + "composer", + "gradle", + "maven", + "mvn", + "go get", + "nuget", + "dotnet", + "mix", + "hex", + "rebar3", + "nix", + "nix-env", + "cabal", + "opam", + ]; + + let pkg_commands = history + .iter() + .filter(|h| { + let cmd = h.command.clone(); + pkg_managers.iter().any(|pm| cmd.starts_with(pm)) + }) + .count(); + + // Error analysis + let mut command_errors: HashMap<String, (usize, usize)> = HashMap::new(); // (total_uses, errors) + let midyear = history[0].timestamp + Duration::days(182); // Split year in half + + let mut first_half_commands: HashMap<String, usize> = HashMap::new(); + let mut second_half_commands: HashMap<String, usize> = HashMap::new(); + let mut hours: HashMap<String, usize> = HashMap::new(); + + for entry in history { + let raw_cmd = entry + .command + .split_whitespace() + .next() + .unwrap_or("") + .to_string(); + let cmd = expand_alias(&raw_cmd); + let (total, errors) = command_errors.entry(cmd.clone()).or_insert((0, 0)); + *total += 1; + if entry.exit != 0 { + *errors += 1; + } + + // Track command evolution + if entry.timestamp < midyear { + *first_half_commands.entry(cmd.clone()).or_default() += 1; + } else { + *second_half_commands.entry(cmd).or_default() += 1; + } + + // Track hourly distribution + let local_time = entry + .timestamp + .to_offset(time::UtcOffset::current_local_offset().unwrap_or(settings.timezone.0)); + let hour = format!("{:02}:00", local_time.time().hour()); + *hours.entry(hour).or_default() += 1; + } + + let total_errors: usize = command_errors.values().map(|(_, errors)| errors).sum(); + let total_commands: usize = command_errors.values().map(|(total, _)| total).sum(); + let error_rate = total_errors as f64 / total_commands as f64; + + // Process command evolution data + let mut first_half: Vec<_> = first_half_commands.into_iter().collect(); + let mut second_half: Vec<_> = second_half_commands.into_iter().collect(); + first_half.sort_by_key(|(_, count)| std::cmp::Reverse(*count)); + second_half.sort_by_key(|(_, count)| std::cmp::Reverse(*count)); + first_half.truncate(5); + second_half.truncate(5); + + // Calculate git percentage + let git_commands: usize = stats + .top + .iter() + .filter(|(cmd, _)| cmd[0].starts_with("git")) + .map(|(_, count)| count) + .sum(); + let git_percentage = git_commands as f64 / stats.total_commands as f64; + + // Find busiest hour + let busiest_hour = hours.into_iter().max_by_key(|(_, count)| *count); + + Self { + nav_commands, + pkg_commands, + error_rate, + first_half_commands: first_half, + second_half_commands: second_half, + git_percentage, + busiest_hour, + } + } +} + +pub fn print_wrapped_header(year: i32) { + let reset = ResetColor; + let bold = SetAttribute(crossterm::style::Attribute::Bold); + + println!("{bold}╭────────────────────────────────────╮{reset}"); + println!("{bold}│ ATUIN WRAPPED {year} │{reset}"); + println!("{bold}│ Your Year in Shell History │{reset}"); + println!("{bold}╰────────────────────────────────────╯{reset}"); + println!(); +} + +#[allow(clippy::cast_precision_loss)] +fn print_fun_facts(wrapped_stats: &WrappedStats, stats: &Stats, year: i32) { + let reset = ResetColor; + let bold = SetAttribute(crossterm::style::Attribute::Bold); + + if wrapped_stats.git_percentage > 0.05 { + println!( + "{bold}🌟 You're a Git Power User!{reset} {bold}{:.1}%{reset} of your commands were Git operations\n", + wrapped_stats.git_percentage * 100.0 + ); + } + // Navigation patterns + let nav_percentage = wrapped_stats.nav_commands as f64 / stats.total_commands as f64 * 100.0; + if nav_percentage > 0.05 { + println!( + "{bold}🚀 You're a Navigator!{reset} {bold}{nav_percentage:.1}%{reset} of your time was spent navigating directories\n", + ); + } + + // Command vocabulary + println!( + "{bold}📚 Command Vocabulary{reset}: You know {bold}{}{reset} unique commands\n", + stats.unique_commands + ); + + // Package management + println!( + "{bold}📦 Package Management{reset}: You ran {bold}{}{reset} package-related commands\n", + wrapped_stats.pkg_commands + ); + + // Error patterns + let error_percentage = wrapped_stats.error_rate * 100.0; + println!( + "{bold}🚨 Error Analysis{reset}: Your commands failed {bold}{error_percentage:.1}%{reset} of the time\n", + ); + + // Command evolution + println!("🔍 Command Evolution:"); + + // print stats for each half and compare + println!(" {bold}Top Commands{reset} in the first half of {year}:"); + for (cmd, count) in wrapped_stats.first_half_commands.iter().take(3) { + println!(" {bold}{cmd}{reset} ({count} times)"); + } + + println!(" {bold}Top Commands{reset} in the second half of {year}:"); + for (cmd, count) in wrapped_stats.second_half_commands.iter().take(3) { + println!(" {bold}{cmd}{reset} ({count} times)"); + } + + // Find new favorite commands (in top 5 of second half but not in first half) + let first_half_set: HashSet<_> = wrapped_stats + .first_half_commands + .iter() + .map(|(cmd, _)| cmd) + .collect(); + let new_favorites: Vec<_> = wrapped_stats + .second_half_commands + .iter() + .filter(|(cmd, _)| !first_half_set.contains(cmd)) + .take(2) + .collect(); + + if !new_favorites.is_empty() { + println!(" {bold}New favorites{reset} in the second half:"); + for (cmd, count) in new_favorites { + println!(" {bold}{cmd}{reset} ({count} times)"); + } + } + + // Time patterns + if let Some((hour, count)) = &wrapped_stats.busiest_hour { + println!("\n🕘 Most Productive Hour: {bold}{hour}{reset} ({count} commands)",); + + // Night owl or early bird + let hour_num = hour + .split(':') + .next() + .unwrap_or("0") + .parse::<u32>() + .unwrap_or(0); + if hour_num >= 22 || hour_num <= 4 { + println!(" You're quite the night owl! 🦉"); + } else if (5..=7).contains(&hour_num) { + println!(" Early bird gets the worm! 🐦"); + } + } + + println!(); +} + +pub async fn run( + year: Option<i32>, + db: &impl Database, + settings: &Settings, + store: SqliteStore, + theme: &Theme, +) -> Result<()> { + let now = OffsetDateTime::now_utc().to_offset(settings.timezone.0); + let month = now.month(); + + // If we're in December, then wrapped is for the current year. If not, it's for the previous year + let year = year.unwrap_or_else(|| { + if month == Month::December { + now.year() + } else { + now.year() - 1 + } + }); + + let start = OffsetDateTime::new_in_offset( + Date::from_calendar_date(year, Month::January, 1).unwrap(), + Time::MIDNIGHT, + now.offset(), + ); + let end = OffsetDateTime::new_in_offset( + Date::from_calendar_date(year, Month::December, 31).unwrap(), + Time::MIDNIGHT + Duration::days(1) - Duration::nanoseconds(1), + now.offset(), + ); + + let history = db.range(start, end).await?; + if history.is_empty() { + println!( + "Your history for {year} is empty!\nMaybe 'atuin import' could help you import your previous history 🪄" + ); + return Ok(()); + } + + // Load aliases for expansion + let alias_map: HashMap<String, String> = if settings.dotfiles.enabled { + if let Ok(encryption_key) = encryption::load_key(settings) { + let encryption_key: [u8; 32] = encryption_key.into(); + let host_id = Settings::host_id().await?; + let alias_store = AliasStore::new(store, host_id, encryption_key); + + alias_store + .aliases() + .await + .unwrap_or_default() + .into_iter() + .map(|a| (a.name, a.value)) + .collect() + } else { + HashMap::new() + } + } else { + HashMap::new() + }; + + // Compute overall stats using existing functionality + let stats = compute(settings, &history, 10, 1).expect("Failed to compute stats"); + let wrapped_stats = WrappedStats::new(settings, &stats, &history, &alias_map); + + // Print wrapped format + print_wrapped_header(year); + + println!("🎉 In {year}, you typed {} commands!", stats.total_commands); + println!( + " That's ~{} commands every day\n", + stats.total_commands / 365 + ); + + println!("Your Top Commands:"); + atuin_history::stats::pretty_print(stats.clone(), 1, theme); + println!(); + + print_fun_facts(&wrapped_stats, &stats, year); + + Ok(()) +} diff --git a/crates/atuin/src/command/contributors.rs b/crates/atuin/src/command/contributors.rs new file mode 100644 index 00000000..452fd335 --- /dev/null +++ b/crates/atuin/src/command/contributors.rs @@ -0,0 +1,5 @@ +static CONTRIBUTORS: &str = include_str!("CONTRIBUTORS"); + +pub fn run() { + println!("\n{CONTRIBUTORS}"); +} diff --git a/crates/atuin/src/command/external.rs b/crates/atuin/src/command/external.rs new file mode 100644 index 00000000..657aea56 --- /dev/null +++ b/crates/atuin/src/command/external.rs @@ -0,0 +1,92 @@ +use std::fmt::Write as _; +use std::process::Command; +use std::{io, process}; + +#[cfg(feature = "client")] +use atuin_client::plugin::OfficialPluginRegistry; +use clap::CommandFactory; +use clap::builder::{StyledStr, Styles}; +use eyre::Result; + +use crate::Atuin; + +pub fn run(args: &[String]) -> Result<()> { + let subcommand = &args[0]; + let bin = format!("atuin-{subcommand}"); + let mut cmd = Command::new(&bin); + cmd.args(&args[1..]); + + let spawn_result = match cmd.spawn() { + Ok(child) => Ok(child), + Err(e) => match e.kind() { + io::ErrorKind::NotFound => { + let output = render_not_found(subcommand, &bin); + Err(output) + } + _ => Err(e.to_string().into()), + }, + }; + + match spawn_result { + Ok(mut child) => { + let status = child.wait()?; + if status.success() { + Ok(()) + } else { + process::exit(status.code().unwrap_or(1)); + } + } + Err(e) => { + eprintln!("{}", e.ansi()); + process::exit(1); + } + } +} + +fn render_not_found(subcommand: &str, bin: &str) -> StyledStr { + let mut output = StyledStr::new(); + let styles = Styles::styled(); + + let error = styles.get_error(); + let invalid = styles.get_invalid(); + let literal = styles.get_literal(); + + #[cfg(feature = "client")] + { + let registry = OfficialPluginRegistry::new(); + + // Check if this is an official plugin + if let Some(install_message) = registry.get_install_message(subcommand) { + let _ = write!(output, "{error}error:{error:#} "); + let _ = write!( + output, + "'{invalid}{subcommand}{invalid:#}' is an official atuin plugin, but it's not installed" + ); + let _ = write!(output, "\n\n"); + let _ = write!(output, "{install_message}"); + return output; + } + } + + let mut atuin_cmd = Atuin::command(); + let usage = atuin_cmd.render_usage(); + + let _ = write!(output, "{error}error:{error:#} "); + let _ = write!( + output, + "unrecognized subcommand '{invalid}{subcommand}{invalid:#}' " + ); + let _ = write!( + output, + "and no executable named '{invalid}{bin}{invalid:#}' found in your PATH" + ); + let _ = write!(output, "\n\n"); + let _ = write!(output, "{usage}"); + let _ = write!(output, "\n\n"); + let _ = write!( + output, + "For more information, try '{literal}--help{literal:#}'." + ); + + output +} diff --git a/crates/atuin/src/command/gen_completions.rs b/crates/atuin/src/command/gen_completions.rs new file mode 100644 index 00000000..10d4f689 --- /dev/null +++ b/crates/atuin/src/command/gen_completions.rs @@ -0,0 +1,84 @@ +use clap::{CommandFactory, Parser, ValueEnum}; +use clap_complete::{Generator, Shell, generate, generate_to}; +use clap_complete_nushell::Nushell; +use eyre::Result; + +// clap put nushell completions into a separate package due to the maintainers +// being a little less committed to support them. +// This means we have to do a tiny bit of legwork to combine these completions +// into one command. +#[derive(Debug, Clone, ValueEnum)] +#[value(rename_all = "lower")] +pub enum GenShell { + Bash, + Elvish, + Fish, + Nushell, + PowerShell, + Zsh, +} + +impl Generator for GenShell { + fn file_name(&self, name: &str) -> String { + match self { + // clap_complete + Self::Bash => Shell::Bash.file_name(name), + Self::Elvish => Shell::Elvish.file_name(name), + Self::Fish => Shell::Fish.file_name(name), + Self::PowerShell => Shell::PowerShell.file_name(name), + Self::Zsh => Shell::Zsh.file_name(name), + + // clap_complete_nushell + Self::Nushell => Nushell.file_name(name), + } + } + + fn generate(&self, cmd: &clap::Command, buf: &mut dyn std::io::prelude::Write) { + match self { + // clap_complete + Self::Bash => Shell::Bash.generate(cmd, buf), + Self::Elvish => Shell::Elvish.generate(cmd, buf), + Self::Fish => Shell::Fish.generate(cmd, buf), + Self::PowerShell => Shell::PowerShell.generate(cmd, buf), + Self::Zsh => Shell::Zsh.generate(cmd, buf), + + // clap_complete_nushell + Self::Nushell => Nushell.generate(cmd, buf), + } + } +} + +#[derive(Debug, Parser)] +pub struct Cmd { + /// Set the shell for generating completions + #[arg(long, short)] + shell: GenShell, + + /// Set the output directory + #[arg(long, short)] + out_dir: Option<String>, +} + +impl Cmd { + pub fn run(self) -> Result<()> { + let Cmd { shell, out_dir } = self; + + let mut cli = crate::Atuin::command(); + + match out_dir { + Some(out_dir) => { + generate_to(shell, &mut cli, env!("CARGO_PKG_NAME"), &out_dir)?; + } + None => { + generate( + shell, + &mut cli, + env!("CARGO_PKG_NAME"), + &mut std::io::stdout(), + ); + } + } + + Ok(()) + } +} diff --git a/crates/atuin/src/command/mod.rs b/crates/atuin/src/command/mod.rs new file mode 100644 index 00000000..7896628d --- /dev/null +++ b/crates/atuin/src/command/mod.rs @@ -0,0 +1,75 @@ +use clap::Subcommand; +use eyre::Result; + +#[cfg(not(windows))] +use rustix::{fs::Mode, process::umask}; + +#[cfg(feature = "client")] +mod client; + +mod contributors; + +mod gen_completions; + +mod external; + +#[derive(Subcommand)] +#[command(infer_subcommands = true)] +#[allow(clippy::large_enum_variant)] +pub enum AtuinCmd { + #[cfg(feature = "client")] + #[command(flatten)] + Client(client::Cmd), + + /// Terminal emulator for atuin + #[cfg(feature = "hex")] + Hex { + #[command(subcommand)] + cmd: Option<atuin_hex::Cmd>, + }, + + /// Generate a UUID + Uuid, + + Contributors, + + /// Generate shell completions + GenCompletions(gen_completions::Cmd), + + #[command(external_subcommand)] + External(Vec<String>), +} + +impl AtuinCmd { + pub fn run(self) -> Result<()> { + #[cfg(not(windows))] + { + // set umask before we potentially open/create files + // or in other words, 077. Do not allow any access to any other user + let mode = Mode::RWXG | Mode::RWXO; + umask(mode); + } + + match self { + #[cfg(feature = "client")] + Self::Client(client) => client.run(), + + #[cfg(feature = "hex")] + Self::Hex { cmd } => { + atuin_hex::run(cmd); + Ok(()) + } + + Self::Contributors => { + contributors::run(); + Ok(()) + } + Self::Uuid => { + println!("{}", atuin_common::utils::uuid_v7().as_simple()); + Ok(()) + } + Self::GenCompletions(gen_completions) => gen_completions.run(), + Self::External(args) => external::run(&args), + } + } +} diff --git a/crates/atuin/src/main.rs b/crates/atuin/src/main.rs new file mode 100644 index 00000000..1a45988a --- /dev/null +++ b/crates/atuin/src/main.rs @@ -0,0 +1,59 @@ +#![warn(clippy::pedantic, clippy::nursery)] +#![allow(clippy::use_self, clippy::missing_const_for_fn)] // not 100% reliable + +use clap::Parser; +use clap::builder::Styles; +use clap::builder::styling::{AnsiColor, Effects}; +use eyre::Result; + +use command::AtuinCmd; + +mod command; + +#[cfg(feature = "sync")] +mod sync; + +const VERSION: &str = env!("CARGO_PKG_VERSION"); +const SHA: &str = env!("GIT_HASH"); + +const LONG_VERSION: &str = concat!(env!("CARGO_PKG_VERSION"), " (", env!("GIT_HASH"), ")"); + +static HELP_TEMPLATE: &str = "\ +{before-help}{name} {version} +{author} +{about} + +{usage-heading} + {usage} + +{all-args}{after-help}"; + +const STYLES: Styles = Styles::styled() + .header(AnsiColor::Yellow.on_default().effects(Effects::BOLD)) + .usage(AnsiColor::Green.on_default().effects(Effects::BOLD)) + .literal(AnsiColor::Green.on_default().effects(Effects::BOLD)) + .placeholder(AnsiColor::Green.on_default()); + +/// Magical shell history +#[derive(Parser)] +#[command( + author = "Ellie Huxtable <ellie@atuin.sh>", + version = VERSION, + long_version = LONG_VERSION, + help_template(HELP_TEMPLATE), + styles = STYLES, +)] +struct Atuin { + #[command(subcommand)] + atuin: AtuinCmd, +} + +impl Atuin { + fn run(self) -> Result<()> { + self.atuin.run() + } +} + +fn main() -> Result<()> { + Atuin::parse().run() +} diff --git a/crates/atuin/src/shell/.gitattributes b/crates/atuin/src/shell/.gitattributes new file mode 100644 index 00000000..fae8897c --- /dev/null +++ b/crates/atuin/src/shell/.gitattributes @@ -0,0 +1 @@ +* eol=lf diff --git a/crates/atuin/src/shell/atuin.bash b/crates/atuin/src/shell/atuin.bash new file mode 100644 index 00000000..f72bfcfc --- /dev/null +++ b/crates/atuin/src/shell/atuin.bash @@ -0,0 +1,687 @@ +# Include guard +if [[ ${__atuin_initialized-} == true ]]; then + false +elif [[ $- != *i* ]]; then + # Enable only in interactive shells + false +elif ((BASH_VERSINFO[0] < 3 || BASH_VERSINFO[0] == 3 && BASH_VERSINFO[1] < 1)); then + # Require bash >= 3.1 + [[ -t 2 ]] && printf 'atuin: requires bash >= 3.1 for the integration.\n' >&2 + false +else # (include guard) beginning of main content +#------------------------------------------------------------------------------ +__atuin_initialized=true + +if [[ -z "${ATUIN_SESSION:-}" || "${ATUIN_SHLVL:-}" != "$SHLVL" ]]; then + ATUIN_SESSION=$(atuin uuid) + export ATUIN_SESSION + export ATUIN_SHLVL=$SHLVL +fi +ATUIN_STTY=$(stty -g) +ATUIN_HISTORY_ID="" + +export ATUIN_PREEXEC_BACKEND=$SHLVL:none +__atuin_update_preexec_backend() { + if [[ ${BLE_ATTACHED-} ]]; then + ATUIN_PREEXEC_BACKEND=$SHLVL:blesh-${BLE_VERSION-} + elif [[ ${bash_preexec_imported-} ]]; then + ATUIN_PREEXEC_BACKEND=$SHLVL:bash-preexec + elif [[ ${__bp_imported-} ]]; then + ATUIN_PREEXEC_BACKEND="$SHLVL:bash-preexec (old)" + else + ATUIN_PREEXEC_BACKEND=$SHLVL:unknown + fi +} + +__atuin_preexec() { + # Workaround for old versions of bash-preexec + if [[ ! ${BLE_ATTACHED-} ]]; then + # In older versions of bash-preexec, the preexec hook may be called + # even for the commands run by keybindings. There is no general and + # robust way to detect the command for keybindings, but at least we + # want to exclude Atuin's keybindings. When the preexec hook is called + # for a keybinding, the preexec hook for the user command will not + # fire, so we instead set a fake ATUIN_HISTORY_ID here to notify + # __atuin_precmd of this failure. + if [[ $BASH_COMMAND != "$1" ]]; then + case $BASH_COMMAND in + '__atuin_history'* | '__atuin_widget_run'* | '__atuin_bash42_dispatch'*) + ATUIN_HISTORY_ID=__bash_preexec_failure__ + return 0 ;; + esac + fi + fi + + # Note: We update ATUIN_PREEXEC_BACKEND on every preexec because blesh's + # attaching state can dynamically change. + __atuin_update_preexec_backend + + local id + id=$(atuin history start -- "$1" 2>/dev/null) + export ATUIN_HISTORY_ID=$id + __atuin_preexec_time=${EPOCHREALTIME-} +} + +__atuin_precmd() { + local EXIT=$? __atuin_precmd_time=${EPOCHREALTIME-} + + [[ ! $ATUIN_HISTORY_ID ]] && return + + # If the previous preexec hook failed, we manually call __atuin_preexec + if [[ $ATUIN_HISTORY_ID == __bash_preexec_failure__ ]]; then + # This is the command extraction code taken from bash-preexec + local previous_command + previous_command=$( + export LC_ALL=C HISTTIMEFORMAT='' + builtin history 1 | sed '1 s/^ *[0-9][0-9]*[* ] //' + ) + __atuin_preexec "$previous_command" + fi + + local duration="" + # shellcheck disable=SC2154,SC2309 + if [[ ${BLE_ATTACHED-} && ${_ble_exec_time_ata-} ]]; then + # With ble.sh, we utilize the shell variable `_ble_exec_time_ata` + # recorded by ble.sh. It is more accurate than the measurements by + # Atuin, which includes the spawn cost of Atuin. ble.sh uses the + # special shell variable `EPOCHREALTIME` in bash >= 5.0 with the + # microsecond resolution, or the builtin `time` in bash < 5.0 with the + # millisecond resolution. + duration=${_ble_exec_time_ata}000 + elif ((BASH_VERSINFO[0] >= 5)); then + # We calculate the high-resolution duration based on EPOCHREALTIME + # (bash >= 5.0) recorded by precmd/preexec, though it might not be as + # accurate as `_ble_exec_time_ata` provided by ble.sh because it + # includes the extra time of the precmd/preexec handling. Since Bash + # does not offer floating-point arithmetic, we remove the non-digit + # characters and perform the integral arithmetic. The fraction part of + # EPOCHREALTIME is fixed to have 6 digits in Bash. We remove all the + # non-digit characters because the decimal point is not necessarily a + # period depending on the locale. + duration=$((${__atuin_precmd_time//[!0-9]} - ${__atuin_preexec_time//[!0-9]})) + if ((duration >= 0)); then + duration=${duration}000 + else + duration="" # clear the result on overflow + fi + fi + + (ATUIN_LOG=error atuin history end --exit "$EXIT" ${duration:+"--duration=$duration"} -- "$ATUIN_HISTORY_ID" &) >/dev/null 2>&1 + export ATUIN_HISTORY_ID="" +} + +__atuin_set_ret_value() { + return ${1:+"$1"} +} + +#------------------------------------------------------------------------------ +# section: __atuin_accept_line +# +# The function "__atuin_accept_line" is kept for backward compatibility of the +# direct use of __atuin_history in keybindings by users. + +# The shell function `__atuin_evaluate_prompt` evaluates prompt sequences in +# $PS1. We switch the implementation of the shell function +# `__atuin_evaluate_prompt` based on the Bash version because the expansion +# ${PS1@P} is only available in bash >= 4.4. +if ((BASH_VERSINFO[0] >= 5 || BASH_VERSINFO[0] == 4 && BASH_VERSINFO[1] >= 4)); then + __atuin_evaluate_prompt() { + __atuin_set_ret_value "${__bp_last_ret_value-}" "${__bp_last_argument_prev_command-}" + __atuin_prompt=${PS1@P} + + # Note: Strip the control characters ^A (\001) and ^B (\002), which + # Bash internally uses to enclose the escape sequences. They are + # produced by '\[' and '\]', respectively, in $PS1 and used to tell + # Bash that the strings inbetween do not contribute to the prompt + # width. After the prompt width calculation, Bash strips those control + # characters before outputting it to the terminal. We here strip these + # characters following Bash's behavior. + __atuin_prompt=${__atuin_prompt//[$'\001\002']} + + # Count the number of newlines contained in $__atuin_prompt + __atuin_prompt_offset=${__atuin_prompt//[!$'\n']} + __atuin_prompt_offset=${#__atuin_prompt_offset} + } +else + __atuin_evaluate_prompt() { + __atuin_prompt='$ ' + __atuin_prompt_offset=0 + } +fi + +# The shell function `__atuin_clear_prompt N` outputs terminal control +# sequences to clear the contents of the current and N previous lines. After +# clearing, the cursor is placed at the beginning of the N-th previous line. +__atuin_clear_prompt_cache=() +__atuin_clear_prompt() { + local offset=$1 + if [[ ! ${__atuin_clear_prompt_cache[offset]+set} ]]; then + if [[ ! ${__atuin_clear_prompt_cache[0]+set} ]]; then + __atuin_clear_prompt_cache[0]=$'\r'$(tput el 2>/dev/null || tput ce 2>/dev/null) + fi + if ((offset > 0)); then + __atuin_clear_prompt_cache[offset]=${__atuin_clear_prompt_cache[0]}$( + tput cuu "$offset" 2>/dev/null || tput UP "$offset" 2>/dev/null + tput dl "$offset" 2>/dev/null || tput DL "$offset" 2>/dev/null + tput il "$offset" 2>/dev/null || tput AL "$offset" 2>/dev/null + ) + fi + fi + printf '%s' "${__atuin_clear_prompt_cache[offset]}" +} + +__atuin_accept_line() { + local __atuin_command=$1 + + # Reprint the prompt, accounting for multiple lines + local __atuin_prompt __atuin_prompt_offset + __atuin_evaluate_prompt + __atuin_clear_prompt "$__atuin_prompt_offset" + printf '%s\n' "$__atuin_prompt$__atuin_command" + + # Add it to the bash history + history -s "$__atuin_command" + + # Assuming bash-preexec + # Invoke every function in the preexec array + local __atuin_preexec_function + local __atuin_preexec_function_ret_value + local __atuin_preexec_ret_value=0 + for __atuin_preexec_function in "${preexec_functions[@]:-}"; do + if type -t "$__atuin_preexec_function" 1>/dev/null; then + __atuin_set_ret_value "${__bp_last_ret_value:-}" + "$__atuin_preexec_function" "$__atuin_command" + __atuin_preexec_function_ret_value=$? + if [[ $__atuin_preexec_function_ret_value != 0 ]]; then + __atuin_preexec_ret_value=$__atuin_preexec_function_ret_value + fi + fi + done + + # If extdebug is turned on and any preexec function returns non-zero + # exit status, we do not run the user command. + if ! { shopt -q extdebug && ((__atuin_preexec_ret_value)); }; then + # Note: When a child Bash session is started by enter_accept, if the + # environment variable READLINE_POINT is present, bash-preexec in the + # child session does not fire preexec at all because it considers we + # are inside Atuin's keybinding of the current session. To avoid + # propagating the environment variable to the child session, we remove + # the export attribute of READLINE_LINE and READLINE_POINT. + export -n READLINE_LINE READLINE_POINT + + # Juggle the terminal settings so that the command can be interacted + # with + local __atuin_stty_backup + __atuin_stty_backup=$(stty -g) + stty "$ATUIN_STTY" + + # Execute the command. Note: We need to record $? and $_ after the + # user command within the same call of "eval" because $_ is otherwise + # overwritten by the last argument of "eval". + __atuin_set_ret_value "${__bp_last_ret_value-}" "${__bp_last_argument_prev_command-}" + eval -- "$__atuin_command"$'\n__bp_last_ret_value=$? __bp_last_argument_prev_command=$_' + + stty "$__atuin_stty_backup" + fi + + # Execute preprompt commands + local __atuin_prompt_command + for __atuin_prompt_command in "${PROMPT_COMMAND[@]}"; do + __atuin_set_ret_value "${__bp_last_ret_value-}" "${__bp_last_argument_prev_command-}" + eval -- "$__atuin_prompt_command" + done + # Bash will redraw only the line with the prompt after we finish, + # so to work for a multiline prompt we need to print it ourselves, + # then go to the beginning of the last line. + __atuin_evaluate_prompt + printf '%s' "$__atuin_prompt" + __atuin_clear_prompt 0 +} + +#------------------------------------------------------------------------------ + +# Check if tmux popup is available (tmux >= 3.2) +__atuin_tmux_popup_check() { + [[ -n "${TMUX-}" ]] || return 1 + [[ "${ATUIN_TMUX_POPUP:-true}" != "false" ]] || return 1 + + # https://github.com/tmux/tmux/wiki/FAQ#how-often-is-tmux-released-what-is-the-version-number-scheme + local tmux_version + tmux_version=$(tmux -V 2>/dev/null | sed -n 's/^[^0-9]*\([0-9][0-9]*\.[0-9][0-9]*\).*/\1/p') # Could have used grep... + [[ -z "$tmux_version" ]] && return 1 + + local m1 m2 + m1=${tmux_version%%.*} + m2=${tmux_version#*.} + m2=${m2%%.*} + [[ "$m1" =~ ^[0-9]+$ ]] || return 1 + [[ "$m2" =~ ^[0-9]+$ ]] || m2=0 + (( m1 > 3 || (m1 == 3 && m2 >= 2) )) +} + +# Use global variable to fix scope issues with traps +__atuin_popup_tmpdir="" +__atuin_tmux_popup_cleanup() { + [[ -n "$__atuin_popup_tmpdir" && -d "$__atuin_popup_tmpdir" ]] && command rm -rf "$__atuin_popup_tmpdir" + __atuin_popup_tmpdir="" +} + +__atuin_search_cmd() { + local -a search_args=("$@") + + if __atuin_tmux_popup_check; then + __atuin_popup_tmpdir=$(mktemp -d) || return 1 + local result_file="$__atuin_popup_tmpdir/result" + + trap '__atuin_tmux_popup_cleanup' EXIT HUP INT TERM + + local escaped_query escaped_args + escaped_query=$(printf '%s' "$READLINE_LINE" | sed "s/'/'\\\\''/g") + escaped_args="" + for arg in "${search_args[@]}"; do + escaped_args+=" '$(printf '%s' "$arg" | sed "s/'/'\\\\''/g")'" + done + + # In the popup, atuin goes to terminal, stderr goes to file + local cdir popup_width popup_height + cdir=$(pwd) + popup_width="${ATUIN_TMUX_POPUP_WIDTH:-80%}" # Keep default value anyways + popup_height="${ATUIN_TMUX_POPUP_HEIGHT:-60%}" + tmux display-popup -d "$cdir" -w "$popup_width" -h "$popup_height" -E -E -- \ + sh -c "PATH='$PATH' ATUIN_SESSION='$ATUIN_SESSION' ATUIN_SHELL=bash ATUIN_LOG=error ATUIN_QUERY='$escaped_query' atuin search $escaped_args -i 2>'$result_file'" + + if [[ -f "$result_file" ]]; then + cat "$result_file" + fi + + __atuin_tmux_popup_cleanup + trap - EXIT HUP INT TERM + else + ATUIN_SHELL=bash ATUIN_LOG=error ATUIN_QUERY=$READLINE_LINE atuin search "${search_args[@]}" -i 3>&1 1>&2 2>&3 + fi +} + +__atuin_history() { + # Default action of the up key: When this function is called with the first + # argument `--shell-up-key-binding`, we perform Atuin's history search only + # when the up key is supposed to cause the history movement in the original + # binding. We do this only for ble.sh because the up key always invokes + # the history movement in the plain Bash. + if [[ ${BLE_ATTACHED-} && ${1-} == --shell-up-key-binding ]]; then + # When the current cursor position is not in the first line, the up key + # should move the cursor to the previous line. While the selection is + # performed, the up key should not start the history search. + # shellcheck disable=SC2154 # Note: these variables are set by ble.sh + if [[ ${_ble_edit_str::_ble_edit_ind} == *$'\n'* || $_ble_edit_mark_active ]]; then + ble/widget/@nomarked backward-line + local status=$? + READLINE_LINE=$_ble_edit_str + READLINE_POINT=$_ble_edit_ind + READLINE_MARK=$_ble_edit_mark + return "$status" + fi + fi + + # READLINE_LINE and READLINE_POINT are only supported by bash >= 4.0 or + # ble.sh. When it is not supported, we clear them to suppress strange + # behaviors. + [[ ${BLE_ATTACHED-} ]] || ((BASH_VERSINFO[0] >= 4)) || + READLINE_LINE="" READLINE_POINT=0 + + local __atuin_output + __atuin_output=$(__atuin_search_cmd "$@") + + # We do nothing when the search is canceled. + [[ $__atuin_output ]] || return 0 + + if [[ $__atuin_output == __atuin_accept__:* ]]; then + __atuin_output=${__atuin_output#__atuin_accept__:} + + if [[ ${BLE_ATTACHED-} ]]; then + ble-edit/content/reset-and-check-dirty "$__atuin_output" + ble/widget/accept-line + READLINE_LINE="" + elif [[ ${__atuin_macro_chain_keymap-} ]]; then + READLINE_LINE=$__atuin_output + bind -m "$__atuin_macro_chain_keymap" '"'"$__atuin_macro_chain"'": '"$__atuin_macro_accept_line" + else + __atuin_accept_line "$__atuin_output" + READLINE_LINE="" + fi + + READLINE_POINT=${#READLINE_LINE} + else + READLINE_LINE=$__atuin_output + READLINE_POINT=${#READLINE_LINE} + if [[ ! ${BLE_ATTACHED-} ]] && ((BASH_VERSINFO[0] < 4)) && [[ ${__atuin_macro_chain_keymap-} ]]; then + bind -m "$__atuin_macro_chain_keymap" '"'"$__atuin_macro_chain"'": '"$__atuin_macro_insert_line" + fi + fi +} + +__atuin_initialize_blesh() { + # shellcheck disable=SC2154 + [[ ${BLE_VERSION-} ]] && ((_ble_version >= 400)) || return 0 + + ble-import contrib/integration/bash-preexec + + # Define and register an autosuggestion source for ble.sh's auto-complete. + # If you'd like to overwrite this, define the same name of shell function + # after the $(atuin init bash) line in your .bashrc. If you do not need + # the auto-complete source by Atuin, please add the following code to + # remove the entry after the $(atuin init bash) line in your .bashrc: + # + # ble/util/import/eval-after-load core-complete ' + # ble/array#remove _ble_complete_auto_source atuin-history' + # + function ble/complete/auto-complete/source:atuin-history { + local suggestion + suggestion=$(ATUIN_QUERY="$_ble_edit_str" atuin search --cmd-only --limit 1 --search-mode prefix 2>/dev/null) + [[ $suggestion == "$_ble_edit_str"?* ]] || return 1 + ble/complete/auto-complete/enter h 0 "${suggestion:${#_ble_edit_str}}" '' "$suggestion" + } + ble/util/import/eval-after-load core-complete ' + ble/array#unshift _ble_complete_auto_source atuin-history' + + # @env BLE_SESSION_ID: `atuin doctor` references the environment variable + # BLE_SESSION_ID. We explicitly export the variable because it was not + # exported in older versions of ble.sh. + [[ ${BLE_SESSION_ID-} ]] && export BLE_SESSION_ID +} +__atuin_initialize_blesh +BLE_ONLOAD+=(__atuin_initialize_blesh) +precmd_functions+=(__atuin_precmd) +preexec_functions+=(__atuin_preexec) + +#------------------------------------------------------------------------------ +# section: atuin-bind + +__atuin_widget=() + +__atuin_widget_save() { + local data=$1 + for REPLY in "${!__atuin_widget[@]}"; do + if [[ ${__atuin_widget[REPLY]} == "$data" ]]; then + return 0 + fi + done + # shellcheck disable=SC2154 + REPLY=${#__atuin_widget[*]} + __atuin_widget[REPLY]=$data +} + +__atuin_widget_run() { + local data=${__atuin_widget[$1]} + local keymap=${data%%:*} widget=${data#*:} + local __atuin_macro_chain_keymap=$keymap + bind -m "$keymap" '"'"$__atuin_macro_chain"'": ""' + builtin eval -- "$widget" +} + +# To realize the enter_accept feature in a robust way, we need to call the +# readline bindable function `accept-line'. However, there is no way to call +# `accept-line' from the shell script. To call the bindable function +# `accept-line', we may utilize string macros of readline. When we bind KEYSEQ +# to a WIDGET that wants to conditionally call `accept-line' at the end, we +# perform two-step dispatching: +# +# 1. [KEYSEQ -> IKEYSEQ1 IKEYSEQ2]---We first translate KEYSEQ to two +# intermediate key sequences IKEYSEQ1 and IKEYSEQ2 using string macros. For +# example, when we bind `__atuin_history` to \C-r, this step can be set up by +# `bind '"\C-r": "IKEYSEQ1IKEYSEQ2"'`. +# +# 2. [IKEYSEQ1 -> WIDGET]---Then, IKEYSEQ1 is bound to the WIDGET, and the +# binding of IKEYSEQ2 is dynamically determined by WIDGET. For example, when +# we bind `__atuin_history` to \C-r, this step can be set up by `bind -x +# '"IKEYSEQ1": WIDGET'`. +# +# 3. [IKEYSEQ2 -> accept-line] or [IKEYSEQ2 -> ""]---To request the execution +# of `accept-line', WIDGET can change the binding of IKEYSEQ2 by running +# `bind '"IKEYSEQ2": accept-line''. Otherwise, WIDGET can change the binding +# of IKEYSEQ2 to no-op by running `bind '"IKEYSEQ2": ""'`. +# +# For the choice of the intermediate key sequences, we want to choose key +# sequences that are unlikely to conflict with others. In addition, we want to +# avoid a key sequence containing \e because keymap "vi-insert" stops +# processing key sequences containing \e in older versions of Bash. We have +# used \e[0;<m>A (a variant of the [up] key with modifier <m>) in Atuin 3.10.0 +# for intermediate key sequences, but this contains \e and caused a problem. +# Instead, we use \C-x\C-_A<n>\a, which starts with \C-x\C-_ (an unlikely +# two-byte combination) and A (represents the initial letter of Atuin), +# followed by the payload <n> and the terminator \a (BEL, \C-g). + +__atuin_macro_chain='\C-x\C-_A0\a' +for __atuin_keymap in emacs vi-insert vi-command; do + bind -m "$__atuin_keymap" "\"$__atuin_macro_chain\": \"\"" +done +unset -v __atuin_keymap + +if ((BASH_VERSINFO[0] >= 5 || BASH_VERSINFO[0] == 4 && BASH_VERSINFO[1] >= 3)); then + # In Bash >= 4.3 + + __atuin_macro_accept_line=accept-line + + __atuin_bind_impl() { + local keymap=$1 keyseq=$2 command=$3 + + # Note: In Bash <= 5.0, the table for `bind -x` from the keyseq to the + # command is shared by all the keymaps (emacs, vi-insert, and + # vi-command), so one cannot safely bind different command strings to + # the same keyseq in different keymaps. Therefore, the command string + # and the keyseq need to be globally in one-to-one correspondence in + # all the keymaps. + local REPLY + __atuin_widget_save "$keymap:$command" + local widget=$REPLY + local ikeyseq1='\C-x\C-_A'$((1 + widget))'\a' + local ikeyseq2=$__atuin_macro_chain + + if ((BASH_VERSINFO[0] == 5 && BASH_VERSINFO[1] == 1)); then + # Workaround for Bash 5.1: Bash 5.1 has a bug that overwriting an + # existing "bind -x" keybinding breaks other existing "bind -x" + # keybindings [1,2]. To work around the problem, we explicitly + # unbind an existing keybinding before overwriting it. + # + # [1] https://lists.gnu.org/archive/html/bug-bash/2021-04/msg00135.html + # [2] https://github.com/atuinsh/atuin/issues/962#issuecomment-3451132291 + bind -m "$keymap" -r "$keyseq" + fi + + bind -m "$keymap" "\"$keyseq\": \"$ikeyseq1$ikeyseq2\"" + bind -m "$keymap" -x "\"$ikeyseq1\": __atuin_widget_run $widget" + } + + __atuin_bind_blesh_onload() { + # In ble.sh, we need to enable unrecognized CSI sequences like \e[0;0A, + # which are discarded by ble.sh by default. Note: In Bash <= 4.2, we + # do not need to unset "decode_error_cseq_discard" because \e[0;<m>A is + # used only for the macro chaining (which is unused by ble.sh) in Bash + # <= 4.2. + bleopt decode_error_cseq_discard= + } + if [[ ${BLE_VERSION-} ]]; then + __atuin_bind_blesh_onload + fi + BLE_ONLOAD+=(__atuin_bind_blesh_onload) +else + # In Bash <= 4.2, "bind -x" cannot bind a shell command to a keyseq having + # more than two bytes, so we need to work with only two-byte sequences. + # + # However, the number of available combinations of two-byte sequences is + # limited. To minimize the number of key sequences used by Atuin, instead + # of specifying a widget by its own intermediate sequence, we specify a + # widget by a fixed-length sequence of multiple two-byte sequences. More + # specifically, instead of IKEYSEQ1, we use IKS1 IKS2 IKS3 [IKS4 IKS5] + # IKSX, where IKS1..IKS5 just stores its information to a global variable, + # and IKSX collects all the information and determine and call the actual + # widget based on the stored information. Each of IKn (n=1..5) is one of + # the two reserved sequences, $__atuin_bash42_code0 and + # $__atuin_bash42_code1. IKSX is fixed to be $__atuin_bash42_code2. + # + # For the choices of the special key sequences, we consider \C-xQ, \C-xR, + # and \C-xS. In the emacs editing mode of Bash, \C-x is used as a prefix + # key, i.e., it is used for the beginning key of the keybindings with + # multiple keys, so \C-x is unlikely to be used for a single-key binding by + # the user. Also, \C-x is not used in the vi editing mode by default. The + # combinations \C-xQ..\C-xS are also unlikely be used because we need to + # switch the modifier keys from Control to Shift to input these sequences, + # and these are not easy to input. + __atuin_bash42_code0='\C-xQ' + __atuin_bash42_code1='\C-xR' + __atuin_bash42_code2='\C-xS' + + __atuin_bash42_encode() { + REPLY= + local n=$1 min_width=${2-} + while + if ((n % 2 == 0)); then + REPLY=$__atuin_bash42_code0$REPLY + else + REPLY=$__atuin_bash42_code1$REPLY + fi + (((n /= 2) || ${#REPLY} / ${#__atuin_bash42_code0} < min_width)) + do :; done + } + + __atuin_bash42_bind() { + local __atuin_keymap + for __atuin_keymap in emacs vi-insert vi-command; do + bind -m "$__atuin_keymap" -x '"'"$__atuin_bash42_code0"'": __atuin_bash42_dispatch_selector+=0' + bind -m "$__atuin_keymap" -x '"'"$__atuin_bash42_code1"'": __atuin_bash42_dispatch_selector+=1' + bind -m "$__atuin_keymap" -x '"'"$__atuin_bash42_code2"'": __atuin_bash42_dispatch' + done + } + __atuin_bash42_bind + # In Bash <= 4.2, there is no way to read users' "bind -x" settings, so we + # need to explicitly perform "bind -x" when ble.sh is loaded. + BLE_ONLOAD+=(__atuin_bash42_bind) + + if ((BASH_VERSINFO[0] >= 4)); then + __atuin_macro_accept_line=accept-line + else + # Note: We rewrite the command line and invoke `accept-line'. In + # bash <= 3.2, there is no way to rewrite the command line from the + # shell script, so we rewrite it using a macro and + # `shell-expand-line'. + # + # Note: Concerning the key sequences to invoke bindable functions + # such as "\C-x\C-_A1\a", another option is to use + # "\exbegginning-of-line\r", etc. to make it consistent with bash + # >= 5.3. However, an older Bash configuration can still conflict + # on [M-x]. The conflict is more likely than \C-x\C-_A1\a. + for __atuin_keymap in emacs vi-insert vi-command; do + bind -m "$__atuin_keymap" '"\C-x\C-_A1\a": beginning-of-line' + bind -m "$__atuin_keymap" '"\C-x\C-_A2\a": kill-line' + # shellcheck disable=SC2016 + bind -m "$__atuin_keymap" '"\C-x\C-_A3\a": "$READLINE_LINE"' + bind -m "$__atuin_keymap" '"\C-x\C-_A4\a": shell-expand-line' + bind -m "$__atuin_keymap" '"\C-x\C-_A5\a": accept-line' + bind -m "$__atuin_keymap" '"\C-x\C-_A6\a": end-of-line' + done + unset -v __atuin_keymap + + bind -m vi-command '"\C-x\C-_A7\a": vi-insertion-mode' + bind -m vi-insert '"\C-x\C-_A7\a": vi-movement-mode' + + # "\C-x\C-_A10\a": Replace the command line with READLINE_LINE. When we are + # in the vi-command keymap, we go to vi-insert, input + # "$READLINE_LINE", and come back to vi-command. + bind -m emacs '"\C-x\C-_A10\a": "\C-x\C-_A1\a\C-x\C-_A2\a\C-x\C-_A3\a\C-x\C-_A4\a"' + bind -m vi-insert '"\C-x\C-_A10\a": "\C-x\C-_A1\a\C-x\C-_A2\a\C-x\C-_A3\a\C-x\C-_A4\a"' + bind -m vi-command '"\C-x\C-_A10\a": "\C-x\C-_A1\a\C-x\C-_A2\a\C-x\C-_A7\a\C-x\C-_A3\a\C-x\C-_A7\a\C-x\C-_A4\a"' + + __atuin_macro_accept_line='"\C-x\C-_A10\a\C-x\C-_A5\a"' + __atuin_macro_insert_line='"\C-x\C-_A10\a\C-x\C-_A6\a"' + fi + + __atuin_bash42_dispatch_selector= + + __atuin_bash42_dispatch() { + local s=$__atuin_bash42_dispatch_selector + __atuin_bash42_dispatch_selector= + __atuin_widget_run "$((2#0$s))" + } + + __atuin_bind_impl() { + local keymap=$1 keyseq=$2 command=$3 + + __atuin_widget_save "$keymap:$command" + __atuin_bash42_encode "$REPLY" + local macro=$REPLY$__atuin_bash42_code2$__atuin_macro_chain + + bind -m "$keymap" "\"$keyseq\": \"$macro\"" + } +fi + +atuin-bind() { + local keymap= + local OPTIND=1 OPTARG="" OPTERR=0 flag + while getopts ':m:' flag "$@"; do + case $flag in + m) keymap=$OPTARG ;; + *) + printf '%s\n' "atuin-bind: unrecognized option '-$flag'" >&2 + return 2 + ;; + esac + done + shift "$((OPTIND - 1))" + + if (($# != 2)); then + printf '%s\n' 'usage: atuin-bind [-m keymap] keyseq widget' >&2 + return 2 + fi + + local keyseq=$1 + [[ $keymap ]] || keymap=$(bind -v | awk '$2 == "keymap" { print $3 }') + case $keymap in + emacs-meta) keymap=emacs keyseq='\e'$keyseq ;; + emacs-ctlx) keymap=emacs keyseq='\C-x'$keyseq ;; + emacs*) keymap=emacs ;; + vi-insert) ;; + vi*) keymap=vi-command ;; + *) + printf '%s\n' "atuin-bind: unknown keymap '$keymap'" >&2 + return 2 ;; + esac + + local command=$2 widget=${2%%[[:blank:]]*} + case $widget in + atuin-search) command=${2/#"$widget"/__atuin_history} ;; + atuin-search-emacs) command=${2/#"$widget"/__atuin_history --keymap-mode=emacs} ;; + atuin-search-viins) command=${2/#"$widget"/__atuin_history --keymap-mode=vim-insert} ;; + atuin-search-vicmd) command=${2/#"$widget"/__atuin_history --keymap-mode=vim-normal} ;; + atuin-up-search) command=${2/#"$widget"/__atuin_history --shell-up-key-binding} ;; + atuin-up-search-emacs) command=${2/#"$widget"/__atuin_history --shell-up-key-binding --keymap-mode=emacs} ;; + atuin-up-search-viins) command=${2/#"$widget"/__atuin_history --shell-up-key-binding --keymap-mode=vim-insert} ;; + atuin-up-search-vicmd) command=${2/#"$widget"/__atuin_history --shell-up-key-binding --keymap-mode=vim-normal} ;; + esac + + __atuin_bind_impl "$keymap" "$keyseq" "$command" +} + +#------------------------------------------------------------------------------ + +# shellcheck disable=SC2154 +if [[ $__atuin_bind_ctrl_r == true ]]; then + # Note: We do not overwrite [C-r] in the vi-command keymap because we do + # not want to overwrite "redo", which is already bound to [C-r] in the + # vi_nmap keymap in ble.sh. + atuin-bind -m emacs '\C-r' atuin-search-emacs + atuin-bind -m vi-insert '\C-r' atuin-search-viins + atuin-bind -m vi-command '/' atuin-search-emacs +fi + +# shellcheck disable=SC2154 +if [[ $__atuin_bind_up_arrow == true ]]; then + atuin-bind -m emacs '\e[A' atuin-up-search-emacs + atuin-bind -m emacs '\eOA' atuin-up-search-emacs + atuin-bind -m vi-insert '\e[A' atuin-up-search-viins + atuin-bind -m vi-insert '\eOA' atuin-up-search-viins + atuin-bind -m vi-command '\e[A' atuin-up-search-vicmd + atuin-bind -m vi-command '\eOA' atuin-up-search-vicmd + atuin-bind -m vi-command 'k' atuin-up-search-vicmd +fi + +#------------------------------------------------------------------------------ +fi # (include guard) end of main content diff --git a/crates/atuin/src/shell/atuin.fish b/crates/atuin/src/shell/atuin.fish new file mode 100644 index 00000000..d1333e46 --- /dev/null +++ b/crates/atuin/src/shell/atuin.fish @@ -0,0 +1,152 @@ +if not set -q ATUIN_SESSION; or test "$ATUIN_SHLVL" != "$SHLVL" + set -gx ATUIN_SESSION (atuin uuid) + set -gx ATUIN_SHLVL $SHLVL +end +set --erase ATUIN_HISTORY_ID + +function _atuin_preexec --on-event fish_preexec + if not test -n "$fish_private_mode" + set -g ATUIN_HISTORY_ID (atuin history start -- "$argv[1]" 2>/dev/null) + end +end + +function _atuin_postexec --on-event fish_postexec + set -l s $status + + if test -n "$ATUIN_HISTORY_ID" + ATUIN_LOG=error atuin history end --exit $s -- $ATUIN_HISTORY_ID &>/dev/null & + disown + end + + set --erase ATUIN_HISTORY_ID +end + +# Check if tmux popup is available (tmux >= 3.2) +function _atuin_tmux_popup_check + if not test -n "$TMUX" + echo 0 + return + end + + if test "$ATUIN_TMUX_POPUP" = false + echo 0 + return + end + + set -l tmux_version (tmux -V 2>/dev/null | string match -r '\d+\.\d+') + if not test -n "$tmux_version" + echo 0 + return + end + + set -l parts (string split '.' $tmux_version) + set -l m1 $parts[1] + set -l m2 0 + if test (count $parts) -ge 2 + set m2 $parts[2] + end + + if not string match -rq '^[0-9]+$' -- "$m1" + echo 0 + return + end + + if not string match -rq '^[0-9]+$' -- "$m2" + set m2 0 + end + + if test "$m1" -gt 3 2>/dev/null; or begin + test "$m1" -eq 3 2>/dev/null; and test "$m2" -ge 2 2>/dev/null + end + echo 1 + else + echo 0 + end +end + +function _atuin_search + set -l keymap_mode + switch $fish_key_bindings + case fish_vi_key_bindings + switch $fish_bind_mode + case default + set keymap_mode vim-normal + case insert + set keymap_mode vim-insert + end + case '*' + set keymap_mode emacs + end + + set -l use_tmux_popup (_atuin_tmux_popup_check) + + set -l ATUIN_H + if test "$use_tmux_popup" -eq 1 + set -l tmpdir (mktemp -d) + if not test -d "$tmpdir" + # if mktemp got errors + set ATUIN_H (ATUIN_SHELL=fish ATUIN_LOG=error ATUIN_QUERY=(commandline -b) atuin search --keymap-mode=$keymap_mode $argv -i 3>&1 1>&2 2>&3 | string collect) + else + set -l result_file "$tmpdir/result" + + set -l query (commandline -b | string replace -a "'" "'\\''") + set -l escaped_args "" + for arg in $argv + set escaped_args "$escaped_args '"(string replace -a "'" "'\\''" -- $arg)"'" + end + + # In the popup, atuin goes to terminal, stderr goes to file + set -l cdir (pwd) + # Keep default value anyways + set -l popup_width (test -n "$ATUIN_TMUX_POPUP_WIDTH" && echo "$ATUIN_TMUX_POPUP_WIDTH" || echo "80%") + set -l popup_height (test -n "$ATUIN_TMUX_POPUP_HEIGHT" && echo "$ATUIN_TMUX_POPUP_HEIGHT" || echo "60%") + tmux display-popup -d "$cdir" -w "$popup_width" -h "$popup_height" -E -E -- \ + sh -c "PATH='$PATH' ATUIN_SESSION='$ATUIN_SESSION' ATUIN_SHELL=fish ATUIN_LOG=error ATUIN_QUERY='$query' atuin search --keymap-mode=$keymap_mode$escaped_args -i 2>'$result_file'" + + if test -f "$result_file" + set ATUIN_H (cat "$result_file" | string collect) + end + + command rm -rf "$tmpdir" + end + else + # In fish 3.4 and above we can use `"$(some command)"` to keep multiple lines separate; + # but to support fish 3.3 we need to use `(some command | string collect)`. + # https://fishshell.com/docs/current/relnotes.html#id24 (fish 3.4 "Notable improvements and fixes") + set ATUIN_H (ATUIN_SHELL=fish ATUIN_LOG=error ATUIN_QUERY=(commandline -b) atuin search --keymap-mode=$keymap_mode $argv -i 3>&1 1>&2 2>&3 | string collect) + end + + set ATUIN_H (string trim -- $ATUIN_H | string collect) # trim whitespace + + if test -n "$ATUIN_H" + if string match --quiet '__atuin_accept__:*' "$ATUIN_H" + set -l ATUIN_HIST (string replace "__atuin_accept__:" "" -- "$ATUIN_H" | string collect) + commandline -r "$ATUIN_HIST" + commandline -f repaint + commandline -f execute + return + else + commandline -r "$ATUIN_H" + end + end + + commandline -f repaint +end + +function _atuin_bind_up + # Fallback to fish's builtin up-or-search if we're in search or paging mode + if commandline --search-mode; or commandline --paging-mode + up-or-search + return + end + + # Only invoke atuin if we're on the top line of the command + set -l lineno (commandline --line) + + switch $lineno + case 1 + _atuin_search --shell-up-key-binding + case '*' + up-or-search + end +end diff --git a/crates/atuin/src/shell/atuin.nu b/crates/atuin/src/shell/atuin.nu new file mode 100644 index 00000000..3bc59edc --- /dev/null +++ b/crates/atuin/src/shell/atuin.nu @@ -0,0 +1,101 @@ +# Source this in your ~/.config/nushell/config.nu +# minimum supported version = 0.93.0 +module compat { + export def --wrapped "random uuid -v 7" [...rest] { atuin uuid } +} +use (if not ( + (version).major > 0 or + (version).minor >= 103 +) { "compat" }) * + +if 'ATUIN_SESSION' not-in $env or ('ATUIN_SHLVL' not-in $env) or ($env.ATUIN_SHLVL != ($env.SHLVL? | default "")) { + $env.ATUIN_SESSION = (random uuid -v 7 | str replace -a "-" "") + $env.ATUIN_SHLVL = ($env.SHLVL? | default "") +} +hide-env -i ATUIN_HISTORY_ID + +# Magic token to make sure we don't record commands run by keybindings +let ATUIN_KEYBINDING_TOKEN = $"# (random uuid)" + +let _atuin_pre_execution = {|| + if ($nu | get history-enabled?) == false { + return + } + let cmd = (commandline) + if ($cmd | is-empty) { + return + } + if not ($cmd | str starts-with $ATUIN_KEYBINDING_TOKEN) { + $env.ATUIN_HISTORY_ID = (atuin history start -- $cmd e>| complete | get stdout | str trim) + } +} + +let _atuin_pre_prompt = {|| + let last_exit = $env.LAST_EXIT_CODE + if 'ATUIN_HISTORY_ID' not-in $env { + return + } + with-env { ATUIN_LOG: error } { + if (version).minor >= 111 or (version).major > 0 { + job spawn -d atuin { + ^atuin history end $'--exit=($env.LAST_EXIT_CODE)' -- $env.ATUIN_HISTORY_ID | complete + } | ignore + } else if (version).minor >= 104 or (version).major > 0 { + job spawn { + ^atuin history end $'--exit=($env.LAST_EXIT_CODE)' -- $env.ATUIN_HISTORY_ID | complete + } | ignore + } else { + do { atuin history end $'--exit=($last_exit)' -- $env.ATUIN_HISTORY_ID } | complete + } + + } + hide-env ATUIN_HISTORY_ID +} + +def _atuin_search_cmd [...flags: string] { + if (version).minor >= 106 or (version).major > 0 { + [ + $ATUIN_KEYBINDING_TOKEN, + ([ + `with-env { ATUIN_LOG: error, ATUIN_QUERY: (commandline), ATUIN_SHELL: nu } {`, + ([ + 'let output = (run-external atuin search', + ($flags | append [--interactive] | each {|e| $'"($e)"'}), + 'e>| str trim)', + ] | flatten | str join ' '), + 'if ($output | str starts-with "__atuin_accept__:") {', + 'commandline edit --accept ($output | str replace "__atuin_accept__:" "")', + '} else {', + 'commandline edit $output', + '}', + `}`, + ] | flatten | str join "\n"), + ] + } else { + [ + $ATUIN_KEYBINDING_TOKEN, + ([ + `with-env { ATUIN_LOG: error, ATUIN_QUERY: (commandline) } {`, + 'commandline edit', + '(run-external atuin search', + ($flags | append [--interactive] | each {|e| $'"($e)"'}), + ' e>| str trim)', + `}`, + ] | flatten | str join ' '), + ] + } | str join "\n" +} + +$env.config = ($env | default {} config).config +$env.config = ($env.config | default {} hooks) +$env.config = ( + $env.config | upsert hooks ( + $env.config.hooks + | upsert pre_execution ( + $env.config.hooks | get pre_execution? | default [] | append $_atuin_pre_execution) + | upsert pre_prompt ( + $env.config.hooks | get pre_prompt? | default [] | append $_atuin_pre_prompt) + ) +) + +$env.config = ($env.config | default [] keybindings) diff --git a/crates/atuin/src/shell/atuin.ps1 b/crates/atuin/src/shell/atuin.ps1 new file mode 100644 index 00000000..7f6ad8cf --- /dev/null +++ b/crates/atuin/src/shell/atuin.ps1 @@ -0,0 +1,238 @@ +# Atuin PowerShell module +# +# This should support PowerShell 5.1 (which is shipped with Windows) and later versions, on Windows and Linux. +# +# Usage: atuin init powershell | Out-String | Invoke-Expression +# +# Settings: +# - $env:ATUIN_POWERSHELL_PROMPT_OFFSET - Number of lines to offset the prompt position after exiting search. +# This is useful when using a multi-line prompt: e.g. set this to -1 when using a 2-line prompt. +# It is initialized from the current prompt line count if not set when the first Atuin search is performed. + +if (Get-Module Atuin -ErrorAction Ignore) { + if ($PSVersionTable.PSVersion.Major -ge 7) { + Write-Warning "The Atuin module is already loaded, replacing it." + Remove-Module Atuin + } else { + Write-Warning "The Atuin module is already loaded, skipping." + return + } +} + +if (!(Get-Command atuin -ErrorAction Ignore)) { + Write-Error "The 'atuin' executable needs to be available in the PATH." + return +} + +if (!(Get-Module PSReadLine -ErrorAction Ignore)) { + Write-Error "Atuin requires the PSReadLine module to be installed." + return +} + +New-Module -Name Atuin -ScriptBlock { + if (-not $env:ATUIN_SESSION -or $env:ATUIN_PID -ne $PID) { + $env:ATUIN_SESSION = atuin uuid + $env:ATUIN_PID = $PID + } + + $script:atuinHistoryId = $null + $script:previousPSConsoleHostReadLine = $Function:PSConsoleHostReadLine + + # The ReadLine overloads changed with breaking changes over time, make sure the one we expect is available. + $script:hasExpectedReadLineOverload = ([Microsoft.PowerShell.PSConsoleReadLine]::ReadLine).OverloadDefinitions.Contains("static string ReadLine(runspace runspace, System.Management.Automation.EngineIntrinsics engineIntrinsics, System.Threading.CancellationToken cancellationToken, System.Nullable[bool] lastRunStatus)") + + function Get-CommandLine { + $commandLine = "" + [Microsoft.PowerShell.PSConsoleReadLine]::GetBufferState([ref]$commandLine, [ref]$null) + return $commandLine + } + + function Set-CommandLine { + param([string]$Text) + + $commandLine = Get-CommandLine + [Microsoft.PowerShell.PSConsoleReadLine]::Replace(0, $commandLine.Length, $Text) + } + + # This function name is called by PSReadLine to read the next command line to execute. + # We replace it with a custom implementation which adds Atuin support. + function PSConsoleHostReadLine { + ## 1. Collect the exit code of the previous command. + + # This needs to be done as the first thing because any script run will flush $?. + $lastRunStatus = $? + + # Exit statuses are maintained separately for native and PowerShell commands, this needs to be taken into account. + $lastNativeExitCode = $global:LASTEXITCODE + $exitCode = if ($lastRunStatus) { 0 } elseif ($lastNativeExitCode) { $lastNativeExitCode } else { 1 } + + ## 2. Report the status of the previous command to Atuin (atuin history end). + + if ($script:atuinHistoryId) { + try { + # The duration is not recorded in old PowerShell versions, let Atuin handle it. $null arguments are ignored. + $duration = (Get-History -Count 1).Duration.Ticks * 100 + $durationArg = if ($duration) { "--duration=$duration" } else { $null } + + # Fire and forget the atuin history end command to avoid blocking the shell during a potential sync. + $process = New-Object System.Diagnostics.Process + $process.StartInfo.FileName = "atuin" + $process.StartInfo.Arguments = "history end --exit=$exitCode $durationArg -- $script:atuinHistoryId" + $process.StartInfo.UseShellExecute = $false + $process.StartInfo.CreateNoWindow = $true + $process.StartInfo.RedirectStandardInput = $true + $process.StartInfo.RedirectStandardOutput = $true + $process.StartInfo.RedirectStandardError = $true + $process.Start() | Out-Null + $process.StandardInput.Close() + $process.BeginOutputReadLine() + $process.BeginErrorReadLine() + } + catch { + # Ignore errors to avoid breaking the shell. + # An error would occur if the user removes atuin from the PATH, for instance. + } + finally { + $script:atuinHistoryId = $null + } + } + + ## 3. Read the next command line to execute. + + # PSConsoleHostReadLine implementation from PSReadLine, adjusted to support old versions. + Microsoft.PowerShell.Core\Set-StrictMode -Off + + $line = if ($script:hasExpectedReadLineOverload) { + # When the overload we expect is available, we can pass $lastRunStatus to it. + [Microsoft.PowerShell.PSConsoleReadLine]::ReadLine($Host.Runspace, $ExecutionContext, [System.Threading.CancellationToken]::None, $lastRunStatus) + } else { + # Either PSReadLine is older than v2.2.0-beta3, or maybe newer than we expect, so use the function from PSReadLine as-is. + & $script:previousPSConsoleHostReadLine + } + + ## 4. Report the next command line to Atuin (atuin history start). + + # PowerShell doesn't handle double quotes in native command line arguments the same way depending on its version, + # and the value of $PSNativeCommandArgumentPassing - see the about_Parsing help page which explains the breaking changes. + # This makes it unreliable, so we go through an environment variable, which should always be consistent across versions. + try { + $env:ATUIN_COMMAND_LINE = $line + $script:atuinHistoryId = atuin history start --command-from-env + } + catch { + # Ignore errors to avoid breaking the shell, see above. + } + finally { + $env:ATUIN_COMMAND_LINE = $null + } + + $global:LASTEXITCODE = $lastNativeExitCode + return $line + } + + function Invoke-AtuinSearch { + param([string]$ExtraArgs = "") + + $previousOutputEncoding = [System.Console]::OutputEncoding + $resultFile = New-TemporaryFile + $suggestion = "" + $errorOutput = "" + + try { + [System.Console]::OutputEncoding = [System.Text.Encoding]::UTF8 + + # Start-Process does some crazy stuff, just use the Process class directly to have more control. + $process = New-Object System.Diagnostics.Process + $process.StartInfo.FileName = "atuin" + $process.StartInfo.Arguments = "search -i --result-file ""$resultFile"" $ExtraArgs" + $process.StartInfo.UseShellExecute = $false + $process.StartInfo.RedirectStandardError = $true + $process.StartInfo.StandardErrorEncoding = [System.Text.Encoding]::UTF8 + $process.StartInfo.EnvironmentVariables["ATUIN_SHELL"] = "powershell" + $process.StartInfo.EnvironmentVariables["ATUIN_QUERY"] = Get-CommandLine + + try { + $process.Start() | Out-Null + + # A single stream is redirected, so we can read it synchronously, but we have to start reading it + # before waiting for the process to exit, otherwise the buffer could fill up and cause a deadlock. + $errorOutput = $process.StandardError.ReadToEnd().Trim() + $process.WaitForExit() + + $suggestion = (Get-Content -Raw $resultFile -Encoding UTF8 | Out-String).Trim() + } + catch { + $errorOutput = $_ + } + + if ($errorOutput) { + Write-Host -ForegroundColor Red "Atuin error:" + Write-Host -ForegroundColor DarkRed $errorOutput + } + + # If no shell prompt offset is set, initialize it from the current prompt line count. + if ($null -eq $env:ATUIN_POWERSHELL_PROMPT_OFFSET) { + try { + $promptLines = (& $Function:prompt | Out-String | Measure-Object -Line).Lines + $env:ATUIN_POWERSHELL_PROMPT_OFFSET = -1 * ($promptLines - 1) + } + catch { + $env:ATUIN_POWERSHELL_PROMPT_OFFSET = 0 + } + } + + # PSReadLine maintains its own cursor position, which will no longer be valid if Atuin scrolls the display in inline mode. + # Fortunately, InvokePrompt can receive a new Y position and reset the internal state. + $y = $Host.UI.RawUI.CursorPosition.Y + [int]$env:ATUIN_POWERSHELL_PROMPT_OFFSET + $y = [System.Math]::Max([System.Math]::Min($y, [System.Console]::BufferHeight - 1), 0) + [Microsoft.PowerShell.PSConsoleReadLine]::InvokePrompt($null, $y) + + if ($suggestion -eq "") { + # The previous input was already rendered by InvokePrompt + return + } + + $acceptPrefix = "__atuin_accept__:" + + if ( $suggestion.StartsWith($acceptPrefix)) { + Set-CommandLine $suggestion.Substring($acceptPrefix.Length) + [Microsoft.PowerShell.PSConsoleReadLine]::AcceptLine() + } else { + Set-CommandLine $suggestion + } + } + finally { + [System.Console]::OutputEncoding = $previousOutputEncoding + Remove-Item $resultFile + } + } + + function Enable-AtuinSearchKeys { + param([bool]$CtrlR = $true, [bool]$UpArrow = $true) + + if ($CtrlR) { + Set-PSReadLineKeyHandler -Chord "Ctrl+r" -BriefDescription "Runs Atuin search" -ScriptBlock { + Invoke-AtuinSearch + } + } + + if ($UpArrow) { + Set-PSReadLineKeyHandler -Chord "UpArrow" -BriefDescription "Runs Atuin search" -ScriptBlock { + $line = Get-CommandLine + + if (!$line.Contains("`n")) { + Invoke-AtuinSearch -ExtraArgs "--shell-up-key-binding" + } else { + [Microsoft.PowerShell.PSConsoleReadLine]::PreviousLine() + } + } + } + } + + $ExecutionContext.SessionState.Module.OnRemove += { + $env:ATUIN_SESSION = $null + $Function:PSConsoleHostReadLine = $script:previousPSConsoleHostReadLine + } + + Export-ModuleMember -Function @("Enable-AtuinSearchKeys", "PSConsoleHostReadLine") +} | Import-Module -Global diff --git a/crates/atuin/src/shell/atuin.xsh b/crates/atuin/src/shell/atuin.xsh new file mode 100644 index 00000000..9f45cc6d --- /dev/null +++ b/crates/atuin/src/shell/atuin.xsh @@ -0,0 +1,85 @@ +import subprocess + +from prompt_toolkit.application.current import get_app +from prompt_toolkit.filters import Condition +from prompt_toolkit.keys import Keys + + +if "ATUIN_SESSION" not in ${...} or ${...}.get("ATUIN_SHLVL", "") != ${...}.get("SHLVL", ""): + $ATUIN_SESSION=$(atuin uuid).rstrip('\n') + $ATUIN_SHLVL = ${...}.get("SHLVL", "") + +@events.on_precommand +def _atuin_precommand(cmd: str): + cmd = cmd.rstrip("\n") + try: + $ATUIN_HISTORY_ID = $(atuin history start -- @(cmd) 2>/dev/null).rstrip("\n") + except: + $ATUIN_HISTORY_ID = "" + + +@events.on_postcommand +def _atuin_postcommand(cmd: str, rtn: int, out, ts): + if "ATUIN_HISTORY_ID" not in ${...}: + return + + duration = ts[1] - ts[0] + # Duration is float representing seconds, but atuin expects integer of nanoseconds + nanos = round(duration * 10 ** 9) + with ${...}.swap(ATUIN_LOG="error"): + # This causes the entire .xonshrc to be re-executed, which is incredibly slow + # This happens when using a subshell and using output redirection at the same time + # For more details, see https://github.com/xonsh/xonsh/issues/5224 + # (atuin history end --exit @(rtn) -- $ATUIN_HISTORY_ID &) > /dev/null 2>&1 + atuin history end --exit @(rtn) --duration @(nanos) -- $ATUIN_HISTORY_ID > /dev/null 2>&1 + del $ATUIN_HISTORY_ID + + +def _search(event, extra_args: list[str]): + buffer = event.current_buffer + cmd = ["atuin", "search", "--interactive", *extra_args] + # We need to explicitly pass in xonsh env, in case user has set XDG_HOME or something else that matters + env = ${...}.detype() + env["ATUIN_SHELL"] = "xonsh" + env["ATUIN_QUERY"] = buffer.text + + p = subprocess.run(cmd, stderr=subprocess.PIPE, encoding="utf-8", env=env) + result = p.stderr.rstrip("\n") + # redraw prompt - necessary if atuin is configured to run inline, rather than fullscreen + event.cli.renderer.erase() + + if not result: + return + + buffer.reset() + if result.startswith("__atuin_accept__:"): + buffer.insert_text(result[17:]) + buffer.validate_and_handle() + else: + buffer.insert_text(result) + + +@events.on_ptk_create +def _custom_keybindings(bindings, **kw): + if _ATUIN_BIND_CTRL_R: + @bindings.add(Keys.ControlR) + def r_search(event): + _search(event, extra_args=[]) + + if _ATUIN_BIND_UP_ARROW: + @Condition + def should_search(): + buffer = get_app().current_buffer + # disable keybind when there is an active completion, so + # that up arrow can be used to navigate completion menu + if buffer.complete_state is not None: + return False + # similarly, disable when buffer text contains multiple lines + if '\n' in buffer.text: + return False + + return True + + @bindings.add(Keys.Up, filter=should_search) + def up_search(event): + _search(event, extra_args=["--shell-up-key-binding"]) diff --git a/crates/atuin/src/shell/atuin.zsh b/crates/atuin/src/shell/atuin.zsh new file mode 100644 index 00000000..8e9b975c --- /dev/null +++ b/crates/atuin/src/shell/atuin.zsh @@ -0,0 +1,176 @@ +# shellcheck disable=SC2034,SC2153,SC2086,SC2155 + +# Above line is because shellcheck doesn't support zsh, per +# https://github.com/koalaman/shellcheck/wiki/SC1071, and the ignore: param in +# ludeeus/action-shellcheck only supports _directories_, not _files_. So +# instead, we manually add any error the shellcheck step finds in the file to +# the above line ... + +# Source this in your ~/.zshrc +autoload -U add-zsh-hook + +zmodload zsh/datetime 2>/dev/null + +# If zsh-autosuggestions is installed, configure it to use Atuin's search. If +# you'd like to override this, then add your config after the $(atuin init zsh) +# in your .zshrc +_zsh_autosuggest_strategy_atuin() { + # silence errors, since we don't want to spam the terminal prompt while typing. + suggestion=$(ATUIN_QUERY="$1" atuin search --cmd-only --limit 1 --search-mode prefix 2>/dev/null) +} + +if [ -n "${ZSH_AUTOSUGGEST_STRATEGY:-}" ]; then + ZSH_AUTOSUGGEST_STRATEGY=("atuin" "${ZSH_AUTOSUGGEST_STRATEGY[@]}") +else + ZSH_AUTOSUGGEST_STRATEGY=("atuin") +fi + +if [[ -z "${ATUIN_SESSION:-}" || "${ATUIN_SHLVL:-}" != "$SHLVL" ]]; then + export ATUIN_SESSION=$(atuin uuid) + export ATUIN_SHLVL=$SHLVL +fi +ATUIN_HISTORY_ID="" + +_atuin_preexec() { + local id + id=$(atuin history start -- "$1" 2>/dev/null) + export ATUIN_HISTORY_ID="$id" + __atuin_preexec_time=${EPOCHREALTIME-} +} + +_atuin_precmd() { + local EXIT="$?" __atuin_precmd_time=${EPOCHREALTIME-} + + [[ -z "${ATUIN_HISTORY_ID:-}" ]] && return + + local duration="" + if [[ -n $__atuin_preexec_time && -n $__atuin_precmd_time ]]; then + printf -v duration %.0f $(((__atuin_precmd_time - __atuin_preexec_time) * 1000000000)) + fi + + (ATUIN_LOG=error atuin history end --exit $EXIT ${duration:+--duration=$duration} -- $ATUIN_HISTORY_ID &) >/dev/null 2>&1 + export ATUIN_HISTORY_ID="" +} + +# Check if tmux popup is available (tmux >= 3.2) +__atuin_tmux_popup_check() { + [[ -n "${TMUX-}" ]] || return 1 + [[ "${ATUIN_TMUX_POPUP:-true}" != "false" ]] || return 1 + + # https://github.com/tmux/tmux/wiki/FAQ#how-often-is-tmux-released-what-is-the-version-number-scheme + local tmux_version + tmux_version=$(tmux -V 2>/dev/null | sed -n 's/^[^0-9]*\([0-9][0-9]*\.[0-9][0-9]*\).*/\1/p') # Could have used grep... + [[ -z "$tmux_version" ]] && return 1 + + local m1 m2 + m1=${tmux_version%%.*} + m2=${tmux_version#*.} + m2=${m2%%.*} + [[ "$m1" =~ ^[0-9]+$ ]] || return 1 + [[ "$m2" =~ ^[0-9]+$ ]] || m2=0 + (( m1 > 3 || (m1 == 3 && m2 >= 2) )) +} + +# Use global variable to fix scope issues with traps +__atuin_popup_tmpdir="" +__atuin_tmux_popup_cleanup() { + [[ -n "$__atuin_popup_tmpdir" && -d "$__atuin_popup_tmpdir" ]] && command rm -rf "$__atuin_popup_tmpdir" + __atuin_popup_tmpdir="" +} + +__atuin_search_cmd() { + local -a search_args=("$@") + + if __atuin_tmux_popup_check; then + __atuin_popup_tmpdir=$(mktemp -d) || return 1 + local result_file="$__atuin_popup_tmpdir/result" + + trap '__atuin_tmux_popup_cleanup' EXIT HUP INT TERM + + local escaped_query escaped_args + escaped_query=$(printf '%s' "$BUFFER" | sed "s/'/'\\\\''/g") + escaped_args="" + for arg in "${search_args[@]}"; do + escaped_args+=" '$(printf '%s' "$arg" | sed "s/'/'\\\\''/g")'" + done + + # In the popup, atuin goes to terminal, stderr goes to file + local cdir popup_width popup_height + cdir=$(pwd) + popup_width="${ATUIN_TMUX_POPUP_WIDTH:-80%}" # Keep default value anyways + popup_height="${ATUIN_TMUX_POPUP_HEIGHT:-60%}" + tmux display-popup -d "$cdir" -w "$popup_width" -h "$popup_height" -E -E -- \ + sh -c "PATH='$PATH' ATUIN_SESSION='$ATUIN_SESSION' ATUIN_SHELL=zsh ATUIN_LOG=error ATUIN_QUERY='$escaped_query' atuin search $escaped_args -i 2>'$result_file'" + + if [[ -f "$result_file" ]]; then + cat "$result_file" + fi + + __atuin_tmux_popup_cleanup + trap - EXIT HUP INT TERM + else + ATUIN_SHELL=zsh ATUIN_LOG=error ATUIN_QUERY=$BUFFER atuin search "${search_args[@]}" -i 3>&1 1>&2 2>&3 + fi +} + +_atuin_search() { + emulate -L zsh + zle -I + + # swap stderr and stdout, so that the tui stuff works + # TODO: not this + local output + # shellcheck disable=SC2048 + output=$(__atuin_search_cmd $*) + + zle reset-prompt + # re-enable bracketed paste + # shellcheck disable=SC2154 + echo -n ${zle_bracketed_paste[1]} >/dev/tty + + if [[ -n $output ]]; then + RBUFFER="" + LBUFFER=$output + + if [[ $LBUFFER == __atuin_accept__:* ]] + then + LBUFFER=${LBUFFER#__atuin_accept__:} + zle accept-line + fi + fi +} +_atuin_search_vicmd() { + _atuin_search --keymap-mode=vim-normal +} +_atuin_search_viins() { + _atuin_search --keymap-mode=vim-insert +} + +_atuin_up_search() { + # Only trigger if the buffer is a single line + if [[ ! $BUFFER == *$'\n'* ]]; then + _atuin_search --shell-up-key-binding "$@" + else + zle up-line + fi +} +_atuin_up_search_vicmd() { + _atuin_up_search --keymap-mode=vim-normal +} +_atuin_up_search_viins() { + _atuin_up_search --keymap-mode=vim-insert +} + +add-zsh-hook preexec _atuin_preexec +add-zsh-hook precmd _atuin_precmd + +zle -N atuin-search _atuin_search +zle -N atuin-search-vicmd _atuin_search_vicmd +zle -N atuin-search-viins _atuin_search_viins +zle -N atuin-up-search _atuin_up_search +zle -N atuin-up-search-vicmd _atuin_up_search_vicmd +zle -N atuin-up-search-viins _atuin_up_search_viins + +# These are compatibility widget names for "atuin <= 17.2.1" users. +zle -N _atuin_search_widget _atuin_search +zle -N _atuin_up_search_widget _atuin_up_search diff --git a/crates/atuin/src/sync.rs b/crates/atuin/src/sync.rs new file mode 100644 index 00000000..26004130 --- /dev/null +++ b/crates/atuin/src/sync.rs @@ -0,0 +1,50 @@ +use atuin_dotfiles::store::{AliasStore, var::VarStore}; +use atuin_scripts::store::ScriptStore; +use eyre::{Context, Result}; + +use atuin_client::{ + database::Database, history::store::HistoryStore, record::sqlite_store::SqliteStore, + settings::Settings, +}; +use atuin_common::record::RecordId; +use atuin_kv::store::KvStore; + +// This is the only crate that ties together all other crates. +// Therefore, it's the only crate where functions tying together all stores can live + +/// Rebuild all stores after a sync +/// Note: for history, this only does an _incremental_ sync. Hence the need to specify downloaded +/// records. +pub async fn build( + settings: &Settings, + store: &SqliteStore, + db: &dyn Database, + downloaded: Option<&[RecordId]>, +) -> Result<()> { + let encryption_key: [u8; 32] = atuin_client::encryption::load_key(settings) + .context("could not load encryption key")? + .into(); + + let host_id = Settings::host_id().await?; + + let downloaded = downloaded.unwrap_or(&[]); + + let kv_db = atuin_kv::database::Database::new(settings.kv.db_path.clone(), 1.0).await?; + + let history_store = HistoryStore::new(store.clone(), host_id, encryption_key); + let alias_store = AliasStore::new(store.clone(), host_id, encryption_key); + let var_store = VarStore::new(store.clone(), host_id, encryption_key); + let kv_store = KvStore::new(store.clone(), kv_db, host_id, encryption_key); + let script_store = ScriptStore::new(store.clone(), host_id, encryption_key); + + history_store.incremental_build(db, downloaded).await?; + + alias_store.build().await?; + var_store.build().await?; + kv_store.build().await?; + + let script_db = + atuin_scripts::database::Database::new(settings.scripts.db_path.clone(), 1.0).await?; + script_store.build(script_db).await?; + Ok(()) +} diff --git a/crates/atuin/tests/common/mod.rs b/crates/atuin/tests/common/mod.rs new file mode 100644 index 00000000..0a7c5275 --- /dev/null +++ b/crates/atuin/tests/common/mod.rs @@ -0,0 +1,117 @@ +use std::{env, time::Duration}; + +use atuin_client::api_client; +use atuin_common::utils::uuid_v7; +use atuin_server::{Settings as ServerSettings, launch_with_tcp_listener}; +use atuin_server_database::DbSettings; +use atuin_server_postgres::Postgres; +use futures_util::TryFutureExt; +use tokio::{net::TcpListener, sync::oneshot, task::JoinHandle}; +use tracing::{Dispatch, dispatcher}; +use tracing_subscriber::{EnvFilter, layer::SubscriberExt}; + +pub async fn start_server(path: &str) -> (String, oneshot::Sender<()>, JoinHandle<()>) { + let formatting_layer = tracing_tree::HierarchicalLayer::default() + .with_writer(tracing_subscriber::fmt::TestWriter::new()) + .with_indent_lines(true) + .with_ansi(true) + .with_targets(true) + .with_indent_amount(2); + + let dispatch: Dispatch = tracing_subscriber::registry() + .with(formatting_layer) + .with(EnvFilter::new("atuin_server=debug,atuin_client=debug,info")) + .into(); + + let db_uri = env::var("ATUIN_DB_URI") + .unwrap_or_else(|_| "postgres://atuin:pass@localhost:5432/atuin".to_owned()); + + let server_settings = ServerSettings { + host: "127.0.0.1".to_owned(), + port: 0, + path: path.to_owned(), + sync_v1_enabled: true, + open_registration: true, + max_history_length: 8192, + max_record_size: 1024 * 1024 * 1024, + page_size: 1100, + register_webhook_url: None, + register_webhook_username: String::new(), + db_settings: DbSettings { + db_uri: db_uri, + read_db_uri: None, + }, + metrics: atuin_server::settings::Metrics::default(), + fake_version: None, + }; + + let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel(); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + let _tracing_guard = dispatcher::set_default(&dispatch); + + if let Err(e) = launch_with_tcp_listener::<Postgres>( + server_settings, + listener, + shutdown_rx.unwrap_or_else(|_| ()), + ) + .await + { + tracing::error!(error=?e, "server error"); + panic!("error running server: {e:?}"); + } + }); + + // let the server come online + tokio::time::sleep(Duration::from_millis(200)).await; + + (format!("http://{addr}{path}"), shutdown_tx, server) +} + +pub async fn register_inner<'a>( + address: &'a str, + username: &str, + password: &str, +) -> api_client::Client<'a> { + let email = format!("{}@example.com", uuid_v7().as_simple()); + + // registration works + let registration_response = api_client::register(address, username, &email, password) + .await + .unwrap(); + + api_client::Client::new( + address, + api_client::AuthToken::Token(registration_response.session), + 5, + 30, + ) + .unwrap() +} + +#[allow(dead_code)] +pub async fn login(address: &str, username: String, password: String) -> api_client::Client<'_> { + // registration works + let login_response = api_client::login( + address, + atuin_common::api::LoginRequest { username, password }, + ) + .await + .unwrap(); + + api_client::Client::new( + address, + api_client::AuthToken::Token(login_response.session), + 5, + 30, + ) + .unwrap() +} + +#[allow(dead_code)] +pub async fn register(address: &str) -> api_client::Client<'_> { + let username = uuid_v7().as_simple().to_string(); + let password = uuid_v7().as_simple().to_string(); + register_inner(address, &username, &password).await +} diff --git a/crates/atuin/tests/sync.rs b/crates/atuin/tests/sync.rs new file mode 100644 index 00000000..7e25d1c2 --- /dev/null +++ b/crates/atuin/tests/sync.rs @@ -0,0 +1,45 @@ +use atuin_common::{api::AddHistoryRequest, utils::uuid_v7}; +use time::OffsetDateTime; + +mod common; + +#[tokio::test] +async fn sync() { + let path = format!("/{}", uuid_v7().as_simple()); + let (address, shutdown, server) = common::start_server(&path).await; + + let client = common::register(&address).await; + let hostname = uuid_v7().as_simple().to_string(); + let now = OffsetDateTime::now_utc(); + + let data1 = uuid_v7().as_simple().to_string(); + let data2 = uuid_v7().as_simple().to_string(); + + client + .post_history(&[ + AddHistoryRequest { + id: uuid_v7().as_simple().to_string(), + timestamp: now, + data: data1.clone(), + hostname: hostname.clone(), + }, + AddHistoryRequest { + id: uuid_v7().as_simple().to_string(), + timestamp: now, + data: data2.clone(), + hostname: hostname.clone(), + }, + ]) + .await + .unwrap(); + + let history = client + .get_history(OffsetDateTime::UNIX_EPOCH, OffsetDateTime::UNIX_EPOCH, None) + .await + .unwrap(); + + assert_eq!(history.history, vec![data1, data2]); + + shutdown.send(()).unwrap(); + server.await.unwrap(); +} diff --git a/crates/atuin/tests/users.rs b/crates/atuin/tests/users.rs new file mode 100644 index 00000000..95fb533b --- /dev/null +++ b/crates/atuin/tests/users.rs @@ -0,0 +1,121 @@ +use atuin_common::utils::uuid_v7; + +mod common; + +#[tokio::test] +async fn registration() { + let path = format!("/{}", uuid_v7().as_simple()); + let (address, shutdown, server) = common::start_server(&path).await; + dbg!(&address); + + // -- REGISTRATION -- + + let username = uuid_v7().as_simple().to_string(); + let password = uuid_v7().as_simple().to_string(); + let client = common::register_inner(&address, &username, &password).await; + + // the session token works + let status = client.status().await.unwrap(); + assert_eq!(status.username, username); + + // -- LOGIN -- + + let client = common::login(&address, username.clone(), password).await; + + // the session token works + let status = client.status().await.unwrap(); + assert_eq!(status.username, username); + + shutdown.send(()).unwrap(); + server.await.unwrap(); +} + +#[tokio::test] +async fn change_password() { + let path = format!("/{}", uuid_v7().as_simple()); + let (address, shutdown, server) = common::start_server(&path).await; + + // -- REGISTRATION -- + + let username = uuid_v7().as_simple().to_string(); + let password = uuid_v7().as_simple().to_string(); + let client = common::register_inner(&address, &username, &password).await; + + // the session token works + let status = client.status().await.unwrap(); + assert_eq!(status.username, username); + + // -- PASSWORD CHANGE -- + + let current_password = password; + let new_password = uuid_v7().as_simple().to_string(); + let result = client + .change_password(current_password, new_password.clone()) + .await; + + // the password change request succeeded + assert!(result.is_ok()); + + // -- LOGIN -- + + let client = common::login(&address, username.clone(), new_password).await; + + // login with new password yields a working token + let status = client.status().await.unwrap(); + assert_eq!(status.username, username); + + shutdown.send(()).unwrap(); + server.await.unwrap(); +} + +#[tokio::test] +async fn multi_user_test() { + let path = format!("/{}", uuid_v7().as_simple()); + let (address, shutdown, server) = common::start_server(&path).await; + dbg!(&address); + + // -- REGISTRATION -- + + let user_one = uuid_v7().as_simple().to_string(); + let password_one = uuid_v7().as_simple().to_string(); + let client_one = common::register_inner(&address, &user_one, &password_one).await; + + // the session token works + let status = client_one.status().await.unwrap(); + assert_eq!(status.username, user_one); + + let user_two = uuid_v7().as_simple().to_string(); + let password_two = uuid_v7().as_simple().to_string(); + let client_two = common::register_inner(&address, &user_two, &password_two).await; + + // the session token works + let status = client_two.status().await.unwrap(); + assert_eq!(status.username, user_two); + + // check that we can change user one's password, and _this does not affect user two_ + + let current_password = password_one; + let new_password = uuid_v7().as_simple().to_string(); + let result = client_one + .change_password(current_password, new_password.clone()) + .await; + + // the password change request succeeded + assert!(result.is_ok()); + + // -- LOGIN -- + + let client_one = common::login(&address, user_one.clone(), new_password).await; + let client_two = common::login(&address, user_two.clone(), password_two).await; + + // login with new password yields a working token + let status = client_one.status().await.unwrap(); + assert_eq!(status.username, user_one); + assert_ne!(status.username, user_two); + + let status = client_two.status().await.unwrap(); + assert_eq!(status.username, user_two); + + shutdown.send(()).unwrap(); + server.await.unwrap(); +} |