yapi_ai/
credentials.rs

1//! Credential values: `$VAR` and `${VAR}` templates, `!command` output, and the
2//! environment variables each provider reads.
3//!
4//! Ports of `resolve-config-value.ts` and `env-api-keys.ts` in pi `v1.0.0`.
5
6use std::collections::HashMap;
7use std::sync::{Mutex, OnceLock};
8
9use indexmap::IndexMap;
10
11/// Extra variables that take priority over the process environment.
12pub type ProviderEnv = IndexMap<String, String>;
13
14enum Part {
15    Literal(String),
16    Env(String),
17}
18
19fn is_name_start(c: char) -> bool {
20    c.is_ascii_alphabetic() || c == '_'
21}
22
23fn is_name(c: char) -> bool {
24    c.is_ascii_alphanumeric() || c == '_'
25}
26
27fn push_literal(parts: &mut Vec<Part>, text: &str) {
28    if text.is_empty() {
29        return;
30    }
31    if let Some(Part::Literal(previous)) = parts.last_mut() {
32        previous.push_str(text);
33    } else {
34        parts.push(Part::Literal(text.to_owned()));
35    }
36}
37
38fn parse_template(config: &str) -> Vec<Part> {
39    let mut parts = Vec::new();
40    let mut rest = config;
41    while let Some(dollar) = rest.find('$') {
42        push_literal(&mut parts, &rest[..dollar]);
43        let after = &rest[dollar + 1..];
44        match after.chars().next() {
45            Some(c @ ('$' | '!')) => {
46                push_literal(&mut parts, &c.to_string());
47                rest = &after[1..];
48            }
49            Some('{') => match after.find('}') {
50                None => {
51                    push_literal(&mut parts, "$");
52                    rest = after;
53                }
54                Some(end) => {
55                    let name = &after[1..end];
56                    let valid =
57                        name.chars().next().is_some_and(is_name_start) && name.chars().all(is_name);
58                    if valid {
59                        parts.push(Part::Env(name.to_owned()));
60                    } else {
61                        push_literal(&mut parts, &rest[dollar..dollar + end + 2]);
62                    }
63                    rest = &after[end + 1..];
64                }
65            },
66            Some(c) if is_name_start(c) => {
67                let length = after.find(|c: char| !is_name(c)).unwrap_or(after.len());
68                parts.push(Part::Env(after[..length].to_owned()));
69                rest = &after[length..];
70            }
71            _ => {
72                push_literal(&mut parts, "$");
73                rest = after;
74            }
75        }
76    }
77    push_literal(&mut parts, rest);
78    parts
79}
80
81/// A provider setting: `env`'s value, else the process environment's. Empty
82/// values count as unset. Port of `getProviderEnvValue`.
83pub fn provider_env_value(name: &str, env: Option<&ProviderEnv>) -> Option<String> {
84    env.and_then(|env| env.get(name))
85        .filter(|value| !value.is_empty())
86        .cloned()
87        .or_else(|| std::env::var(name).ok().filter(|value| !value.is_empty()))
88}
89
90/// Whether the value runs a command.
91pub fn is_command(config: &str) -> bool {
92    config.starts_with('!')
93}
94
95/// The environment variables a template value reads.
96pub fn env_var_names(config: &str) -> Vec<String> {
97    if is_command(config) {
98        return Vec::new();
99    }
100    let mut names = Vec::new();
101    for part in parse_template(config) {
102        if let Part::Env(name) = part
103            && !names.contains(&name)
104        {
105            names.push(name);
106        }
107    }
108    names
109}
110
111/// The variables the value reads that are not set.
112pub fn missing_env_vars(config: &str, env: Option<&ProviderEnv>) -> Vec<String> {
113    env_var_names(config)
114        .into_iter()
115        .filter(|name| provider_env_value(name, env).is_none())
116        .collect()
117}
118
119/// pi's `resolveConfigValueOrThrow` message for a value that did not
120/// resolve, describing it as `description`.
121pub fn unresolved_message(config: &str, description: &str, env: Option<&ProviderEnv>) -> String {
122    if let Some(command) = config.strip_prefix('!') {
123        return format!("Failed to resolve {description} from shell command: {command}");
124    }
125    match missing_env_vars(config, env).as_slice() {
126        [] => format!("Failed to resolve {description}"),
127        [name] => format!("Failed to resolve {description} from environment variable: {name}"),
128        names => format!(
129            "Failed to resolve {description} from environment variables: {}",
130            names.join(", ")
131        ),
132    }
133}
134
135/// Whether every variable the value reads is set.
136pub fn is_configured(config: &str, env: Option<&ProviderEnv>) -> bool {
137    env_var_names(config)
138        .iter()
139        .all(|name| provider_env_value(name, env).is_some())
140}
141
142/// Resolves a value: `!command` runs through the shell (cached when `cache`),
143/// anything else is a template whose variables must all be set.
144pub async fn resolve(config: &str, env: Option<&ProviderEnv>, cache: bool) -> Option<String> {
145    if let Some(command) = config.strip_prefix('!') {
146        return run_command(command, cache).await;
147    }
148    resolve_template(config, env)
149}
150
151/// Resolves a template value, or returns `None` for a `!command` or when a
152/// variable it reads is not set.
153pub fn resolve_template(config: &str, env: Option<&ProviderEnv>) -> Option<String> {
154    if is_command(config) {
155        return None;
156    }
157    let mut resolved = String::new();
158    for part in parse_template(config) {
159        match part {
160            Part::Literal(text) => resolved.push_str(&text),
161            Part::Env(name) => resolved.push_str(&provider_env_value(&name, env)?),
162        }
163    }
164    Some(resolved)
165}
166
167async fn run_command(command: &str, cache: bool) -> Option<String> {
168    static CACHE: OnceLock<Mutex<HashMap<String, Option<String>>>> = OnceLock::new();
169    let cache_map = CACHE.get_or_init(Default::default);
170    if cache
171        && let Some(value) = cache_map
172            .lock()
173            .ok()
174            .and_then(|map| map.get(command).cloned())
175    {
176        return value;
177    }
178    let output = tokio::time::timeout(
179        std::time::Duration::from_secs(10),
180        tokio::process::Command::new("sh")
181            .arg("-c")
182            .arg(command)
183            .stdin(std::process::Stdio::null())
184            .stderr(std::process::Stdio::null())
185            .kill_on_drop(true)
186            .output(),
187    )
188    .await;
189    let value = match output {
190        Ok(Ok(output)) if output.status.success() => {
191            let text = String::from_utf8_lossy(&output.stdout).trim().to_owned();
192            (!text.is_empty()).then_some(text)
193        }
194        _ => None,
195    };
196    if cache && let Ok(mut map) = cache_map.lock() {
197        map.insert(command.to_owned(), value.clone());
198    }
199    value
200}
201
202/// The environment variables holding a provider's API key, in priority order.
203/// Anthropic's `ANTHROPIC_AUTH_TOKEN` is a bearer token, not a key (see
204/// [`BEARER_TOKEN_ENV`]).
205pub fn api_key_env_vars(provider: &str) -> &'static [&'static str] {
206    crate::providers::info(provider)
207        .and_then(|info| info.api_key)
208        .map_or(&[], |method| method.env)
209}
210
211/// The variable Anthropic reads as a bearer token rather than an API key.
212pub const BEARER_TOKEN_ENV: &str = "ANTHROPIC_AUTH_TOKEN";
213
214/// The key pi reports for providers authenticated by ambient credentials.
215pub const AMBIENT_CREDENTIALS: &str = "<authenticated>";
216
217/// Whether Google Application Default Credentials exist: the file named by
218/// `GOOGLE_APPLICATION_CREDENTIALS`, or gcloud's default file.
219fn has_vertex_adc_credentials(env: Option<&ProviderEnv>) -> bool {
220    match provider_env_value("GOOGLE_APPLICATION_CREDENTIALS", env) {
221        Some(path) => std::path::Path::new(&path).exists(),
222        // Node's `os.homedir()`: `HOME`, even when empty, else the user database.
223        None => std::env::var_os("HOME")
224            .map(std::path::PathBuf::from)
225            .or_else(std::env::home_dir)
226            .is_some_and(|home| {
227                home.join(".config/gcloud/application_default_credentials.json")
228                    .exists()
229            }),
230    }
231}
232
233/// Ambient credentials that authenticate a provider without an API key: AWS
234/// profiles, keys, tokens and roles for Bedrock, and Application Default
235/// Credentials with a project and location for Vertex. Returns the variable that
236/// makes it so.
237fn ambient_credentials(provider: &str, env: Option<&ProviderEnv>) -> Option<&'static str> {
238    let set = |name: &str| provider_env_value(name, env).is_some();
239    match provider {
240        "google-vertex" => (has_vertex_adc_credentials(env)
241            && (set("GOOGLE_CLOUD_PROJECT") || set("GCLOUD_PROJECT"))
242            && set("GOOGLE_CLOUD_LOCATION"))
243        .then_some("GOOGLE_CLOUD_LOCATION"),
244        "amazon-bedrock" => {
245            if set("AWS_PROFILE") {
246                Some("AWS_PROFILE")
247            } else if set("AWS_ACCESS_KEY_ID") && set("AWS_SECRET_ACCESS_KEY") {
248                Some("AWS_ACCESS_KEY_ID")
249            } else {
250                [
251                    "AWS_BEARER_TOKEN_BEDROCK",
252                    "AWS_CONTAINER_CREDENTIALS_RELATIVE_URI",
253                    "AWS_CONTAINER_CREDENTIALS_FULL_URI",
254                    "AWS_WEB_IDENTITY_TOKEN_FILE",
255                ]
256                .into_iter()
257                .find(|name| set(name))
258            }
259        }
260        _ => None,
261    }
262}
263
264/// The first set environment variable for a provider's key, with its name. For
265/// providers with ambient credentials (Bedrock, Vertex) the key is
266/// [`AMBIENT_CREDENTIALS`].
267pub fn env_api_key(provider: &str, env: Option<&ProviderEnv>) -> Option<(&'static str, String)> {
268    api_key_env_vars(provider)
269        .iter()
270        .find_map(|name| provider_env_value(name, env).map(|value| (*name, value)))
271        .or_else(|| {
272            ambient_credentials(provider, env).map(|name| (name, AMBIENT_CREDENTIALS.to_owned()))
273        })
274}
275
276#[cfg(test)]
277mod tests {
278    use super::*;
279
280    #[tokio::test]
281    async fn resolves_templates_and_commands() {
282        let mut env = ProviderEnv::new();
283        env.insert("YAPI_TEST_A".into(), "a".into());
284        assert_eq!(
285            resolve("x-$YAPI_TEST_A-${YAPI_TEST_A}$$!", Some(&env), false)
286                .await
287                .as_deref(),
288            Some("x-a-a$!")
289        );
290        assert_eq!(
291            resolve("$YAPI_TEST_UNSET_VALUE", Some(&env), false).await,
292            None
293        );
294        assert_eq!(
295            resolve("${not valid}", None, false).await.as_deref(),
296            Some("${not valid}")
297        );
298        assert_eq!(resolve("!echo hi", None, true).await.as_deref(), Some("hi"));
299        assert_eq!(resolve("!exit 1", None, false).await, None);
300        assert_eq!(env_var_names("$A and ${B} $A"), ["A", "B"]);
301        assert!(!is_configured("$YAPI_TEST_UNSET_VALUE", None));
302        let mut aws = ProviderEnv::new();
303        aws.insert("AWS_ACCESS_KEY_ID".into(), "id".into());
304        aws.insert("AWS_SECRET_ACCESS_KEY".into(), "secret".into());
305        assert_eq!(
306            env_api_key("amazon-bedrock", Some(&aws)).map(|(_, key)| key),
307            Some(AMBIENT_CREDENTIALS.to_owned())
308        );
309    }
310
311    #[cfg(unix)]
312    #[tokio::test(start_paused = true)]
313    async fn command_killed_on_timeout() {
314        let marker = std::env::temp_dir().join(format!("yapi-command-{}", std::process::id()));
315        let _ = std::fs::remove_file(&marker);
316        // The paused clock reaches the timeout while the command sleeps.
317        let command = format!("!sleep 1; touch '{}'", marker.display());
318        assert_eq!(resolve(&command, None, false).await, None);
319        std::thread::sleep(std::time::Duration::from_millis(1500));
320        assert!(!marker.exists());
321    }
322}
yapi_ai/credentials.rs

yapi_ai/
credentials.rs
