aboutsummaryrefslogtreecommitdiffstats
path: root/crates/atuin-ai/src/context.rs
blob: f891a9fc04631e4ca2ad14bdc205a4cc0de52310 (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
use std::path::PathBuf;
use std::sync::Arc;

use atuin_client::distro::detect_linux_distribution;
use atuin_client::history::History;
use atuin_client::settings::AiCapabilities;

/// Session-scoped context for the AI chat session.
/// Holds the API configuration and client settings needed by the event loop and stream task.
#[derive(Clone, Debug)]
pub(crate) struct AppContext {
    pub endpoint: String,
    pub token: String,
    pub send_cwd: bool,
    pub last_command: Option<History>,
    pub history_db: Arc<atuin_client::database::Sqlite>,
    /// Git root of the current working directory, if inside a git repo.
    /// Resolves through worktrees to the main repo root.
    pub git_root: Option<PathBuf>,
    pub capabilities: AiCapabilities,
    pub daemon_enabled: bool,
}

pub(crate) fn history_output_capability_available(daemon_enabled: bool) -> bool {
    cfg!(feature = "daemon") && daemon_enabled
}

impl AppContext {
    pub(crate) fn capabilities_as_strings(&self) -> Vec<String> {
        let mut caps = vec!["client_invocations".to_string()];
        if self.capabilities.enable_history_search.unwrap_or(true) {
            caps.push("client_v1_atuin_history".to_string());
        }
        if self.capabilities.enable_file_tools.unwrap_or(true) {
            caps.push("client_v1_read_file".to_string());
            caps.push("client_v1_edit_file".to_string());
            caps.push("client_v1_write_file".to_string());
        }
        if self.capabilities.enable_command_execution.unwrap_or(true) {
            caps.push("client_v1_execute_shell_command".to_string());
        }
        if history_output_capability_available(self.daemon_enabled)
            && self.capabilities.enable_history_output.unwrap_or(true)
        {
            caps.push("client_v1_atuin_output".to_string());
        }
        caps.push("client_v1_load_skill".to_string());
        if let Ok(extra) = std::env::var("ATUIN_AI__ADDITIONAL_CAPS") {
            caps.extend(
                extra
                    .split(',')
                    .map(|s| s.trim().to_string())
                    .filter(|s| !s.is_empty()),
            );
        }
        caps
    }
}

/// Machine identity — computed once per session.
#[derive(Clone, Debug)]
pub(crate) struct ClientContext {
    pub os: String,
    pub shell: Option<String>,
    pub distro: Option<String>,
}

impl ClientContext {
    pub(crate) fn detect() -> Self {
        let os = detect_os();
        let shell = crate::commands::detect_shell();
        let distro = if os == "linux" {
            Some(detect_linux_distribution())
        } else {
            None
        };
        Self { os, shell, distro }
    }

    /// Serialize to the JSON format the API expects for the "context" field.
    /// The `pwd` field is always dynamic (current working directory), so it's
    /// computed fresh on each call if `send_cwd` is true.
    pub(crate) fn to_json(
        &self,
        send_cwd: bool,
        last_command: Option<&History>,
    ) -> serde_json::Value {
        let mut ctx = serde_json::json!({
            "os": self.os,
            "shell": self.shell,
            "pwd": if send_cwd {
                std::env::current_dir().ok().map(|p| p.to_string_lossy().into_owned())
            } else {
                None
            },
        });

        if let Some(history) = last_command {
            ctx["last_command"] = serde_json::json!(crate::history_format::format_last_command(
                history,
                crate::history_format::current_local_offset(),
            ));
        }

        if let Some(ref distro) = self.distro {
            ctx["distro"] = serde_json::json!(distro);
        }

        ctx
    }
}

/// Move the `detect_os` function here since it's about client identity.
fn detect_os() -> String {
    match std::env::consts::OS {
        "macos" => "macos".to_string(),
        "linux" => "linux".to_string(),
        "windows" => "windows".to_string(),
        other => format!("Other: {other}"),
    }
}