agent_tools_interface/core/
auth_generator.rs

1//! Dynamic credential generator — produces short-lived auth tokens at call time.
2//!
3//! Generators run where secrets live: on the proxy server in proxy mode,
4//! on the local machine in local mode. Signing keys never enter the sandbox.
5//!
6//! Two generator types:
7//! - `command`: runs an external command, captures stdout
8//! - `script`: writes an inline script to a temp file, runs via interpreter
9//!
10//! Results are cached per (provider, agent_sub) with configurable TTL.
11
12use std::collections::HashMap;
13use std::sync::Mutex;
14use std::time::{Duration, Instant};
15use thiserror::Error;
16
17use crate::core::keyring::Keyring;
18use crate::core::manifest::{AuthGenType, AuthGenerator, AuthOutputFormat, Provider};
19
20// ---------------------------------------------------------------------------
21// Errors
22// ---------------------------------------------------------------------------
23
24#[derive(Error, Debug)]
25pub enum AuthGenError {
26    #[error("Auth generator config error: {0}")]
27    Config(String),
28    #[error("Failed to spawn generator process: {0}")]
29    Spawn(String),
30    #[error("Generator timed out after {0}s")]
31    Timeout(u64),
32    #[error("Generator exited with code {code}: {stderr}")]
33    NonZeroExit { code: i32, stderr: String },
34    #[error("Failed to parse generator output: {0}")]
35    OutputParse(String),
36    #[error("Keyring key '{0}' not found (required by auth_generator)")]
37    KeyringMissing(String),
38    #[error("IO error: {0}")]
39    Io(#[from] std::io::Error),
40}
41
42// ---------------------------------------------------------------------------
43// Context for variable expansion
44// ---------------------------------------------------------------------------
45
46/// Context for expanding `${VAR}` placeholders in generator args/env.
47pub struct GenContext {
48    pub jwt_sub: String,
49    pub jwt_scope: String,
50    pub tool_name: String,
51    pub timestamp: u64,
52}
53
54impl Default for GenContext {
55    fn default() -> Self {
56        GenContext {
57            jwt_sub: "dev".into(),
58            jwt_scope: "*".into(),
59            tool_name: String::new(),
60            timestamp: std::time::SystemTime::now()
61                .duration_since(std::time::UNIX_EPOCH)
62                .unwrap_or_default()
63                .as_secs(),
64        }
65    }
66}
67
68// ---------------------------------------------------------------------------
69// Generated credential
70// ---------------------------------------------------------------------------
71
72/// Result of running a generator — primary token + optional extra injections.
73#[derive(Debug, Clone)]
74pub struct GeneratedCredential {
75    /// Primary token value (used for bearer/header/query auth).
76    pub value: String,
77    /// Extra headers from JSON inject targets with type="header".
78    pub extra_headers: HashMap<String, String>,
79    /// Extra env vars from JSON inject targets with type="env".
80    pub extra_env: HashMap<String, String>,
81}
82
83// ---------------------------------------------------------------------------
84// Cache
85// ---------------------------------------------------------------------------
86
87struct CachedCredential {
88    cred: GeneratedCredential,
89    expires_at: Instant,
90}
91
92/// TTL-based credential cache, keyed by (provider_name, agent_sub).
93pub struct AuthCache {
94    entries: Mutex<HashMap<(String, String), CachedCredential>>,
95}
96
97impl Default for AuthCache {
98    fn default() -> Self {
99        AuthCache {
100            entries: Mutex::new(HashMap::new()),
101        }
102    }
103}
104
105impl AuthCache {
106    pub fn new() -> Self {
107        Self::default()
108    }
109
110    pub fn get(&self, provider: &str, sub: &str) -> Option<GeneratedCredential> {
111        let cache = self.entries.lock().unwrap();
112        let key = (provider.to_string(), sub.to_string());
113        match cache.get(&key) {
114            Some(entry) if Instant::now() < entry.expires_at => Some(entry.cred.clone()),
115            _ => None,
116        }
117    }
118
119    pub fn insert(&self, provider: &str, sub: &str, cred: GeneratedCredential, ttl_secs: u64) {
120        if ttl_secs == 0 {
121            return; // No caching
122        }
123        let mut cache = self.entries.lock().unwrap();
124        let key = (provider.to_string(), sub.to_string());
125        cache.insert(
126            key,
127            CachedCredential {
128                cred,
129                expires_at: Instant::now() + Duration::from_secs(ttl_secs),
130            },
131        );
132    }
133}
134
135// ---------------------------------------------------------------------------
136// Main generate function
137// ---------------------------------------------------------------------------
138
139/// Generate a credential by running the provider's auth_generator.
140///
141/// 1. Check cache → return if hit
142/// 2. Expand variables in args and env
143/// 3. Spawn subprocess (command or script)
144/// 4. Parse output (text or JSON)
145/// 5. Cache and return
146pub async fn generate(
147    provider: &Provider,
148    gen: &AuthGenerator,
149    ctx: &GenContext,
150    keyring: &Keyring,
151    cache: &AuthCache,
152) -> Result<GeneratedCredential, AuthGenError> {
153    // 1. Check cache
154    if gen.cache_ttl_secs > 0 {
155        if let Some(cached) = cache.get(&provider.name, &ctx.jwt_sub) {
156            return Ok(cached);
157        }
158    }
159
160    // 2. Expand variables in args and env
161    let expanded_args: Vec<String> = gen
162        .args
163        .iter()
164        .map(|a| expand_variables(a, ctx, keyring))
165        .collect::<Result<Vec<_>, _>>()?;
166
167    let mut expanded_env: HashMap<String, String> = HashMap::new();
168    for (k, v) in &gen.env {
169        expanded_env.insert(k.clone(), expand_variables(v, ctx, keyring)?);
170    }
171
172    // 3. Build curated env (don't leak host secrets)
173    let mut final_env: HashMap<String, String> = HashMap::new();
174    for var in &["PATH", "HOME", "TMPDIR"] {
175        if let Ok(val) = std::env::var(var) {
176            final_env.insert(var.to_string(), val);
177        }
178    }
179    final_env.extend(expanded_env);
180
181    // 4. Spawn subprocess
182    let output = match gen.gen_type {
183        AuthGenType::Command => {
184            let command = gen
185                .command
186                .as_deref()
187                .ok_or_else(|| AuthGenError::Config("command required for type=command".into()))?;
188
189            let child = tokio::process::Command::new(command)
190                .args(&expanded_args)
191                .env_clear()
192                .envs(&final_env)
193                .stdout(std::process::Stdio::piped())
194                .stderr(std::process::Stdio::piped())
195                .kill_on_drop(true)
196                .spawn()
197                .map_err(|e| AuthGenError::Spawn(format!("{command}: {e}")))?;
198
199            let timeout = Duration::from_secs(gen.timeout_secs);
200            tokio::time::timeout(timeout, child.wait_with_output())
201                .await
202                .map_err(|_| AuthGenError::Timeout(gen.timeout_secs))?
203                .map_err(AuthGenError::Io)?
204        }
205        AuthGenType::Script => {
206            let interpreter = gen
207                .interpreter
208                .as_deref()
209                .ok_or_else(|| AuthGenError::Config("interpreter required for type=script".into()))?;
210            let script = gen
211                .script
212                .as_deref()
213                .ok_or_else(|| AuthGenError::Config("script required for type=script".into()))?;
214
215            // Write script to a temp file
216            let suffix: u32 = rand::random();
217            let tmp_path = std::env::temp_dir().join(format!("ati_gen_{suffix}.tmp"));
218            std::fs::write(&tmp_path, script).map_err(AuthGenError::Io)?;
219
220            let child = tokio::process::Command::new(interpreter)
221                .arg(&tmp_path)
222                .env_clear()
223                .envs(&final_env)
224                .stdout(std::process::Stdio::piped())
225                .stderr(std::process::Stdio::piped())
226                .kill_on_drop(true)
227                .spawn()
228                .map_err(|e| AuthGenError::Spawn(format!("{interpreter}: {e}")))?;
229
230            let timeout = Duration::from_secs(gen.timeout_secs);
231            let result = tokio::time::timeout(timeout, child.wait_with_output())
232                .await
233                .map_err(|_| AuthGenError::Timeout(gen.timeout_secs))?
234                .map_err(AuthGenError::Io)?;
235
236            // Clean up temp file
237            let _ = std::fs::remove_file(&tmp_path);
238            result
239        }
240    };
241
242    if !output.status.success() {
243        let code = output.status.code().unwrap_or(-1);
244        let stderr = String::from_utf8_lossy(&output.stderr).to_string();
245        return Err(AuthGenError::NonZeroExit { code, stderr });
246    }
247
248    let stdout = String::from_utf8_lossy(&output.stdout);
249
250    // 5. Parse output
251    let cred = match gen.output_format {
252        AuthOutputFormat::Text => GeneratedCredential {
253            value: stdout.trim().to_string(),
254            extra_headers: HashMap::new(),
255            extra_env: HashMap::new(),
256        },
257        AuthOutputFormat::Json => {
258            let json: serde_json::Value = serde_json::from_str(stdout.trim())
259                .map_err(|e| AuthGenError::OutputParse(format!("invalid JSON: {e}")))?;
260
261            let mut extra_headers = HashMap::new();
262            let mut extra_env = HashMap::new();
263            let mut primary_value = stdout.trim().to_string();
264
265            // If no inject map, use the whole output as the primary value
266            if gen.inject.is_empty() {
267                // Try to extract a "token" or "access_token" field as primary
268                if let Some(tok) = json.get("token").or(json.get("access_token")) {
269                    if let Some(s) = tok.as_str() {
270                        primary_value = s.to_string();
271                    }
272                }
273            } else {
274                // Extract fields per inject map
275                let mut found_primary = false;
276                for (json_path, target) in &gen.inject {
277                    let extracted = extract_json_path(&json, json_path)
278                        .ok_or_else(|| {
279                            AuthGenError::OutputParse(format!(
280                                "JSON path '{}' not found in output",
281                                json_path
282                            ))
283                        })?;
284
285                    match target.inject_type.as_str() {
286                        "header" => {
287                            extra_headers.insert(target.name.clone(), extracted);
288                        }
289                        "env" => {
290                            extra_env.insert(target.name.clone(), extracted);
291                        }
292                        "query" => {
293                            // For query injection, use as primary value
294                            if !found_primary {
295                                primary_value = extracted;
296                                found_primary = true;
297                            }
298                        }
299                        _ => {
300                            // Default: treat as primary value
301                            if !found_primary {
302                                primary_value = extracted;
303                                found_primary = true;
304                            }
305                        }
306                    }
307                }
308            }
309
310            GeneratedCredential {
311                value: primary_value,
312                extra_headers,
313                extra_env,
314            }
315        }
316    };
317
318    // 6. Cache
319    cache.insert(&provider.name, &ctx.jwt_sub, cred.clone(), gen.cache_ttl_secs);
320
321    Ok(cred)
322}
323
324// ---------------------------------------------------------------------------
325// Variable expansion
326// ---------------------------------------------------------------------------
327
328/// Expand `${VAR}` placeholders in a string.
329///
330/// Recognized variables:
331/// - `${JWT_SUB}`, `${JWT_SCOPE}`, `${TOOL_NAME}`, `${TIMESTAMP}` — from GenContext
332/// - `${anything_else}` — looked up in the keyring
333fn expand_variables(
334    input: &str,
335    ctx: &GenContext,
336    keyring: &Keyring,
337) -> Result<String, AuthGenError> {
338    let mut result = input.to_string();
339    // Process all ${...} patterns
340    while let Some(start) = result.find("${") {
341        let rest = &result[start + 2..];
342        let end = match rest.find('}') {
343            Some(e) => e,
344            None => break,
345        };
346        let var_name = &rest[..end];
347
348        let replacement = match var_name {
349            "JWT_SUB" => ctx.jwt_sub.clone(),
350            "JWT_SCOPE" => ctx.jwt_scope.clone(),
351            "TOOL_NAME" => ctx.tool_name.clone(),
352            "TIMESTAMP" => ctx.timestamp.to_string(),
353            _ => {
354                // Keyring lookup
355                match keyring.get(var_name) {
356                    Some(val) => val.to_string(),
357                    None => return Err(AuthGenError::KeyringMissing(var_name.to_string())),
358                }
359            }
360        };
361
362        result = format!("{}{}{}", &result[..start], replacement, &rest[end + 1..]);
363    }
364    Ok(result)
365}
366
367// ---------------------------------------------------------------------------
368// JSON path extraction (dot-notation)
369// ---------------------------------------------------------------------------
370
371/// Extract a value from a JSON object using dot-notation path.
372///
373/// Example: `extract_json_path(json, "Credentials.AccessKeyId")`
374/// navigates `json["Credentials"]["AccessKeyId"]`.
375fn extract_json_path(value: &serde_json::Value, path: &str) -> Option<String> {
376    let mut current = value;
377    for segment in path.split('.') {
378        current = current.get(segment)?;
379    }
380    match current {
381        serde_json::Value::String(s) => Some(s.clone()),
382        serde_json::Value::Number(n) => Some(n.to_string()),
383        serde_json::Value::Bool(b) => Some(b.to_string()),
384        other => Some(other.to_string()),
385    }
386}
387
388// ---------------------------------------------------------------------------
389// Tests
390// ---------------------------------------------------------------------------
391
392#[cfg(test)]
393mod tests {
394    use super::*;
395
396    #[test]
397    fn test_expand_variables_context() {
398        let ctx = GenContext {
399            jwt_sub: "agent-7".into(),
400            jwt_scope: "tool:brain__*".into(),
401            tool_name: "brain__query".into(),
402            timestamp: 1773096459,
403        };
404        let keyring = Keyring::empty();
405
406        assert_eq!(
407            expand_variables("${JWT_SUB}", &ctx, &keyring).unwrap(),
408            "agent-7"
409        );
410        assert_eq!(
411            expand_variables("${TOOL_NAME}", &ctx, &keyring).unwrap(),
412            "brain__query"
413        );
414        assert_eq!(
415            expand_variables("${TIMESTAMP}", &ctx, &keyring).unwrap(),
416            "1773096459"
417        );
418        assert_eq!(
419            expand_variables("sub=${JWT_SUB}&tool=${TOOL_NAME}", &ctx, &keyring).unwrap(),
420            "sub=agent-7&tool=brain__query"
421        );
422    }
423
424    #[test]
425    fn test_expand_variables_keyring() {
426        let dir = tempfile::TempDir::new().unwrap();
427        let path = dir.path().join("creds");
428        std::fs::write(&path, r#"{"my_secret":"s3cr3t"}"#).unwrap();
429        let keyring = Keyring::load_credentials(&path).unwrap();
430
431        let ctx = GenContext::default();
432        assert_eq!(
433            expand_variables("${my_secret}", &ctx, &keyring).unwrap(),
434            "s3cr3t"
435        );
436    }
437
438    #[test]
439    fn test_expand_variables_missing_key() {
440        let keyring = Keyring::empty();
441        let ctx = GenContext::default();
442        let err = expand_variables("${nonexistent}", &ctx, &keyring).unwrap_err();
443        assert!(matches!(err, AuthGenError::KeyringMissing(_)));
444    }
445
446    #[test]
447    fn test_expand_variables_no_placeholder() {
448        let keyring = Keyring::empty();
449        let ctx = GenContext::default();
450        assert_eq!(
451            expand_variables("plain text", &ctx, &keyring).unwrap(),
452            "plain text"
453        );
454    }
455
456    #[test]
457    fn test_extract_json_path_simple() {
458        let json: serde_json::Value =
459            serde_json::json!({"token": "abc123", "expires_in": 3600});
460        assert_eq!(extract_json_path(&json, "token"), Some("abc123".into()));
461        assert_eq!(extract_json_path(&json, "expires_in"), Some("3600".into()));
462    }
463
464    #[test]
465    fn test_extract_json_path_nested() {
466        let json: serde_json::Value = serde_json::json!({
467            "Credentials": {
468                "AccessKeyId": "AKIA...",
469                "SecretAccessKey": "wJalrX...",
470                "SessionToken": "FwoGZ..."
471            }
472        });
473        assert_eq!(
474            extract_json_path(&json, "Credentials.AccessKeyId"),
475            Some("AKIA...".into())
476        );
477        assert_eq!(
478            extract_json_path(&json, "Credentials.SessionToken"),
479            Some("FwoGZ...".into())
480        );
481    }
482
483    #[test]
484    fn test_extract_json_path_missing() {
485        let json: serde_json::Value = serde_json::json!({"a": "b"});
486        assert_eq!(extract_json_path(&json, "nonexistent"), None);
487        assert_eq!(extract_json_path(&json, "a.b.c"), None);
488    }
489
490    #[test]
491    fn test_auth_cache_basic() {
492        let cache = AuthCache::new();
493        assert!(cache.get("provider", "sub").is_none());
494
495        let cred = GeneratedCredential {
496            value: "token123".into(),
497            extra_headers: HashMap::new(),
498            extra_env: HashMap::new(),
499        };
500        cache.insert("provider", "sub", cred.clone(), 300);
501
502        let cached = cache.get("provider", "sub").unwrap();
503        assert_eq!(cached.value, "token123");
504    }
505
506    #[test]
507    fn test_auth_cache_zero_ttl_no_cache() {
508        let cache = AuthCache::new();
509        let cred = GeneratedCredential {
510            value: "token".into(),
511            extra_headers: HashMap::new(),
512            extra_env: HashMap::new(),
513        };
514        cache.insert("provider", "sub", cred, 0);
515        assert!(cache.get("provider", "sub").is_none());
516    }
517
518    #[test]
519    fn test_auth_cache_different_keys() {
520        let cache = AuthCache::new();
521        let cred1 = GeneratedCredential {
522            value: "token-a".into(),
523            extra_headers: HashMap::new(),
524            extra_env: HashMap::new(),
525        };
526        let cred2 = GeneratedCredential {
527            value: "token-b".into(),
528            extra_headers: HashMap::new(),
529            extra_env: HashMap::new(),
530        };
531        cache.insert("provider", "agent-1", cred1, 300);
532        cache.insert("provider", "agent-2", cred2, 300);
533
534        assert_eq!(cache.get("provider", "agent-1").unwrap().value, "token-a");
535        assert_eq!(cache.get("provider", "agent-2").unwrap().value, "token-b");
536    }
537
538    #[tokio::test]
539    async fn test_generate_command_text() {
540        let provider = Provider {
541            name: "test".into(),
542            description: "test provider".into(),
543            base_url: String::new(),
544            auth_type: crate::core::manifest::AuthType::Bearer,
545            auth_key_name: None,
546            auth_header_name: None,
547            auth_query_name: None,
548            auth_value_prefix: None,
549            extra_headers: HashMap::new(),
550            oauth2_token_url: None,
551            auth_secret_name: None,
552            oauth2_basic_auth: false,
553            internal: false,
554            handler: "http".into(),
555            mcp_transport: None,
556            mcp_command: None,
557            mcp_args: vec![],
558            mcp_url: None,
559            mcp_env: HashMap::new(),
560            cli_command: None,
561            cli_default_args: vec![],
562            cli_env: HashMap::new(),
563            cli_timeout_secs: None,
564            openapi_spec: None,
565            openapi_include_tags: vec![],
566            openapi_exclude_tags: vec![],
567            openapi_include_operations: vec![],
568            openapi_exclude_operations: vec![],
569            openapi_max_operations: None,
570            openapi_overrides: HashMap::new(),
571            auth_generator: None,
572            category: None,
573            skills: vec![],
574        };
575
576        let gen = AuthGenerator {
577            gen_type: AuthGenType::Command,
578            command: Some("echo".into()),
579            args: vec!["hello-token".into()],
580            interpreter: None,
581            script: None,
582            cache_ttl_secs: 0,
583            output_format: AuthOutputFormat::Text,
584            env: HashMap::new(),
585            inject: HashMap::new(),
586            timeout_secs: 5,
587        };
588
589        let ctx = GenContext::default();
590        let keyring = Keyring::empty();
591        let cache = AuthCache::new();
592
593        let cred = generate(&provider, &gen, &ctx, &keyring, &cache)
594            .await
595            .unwrap();
596        assert_eq!(cred.value, "hello-token");
597        assert!(cred.extra_headers.is_empty());
598    }
599
600    #[tokio::test]
601    async fn test_generate_command_json() {
602        let provider = Provider {
603            name: "test".into(),
604            description: "test".into(),
605            base_url: String::new(),
606            auth_type: crate::core::manifest::AuthType::Bearer,
607            auth_key_name: None,
608            auth_header_name: None,
609            auth_query_name: None,
610            auth_value_prefix: None,
611            extra_headers: HashMap::new(),
612            oauth2_token_url: None,
613            auth_secret_name: None,
614            oauth2_basic_auth: false,
615            internal: false,
616            handler: "http".into(),
617            mcp_transport: None,
618            mcp_command: None,
619            mcp_args: vec![],
620            mcp_url: None,
621            mcp_env: HashMap::new(),
622            cli_command: None,
623            cli_default_args: vec![],
624            cli_env: HashMap::new(),
625            cli_timeout_secs: None,
626            openapi_spec: None,
627            openapi_include_tags: vec![],
628            openapi_exclude_tags: vec![],
629            openapi_include_operations: vec![],
630            openapi_exclude_operations: vec![],
631            openapi_max_operations: None,
632            openapi_overrides: HashMap::new(),
633            auth_generator: None,
634            category: None,
635            skills: vec![],
636        };
637
638        let mut inject = HashMap::new();
639        inject.insert(
640            "Credentials.AccessKeyId".into(),
641            crate::core::manifest::InjectTarget {
642                inject_type: "header".into(),
643                name: "X-Access-Key".into(),
644            },
645        );
646        inject.insert(
647            "Credentials.Secret".into(),
648            crate::core::manifest::InjectTarget {
649                inject_type: "env".into(),
650                name: "AWS_SECRET".into(),
651            },
652        );
653
654        let gen = AuthGenerator {
655            gen_type: AuthGenType::Command,
656            command: Some("echo".into()),
657            args: vec![
658                r#"{"Credentials":{"AccessKeyId":"AKIA123","Secret":"wJalr","SessionToken":"FwoG"}}"#.into(),
659            ],
660            interpreter: None,
661            script: None,
662            cache_ttl_secs: 0,
663            output_format: AuthOutputFormat::Json,
664            env: HashMap::new(),
665            inject,
666            timeout_secs: 5,
667        };
668
669        let ctx = GenContext::default();
670        let keyring = Keyring::empty();
671        let cache = AuthCache::new();
672
673        let cred = generate(&provider, &gen, &ctx, &keyring, &cache)
674            .await
675            .unwrap();
676        assert_eq!(cred.extra_headers.get("X-Access-Key").unwrap(), "AKIA123");
677        assert_eq!(cred.extra_env.get("AWS_SECRET").unwrap(), "wJalr");
678    }
679
680    #[tokio::test]
681    async fn test_generate_script() {
682        let provider = Provider {
683            name: "test".into(),
684            description: "test".into(),
685            base_url: String::new(),
686            auth_type: crate::core::manifest::AuthType::Bearer,
687            auth_key_name: None,
688            auth_header_name: None,
689            auth_query_name: None,
690            auth_value_prefix: None,
691            extra_headers: HashMap::new(),
692            oauth2_token_url: None,
693            auth_secret_name: None,
694            oauth2_basic_auth: false,
695            internal: false,
696            handler: "http".into(),
697            mcp_transport: None,
698            mcp_command: None,
699            mcp_args: vec![],
700            mcp_url: None,
701            mcp_env: HashMap::new(),
702            cli_command: None,
703            cli_default_args: vec![],
704            cli_env: HashMap::new(),
705            cli_timeout_secs: None,
706            openapi_spec: None,
707            openapi_include_tags: vec![],
708            openapi_exclude_tags: vec![],
709            openapi_include_operations: vec![],
710            openapi_exclude_operations: vec![],
711            openapi_max_operations: None,
712            openapi_overrides: HashMap::new(),
713            auth_generator: None,
714            category: None,
715            skills: vec![],
716        };
717
718        let gen = AuthGenerator {
719            gen_type: AuthGenType::Script,
720            command: None,
721            args: vec![],
722            interpreter: Some("bash".into()),
723            script: Some("echo script-token-42".into()),
724            cache_ttl_secs: 0,
725            output_format: AuthOutputFormat::Text,
726            env: HashMap::new(),
727            inject: HashMap::new(),
728            timeout_secs: 5,
729        };
730
731        let ctx = GenContext::default();
732        let keyring = Keyring::empty();
733        let cache = AuthCache::new();
734
735        let cred = generate(&provider, &gen, &ctx, &keyring, &cache)
736            .await
737            .unwrap();
738        assert_eq!(cred.value, "script-token-42");
739    }
740
741    #[tokio::test]
742    async fn test_generate_caches_result() {
743        let provider = Provider {
744            name: "cached_provider".into(),
745            description: "test".into(),
746            base_url: String::new(),
747            auth_type: crate::core::manifest::AuthType::Bearer,
748            auth_key_name: None,
749            auth_header_name: None,
750            auth_query_name: None,
751            auth_value_prefix: None,
752            extra_headers: HashMap::new(),
753            oauth2_token_url: None,
754            auth_secret_name: None,
755            oauth2_basic_auth: false,
756            internal: false,
757            handler: "http".into(),
758            mcp_transport: None,
759            mcp_command: None,
760            mcp_args: vec![],
761            mcp_url: None,
762            mcp_env: HashMap::new(),
763            cli_command: None,
764            cli_default_args: vec![],
765            cli_env: HashMap::new(),
766            cli_timeout_secs: None,
767            openapi_spec: None,
768            openapi_include_tags: vec![],
769            openapi_exclude_tags: vec![],
770            openapi_include_operations: vec![],
771            openapi_exclude_operations: vec![],
772            openapi_max_operations: None,
773            openapi_overrides: HashMap::new(),
774            auth_generator: None,
775            category: None,
776            skills: vec![],
777        };
778
779        let gen = AuthGenerator {
780            gen_type: AuthGenType::Command,
781            command: Some("date".into()),
782            args: vec!["+%s%N".into()],
783            interpreter: None,
784            script: None,
785            cache_ttl_secs: 300,
786            output_format: AuthOutputFormat::Text,
787            env: HashMap::new(),
788            inject: HashMap::new(),
789            timeout_secs: 5,
790        };
791
792        let ctx = GenContext {
793            jwt_sub: "test-agent".into(),
794            ..GenContext::default()
795        };
796        let keyring = Keyring::empty();
797        let cache = AuthCache::new();
798
799        let cred1 = generate(&provider, &gen, &ctx, &keyring, &cache)
800            .await
801            .unwrap();
802        let cred2 = generate(&provider, &gen, &ctx, &keyring, &cache)
803            .await
804            .unwrap();
805        // Second call should return cached value (same value)
806        assert_eq!(cred1.value, cred2.value);
807    }
808
809    #[tokio::test]
810    async fn test_generate_with_variable_expansion() {
811        let provider = Provider {
812            name: "test".into(),
813            description: "test".into(),
814            base_url: String::new(),
815            auth_type: crate::core::manifest::AuthType::Bearer,
816            auth_key_name: None,
817            auth_header_name: None,
818            auth_query_name: None,
819            auth_value_prefix: None,
820            extra_headers: HashMap::new(),
821            oauth2_token_url: None,
822            auth_secret_name: None,
823            oauth2_basic_auth: false,
824            internal: false,
825            handler: "http".into(),
826            mcp_transport: None,
827            mcp_command: None,
828            mcp_args: vec![],
829            mcp_url: None,
830            mcp_env: HashMap::new(),
831            cli_command: None,
832            cli_default_args: vec![],
833            cli_env: HashMap::new(),
834            cli_timeout_secs: None,
835            openapi_spec: None,
836            openapi_include_tags: vec![],
837            openapi_exclude_tags: vec![],
838            openapi_include_operations: vec![],
839            openapi_exclude_operations: vec![],
840            openapi_max_operations: None,
841            openapi_overrides: HashMap::new(),
842            auth_generator: None,
843            category: None,
844            skills: vec![],
845        };
846
847        let gen = AuthGenerator {
848            gen_type: AuthGenType::Command,
849            command: Some("echo".into()),
850            args: vec!["${JWT_SUB}".into()],
851            interpreter: None,
852            script: None,
853            cache_ttl_secs: 0,
854            output_format: AuthOutputFormat::Text,
855            env: HashMap::new(),
856            inject: HashMap::new(),
857            timeout_secs: 5,
858        };
859
860        let ctx = GenContext {
861            jwt_sub: "agent-42".into(),
862            jwt_scope: "*".into(),
863            tool_name: "brain__query".into(),
864            timestamp: 1234567890,
865        };
866        let keyring = Keyring::empty();
867        let cache = AuthCache::new();
868
869        let cred = generate(&provider, &gen, &ctx, &keyring, &cache)
870            .await
871            .unwrap();
872        assert_eq!(cred.value, "agent-42");
873    }
874
875    #[tokio::test]
876    async fn test_generate_timeout() {
877        let provider = Provider {
878            name: "test".into(),
879            description: "test".into(),
880            base_url: String::new(),
881            auth_type: crate::core::manifest::AuthType::Bearer,
882            auth_key_name: None,
883            auth_header_name: None,
884            auth_query_name: None,
885            auth_value_prefix: None,
886            extra_headers: HashMap::new(),
887            oauth2_token_url: None,
888            auth_secret_name: None,
889            oauth2_basic_auth: false,
890            internal: false,
891            handler: "http".into(),
892            mcp_transport: None,
893            mcp_command: None,
894            mcp_args: vec![],
895            mcp_url: None,
896            mcp_env: HashMap::new(),
897            cli_command: None,
898            cli_default_args: vec![],
899            cli_env: HashMap::new(),
900            cli_timeout_secs: None,
901            openapi_spec: None,
902            openapi_include_tags: vec![],
903            openapi_exclude_tags: vec![],
904            openapi_include_operations: vec![],
905            openapi_exclude_operations: vec![],
906            openapi_max_operations: None,
907            openapi_overrides: HashMap::new(),
908            auth_generator: None,
909            category: None,
910            skills: vec![],
911        };
912
913        let gen = AuthGenerator {
914            gen_type: AuthGenType::Command,
915            command: Some("sleep".into()),
916            args: vec!["10".into()],
917            interpreter: None,
918            script: None,
919            cache_ttl_secs: 0,
920            output_format: AuthOutputFormat::Text,
921            env: HashMap::new(),
922            inject: HashMap::new(),
923            timeout_secs: 1,
924        };
925
926        let ctx = GenContext::default();
927        let keyring = Keyring::empty();
928        let cache = AuthCache::new();
929
930        let err = generate(&provider, &gen, &ctx, &keyring, &cache)
931            .await
932            .unwrap_err();
933        assert!(matches!(err, AuthGenError::Timeout(1)));
934    }
935}
agent_tools_interface/core/auth_generator.rs

agent_tools_interface/core/
auth_generator.rs
