Skip to main content

comindly_engine/
credentials.rs

1//! `~/.config/comindly/credentials`: your LLM API keys, a default and optionally one per team name,
2//! and the server `comindly login` signed you in to. Readable only by you and never shared;
3//! `OPENAI_API_KEY` is the fallback key.
4
5use std::path::PathBuf;
6
7use anyhow::{Context, Result};
8use serde::{Deserialize, Serialize};
9
10#[derive(Debug, Default, Serialize, Deserialize)]
11struct File {
12    #[serde(default, skip_serializing_if = "Option::is_none")]
13    llm_api_key: Option<String>,
14    #[serde(default, skip_serializing_if = "Vec::is_empty")]
15    team: Vec<Entry>,
16    #[serde(default, skip_serializing_if = "Option::is_none")]
17    server: Option<Login>,
18}
19
20/// A server and the personal access token `comindly login` received for it.
21#[derive(Debug, Clone, Serialize, Deserialize)]
22pub struct Login {
23    pub url: String,
24    pub token: String,
25    pub email: String,
26}
27
28/// The server you are signed in to, if any.
29pub fn login() -> Option<Login> {
30    read().server
31}
32
33/// Remember the server you signed in to, or forget it with `None`.
34pub fn set_login(login: Option<Login>) -> Result<()> {
35    let mut f = read();
36    f.server = login;
37    write(&f)
38}
39
40#[derive(Debug, Clone, Serialize, Deserialize)]
41struct Entry {
42    name: String,
43    llm_api_key: String,
44}
45
46pub fn path() -> PathBuf {
47    crate::config::config_dir().join("credentials")
48}
49
50fn read() -> File {
51    std::fs::read_to_string(path())
52        .ok()
53        .and_then(|t| toml::from_str(&t).ok())
54        .unwrap_or_default()
55}
56
57pub use comindly_api::credentials::Source;
58
59/// The API key for `team` (a team name): its entry, else the default, else `OPENAI_API_KEY`;
60/// on a server, the deployment's key or the one stored there.
61pub fn llm_key(
62    team: Option<&str>,
63    server: Option<&crate::server::Server>,
64) -> (Option<String>, Source) {
65    if let Some(server) = server {
66        // One an admin stored in the app replaces the deployment's.
67        if let Some(k) = server.secrets.secret(crate::server::LLM_SECRET) {
68            return (Some(k), Source::Server);
69        }
70        return match server.config.llm_key.clone() {
71            Some(k) => (Some(k), Source::Environment),
72            None => (None, Source::None),
73        };
74    }
75    let f = read();
76    if let Some(k) = team.and_then(|t| f.team.iter().find(|e| e.name == t)) {
77        return (Some(k.llm_api_key.clone()), Source::Team);
78    }
79    if let Some(k) = f.llm_api_key.filter(|k| !k.is_empty()) {
80        return (Some(k), Source::Default);
81    }
82    match std::env::var("OPENAI_API_KEY")
83        .ok()
84        .filter(|k| !k.is_empty())
85    {
86        Some(k) => (Some(k), Source::Environment),
87        None => (None, Source::None),
88    }
89}
90
91/// Store `key` for `team`, or as the default without one; an empty key removes it.
92pub fn set_llm_key(team: Option<&str>, key: &str) -> Result<()> {
93    let mut f = read();
94    let key = key.trim();
95    match team {
96        Some(t) => {
97            f.team.retain(|e| e.name != t);
98            if !key.is_empty() {
99                f.team.push(Entry {
100                    name: t.to_string(),
101                    llm_api_key: key.to_string(),
102                });
103            }
104        }
105        None => f.llm_api_key = (!key.is_empty()).then(|| key.to_string()),
106    }
107    write(&f)
108}
109
110fn write(f: &File) -> Result<()> {
111    let p = path();
112    std::fs::create_dir_all(p.parent().context("config dir")?)?;
113    let text = format!(
114        "# comindly API keys and server sign-in. Keep this file private.\n\n{}",
115        toml::to_string(f)?
116    );
117    let tmp = p.with_extension("tmp");
118    let mut open = std::fs::OpenOptions::new();
119    open.write(true).create(true).truncate(true);
120    #[cfg(unix)]
121    std::os::unix::fs::OpenOptionsExt::mode(&mut open, 0o600);
122    std::io::Write::write_all(&mut open.open(&tmp)?, text.as_bytes())?;
123    std::fs::rename(&tmp, &p).with_context(|| format!("write {}", p.display()))
124}
125
126#[cfg(test)]
127mod tests {
128    use crate::config::{export, Config};
129    use comindly_core::pipeline::RepoEntry;
130
131    // Written by toml 0.8; later versions must read and write the same files.
132    const CREDENTIALS: &str = r#"llm_api_key = "k0"
133
134[[team]]
135name = "a"
136llm_api_key = "k1"
137
138[server]
139url = "https://x.example"
140token = "t"
141email = "e@example.com"
142"#;
143
144    const PREFERENCES: &str = r#"# comindly preferences: `comindly config import <file> --override` applies this file.
145
146index_dir = "s3://b/i"
147signing_key = "~/.ssh/k"
148cache_ttl_secs = 300
149author_name = "A \"q\" B"
150record_names = false
151
152[[repos]]
153path = "/r/a"
154remote = "git@h:a.git"
155
156[[repos]]
157path = "/r/b"
158"#;
159
160    #[test]
161    fn credentials_file_round_trips() {
162        let f: super::File = toml::from_str(CREDENTIALS).unwrap();
163        assert_eq!(f.team[0].name, "a");
164        assert_eq!(f.server.as_ref().unwrap().email, "e@example.com");
165        assert_eq!(toml::to_string(&f).unwrap(), CREDENTIALS);
166    }
167
168    #[test]
169    fn preferences_file_round_trips() {
170        let c: Config = toml::from_str(PREFERENCES).unwrap();
171        assert_eq!(
172            c,
173            Config {
174                index_dir: Some("s3://b/i".into()),
175                signing_key: Some("~/.ssh/k".into()),
176                allowed_signers: None,
177                cache_ttl_secs: Some(300),
178                author_name: Some("A \"q\" B".into()),
179                record_names: Some(false),
180                repos: vec![
181                    RepoEntry {
182                        path: "/r/a".into(),
183                        remote: Some("git@h:a.git".into()),
184                    },
185                    RepoEntry {
186                        path: "/r/b".into(),
187                        remote: None,
188                    },
189                ],
190            }
191        );
192        assert_eq!(toml::from_str::<Config>(&export(&c).unwrap()).unwrap(), c);
193    }
194}