feat: Add `atuin setup` (#3257)

8 files changed, 174 insertions, 23 deletions

diff --git a/Cargo.lock b/Cargo.lock
index 368f4c3f..88180998 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -265,6 +265,7 @@ dependencies = [
  "time",
  "tiny-bip39",
  "tokio",
+ "toml_edit",
  "tracing",
  "tracing-appender",
  "tracing-subscriber",
@@ -5335,7 +5336,7 @@ checksum = "cf92845e79fc2e2def6a5d828f0801e29a2f8acc037becc5ab08595c7d5e9863"
 dependencies = [
  "serde_core",
  "serde_spanned",
- "toml_datetime",
+ "toml_datetime 0.7.5+spec-1.1.0",
  "toml_parser",
  "winnow",
 ]
@@ -5350,6 +5351,28 @@ dependencies = [
 ]
 
 [[package]]
+name = "toml_datetime"
+version = "1.0.0+spec-1.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "32c2555c699578a4f59f0cc68e5116c8d7cabbd45e1409b989d4be085b53f13e"
+dependencies = [
+ "serde_core",
+]
+
+[[package]]
+name = "toml_edit"
+version = "0.25.4+spec-1.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7193cbd0ce53dc966037f54351dbbcf0d5a642c7f0038c382ef9e677ce8c13f2"
+dependencies = [
+ "indexmap 2.13.0",
+ "toml_datetime 1.0.0+spec-1.1.0",
+ "toml_parser",
+ "toml_writer",
+ "winnow",
+]
+
+[[package]]
 name = "toml_parser"
 version = "1.0.9+spec-1.1.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -5359,6 +5382,12 @@ dependencies = [
 ]
 
 [[package]]
+name = "toml_writer"
+version = "1.0.6+spec-1.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ab16f14aed21ee8bfd8ec22513f7287cd4a91aa92e44edfe2c17ddd004e92607"
+
+[[package]]
 name = "tonic"
 version = "0.14.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
diff --git a/crates/atuin-ai/src/commands/init.rs b/crates/atuin-ai/src/commands/init.rs
index caf4c8d9..6b23e936 100644
--- a/crates/atuin-ai/src/commands/init.rs
+++ b/crates/atuin-ai/src/commands/init.rs
@@ -43,7 +43,10 @@ _atuin_ai_question_mark() {
         # Clean up the inline viewport
         _atuin_ai_cleanup
 
-        if [[ $output == __atuin_ai_cancel__ ]]; then
+        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=""
@@ -86,8 +89,11 @@ _atuin_ai_question_mark() {
         local output
         output=$(atuin ai inline --hook 3>&1 1>&2 2>&3)
 
-        if [[ $output == __atuin_ai_cancel__ ]]; then
-            # User cancelled, do nothing
+        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
@@ -145,8 +151,10 @@ function _atuin_ai_question_mark
         # Run atuin ai inline, swapping stdout and stderr
         set -l output (atuin ai inline --hook 3>&1 1>&2 2>&3 | string collect)
 
-        if test "$output" = "__atuin_ai_cancel__"
-            # User cancelled, do nothing
+        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
@@ -188,6 +196,7 @@ mod tests {
         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__"));
@@ -200,6 +209,7 @@ mod tests {
         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__"));
@@ -212,6 +222,7 @@ mod tests {
         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
index 803c7d72..ce566be1 100644
--- a/crates/atuin-ai/src/commands/inline.rs
+++ b/crates/atuin-ai/src/commands/inline.rs
@@ -26,6 +26,17 @@ pub async fn run(
     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();
 
@@ -62,7 +73,7 @@ pub async fn run(
         settings,
     )
     .await?;
-    emit_shell_result(action.0, &action.1, output_for_hook);
+    emit_shell_result(action, output_for_hook);
 
     Ok(())
 }
@@ -326,10 +337,11 @@ fn detect_os() -> String {
     }
 }
 
-#[derive(Clone, Copy)]
+#[derive(Clone)]
 enum Action {
-    Execute,
-    Insert,
+    Execute(String),
+    Insert(String),
+    Print(String),
     Cancel,
 }
 
@@ -432,7 +444,7 @@ async fn run_inline_tui(
     keep_output: bool,
     debug_state_file: Option<String>,
     settings: &atuin_client::settings::Settings,
-) -> Result<(Action, String)> {
+) -> 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();
@@ -686,9 +698,9 @@ async fn run_inline_tui(
 
     // 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, String::new()),
+        Some(ExitAction::Execute(cmd)) => Action::Execute(cmd),
+        Some(ExitAction::Insert(cmd)) => Action::Insert(cmd),
+        _ => Action::Cancel,
     };
 
     Ok(result)
@@ -702,17 +714,19 @@ impl Drop for RawModeGuard {
     }
 }
 
-fn emit_shell_result(action: Action, command: &str, output_for_hook: bool) {
+fn emit_shell_result(action: Action, output_for_hook: bool) {
     if output_for_hook {
         match action {
-            Action::Execute => eprintln!("__atuin_ai_execute__:{command}"),
-            Action::Insert => eprintln!("__atuin_ai_insert__:{command}"),
+            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 => eprintln!("{command}"),
-            Action::Insert => eprintln!("{command}"),
+            Action::Execute(output) => eprintln!("{output}"),
+            Action::Insert(output) => eprintln!("{output}"),
+            Action::Print(output) => eprintln!("{output}"),
             Action::Cancel => eprintln!(),
         }
     }
diff --git a/crates/atuin-client/src/settings.rs b/crates/atuin-client/src/settings.rs
index 62b3a098..2a96a2b3 100644
--- a/crates/atuin-client/src/settings.rs
+++ b/crates/atuin-client/src/settings.rs
@@ -594,6 +594,9 @@ pub struct Logs {
 
 #[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>,
@@ -1433,6 +1436,8 @@ impl Settings {
             .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![
@@ -1463,7 +1468,7 @@ impl Settings {
             ))
     }
 
-    pub fn new() -> Result<Self> {
+    pub fn get_config_path() -> Result<PathBuf> {
         let config_dir = atuin_common::utils::config_dir();
 
         create_dir_all(&config_dir)
@@ -1479,6 +1484,12 @@ impl Settings {
 
         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)]
diff --git a/crates/atuin/Cargo.toml b/crates/atuin/Cargo.toml
index a7bd3330..5924302b 100644
--- a/crates/atuin/Cargo.toml
+++ b/crates/atuin/Cargo.toml
@@ -87,10 +87,13 @@ 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" }
+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 }
 
diff --git a/crates/atuin/src/command/client.rs b/crates/atuin/src/command/client.rs
index 6e197604..02d64205 100644
--- a/crates/atuin/src/command/client.rs
+++ b/crates/atuin/src/command/client.rs
@@ -59,6 +59,7 @@ mod init;
 mod kv;
 mod scripts;
 mod search;
+mod setup;
 mod stats;
 mod store;
 mod wrapped;
@@ -66,6 +67,10 @@ 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),
@@ -333,6 +338,7 @@ impl Cmd {
         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,
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/docs/docs/ai/settings.md b/docs/docs/ai/settings.md
index 0a1f5b45..be27261f 100644
--- a/docs/docs/ai/settings.md
+++ b/docs/docs/ai/settings.md
@@ -2,6 +2,12 @@
 
 All the settings that control the behavior of [Atuin AI](./introduction.md) are specified in an `[ai]` section in your `config.toml`. See [the configuration documentation](../../configuration/config/) for more detailed information about Atuin's configuration system.
 
+### enabled
+
+Default: `false`
+
+Whether or not the AI feature are enabled. When set to `false`, the question mark keybinding will output a message with instructions to run `atuin setup` to enable the feature.
+
 ### send_cwd
 
 Default: `false`
@@ -19,10 +25,10 @@ send_cwd = true
 
 Default: `null`
 
-The address of the Atuin AI endpoint. Used for AI features like command generation. Only necessary for custom AI endpoints.
+The address of the Atuin AI endpoint. Used for AI features like command generation. Most users will not need this setting; it is only necessary for custom AI endpoints.
 
 ### api_token
 
 Default: `null`
 
-The API token for the Atuin AI endpoint. Used for AI features like command generation. Only necessary for custom AI endpoints.
+The API token for the Atuin AI endpoint. Used for AI features like command generation. Most users will not need this setting; it is only necessary for custom AI endpoints.