aboutsummaryrefslogtreecommitdiffstats
path: root/crates/atuin-ai/src/commands.rs
blob: cdbc8f2d7eb0c4c3d7a423eec3ff2c561cc509bc (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
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};
pub mod init;
pub(crate) 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>,

        /// Use the hook mode
        #[arg(long, hide = true)]
        hook: bool,
    },
}

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,
            hook,
            args,
            ..
        } => {
            if settings.logs.ai_enabled() {
                init_logging(settings, args.verbose)?;
            }

            inline::run(command, args.api_endpoint, args.api_token, settings, hook).await
        }
    }
}

pub(crate) 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);
        }
    }
}