vtcode_core/auth/
auth_handler.rs

1//! Authentication handler for different ACP auth methods.
2//!
3//! This module provides a unified interface for handling authentication
4//! across different auth methods specified in the ACP protocol.
5
6use anyhow::Result;
7use hashbrown::HashMap;
8use serde::{Deserialize, Serialize};
9
10/// ACP authentication methods mirrored locally so `vtcode-core` does not
11/// depend on the ACP client crate.
12#[derive(Debug, Clone, Serialize, Deserialize)]
13#[serde(tag = "type", rename_all = "snake_case")]
14pub enum AuthMethod {
15    #[serde(rename = "agent")]
16    Agent {
17        id: String,
18        name: String,
19        #[serde(skip_serializing_if = "Option::is_none")]
20        description: Option<String>,
21    },
22    #[serde(rename = "env_var")]
23    EnvVar {
24        id: String,
25        name: String,
26        #[serde(skip_serializing_if = "Option::is_none")]
27        description: Option<String>,
28        var_name: String,
29        #[serde(skip_serializing_if = "Option::is_none")]
30        link: Option<String>,
31    },
32    #[serde(rename = "terminal")]
33    Terminal {
34        id: String,
35        name: String,
36        #[serde(skip_serializing_if = "Option::is_none")]
37        description: Option<String>,
38        #[serde(default, skip_serializing_if = "Vec::is_empty")]
39        args: Vec<String>,
40        #[serde(default, skip_serializing_if = "HashMap::is_empty")]
41        env: HashMap<String, String>,
42    },
43    #[serde(rename = "api_key")]
44    ApiKey,
45    #[serde(rename = "oauth2")]
46    OAuth2,
47    #[serde(rename = "bearer")]
48    Bearer,
49    #[serde(rename = "custom")]
50    Custom(String),
51}
52
53/// Handles authentication based on the auth method type
54#[derive(Debug, Clone)]
55pub struct AuthHandler {
56    /// Environment variables to set for the agent process
57    pub env_vars: HashMap<String, String>,
58
59    /// Arguments to pass to the agent process
60    pub args: Vec<String>,
61}
62
63impl AuthHandler {
64    /// Create a new auth handler for the given auth method
65    pub fn new(method: &AuthMethod) -> Result<Self> {
66        match method {
67            AuthMethod::Agent { .. } => {
68                // Agent handles auth itself - no special configuration needed
69                Ok(Self { env_vars: HashMap::new(), args: Vec::new() })
70            }
71
72            AuthMethod::EnvVar { var_name, .. } => {
73                // For env var auth, the client is responsible for setting the variable
74                // We just validate the variable name here
75                if var_name.is_empty() {
76                    anyhow::bail!("Environment variable name cannot be empty");
77                }
78                if !var_name.chars().all(|c: char| c.is_alphanumeric() || c == '_') {
79                    anyhow::bail!(
80                        "Invalid environment variable name: '{var_name}'. Must contain only alphanumeric characters and underscores."
81                    );
82                }
83
84                Ok(Self { env_vars: HashMap::new(), args: Vec::new() })
85            }
86
87            AuthMethod::Terminal { args, env, .. } => {
88                // Terminal auth: pass args and env to the agent process
89                Ok(Self { env_vars: env.clone(), args: args.clone() })
90            }
91
92            // Legacy methods
93            AuthMethod::ApiKey => Ok(Self { env_vars: HashMap::new(), args: Vec::new() }),
94
95            AuthMethod::OAuth2 => Ok(Self { env_vars: HashMap::new(), args: Vec::new() }),
96
97            AuthMethod::Bearer => Ok(Self { env_vars: HashMap::new(), args: Vec::new() }),
98
99            AuthMethod::Custom(_) => Ok(Self { env_vars: HashMap::new(), args: Vec::new() }),
100        }
101    }
102
103    /// Set an environment variable for this auth handler
104    pub fn with_env(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
105        self.env_vars.insert(key.into(), value.into());
106        self
107    }
108
109    /// Add an argument for this auth handler
110    pub fn with_arg(mut self, arg: impl Into<String>) -> Self {
111        self.args.push(arg.into());
112        self
113    }
114
115    /// Merge another auth handler's configuration into this one
116    pub fn merge(&mut self, other: &AuthHandler) {
117        for (key, value) in &other.env_vars {
118            self.env_vars.insert(key.clone(), value.clone());
119        }
120        self.args.extend(other.args.iter().cloned());
121    }
122}
123
124#[cfg(test)]
125mod tests {
126    use super::*;
127
128    #[test]
129    fn test_auth_handler_agent() {
130        let method = AuthMethod::Agent {
131            id: "agent".to_string(),
132            name: "Agent".to_string(),
133            description: None,
134        };
135        let handler = AuthHandler::new(&method).unwrap();
136        assert!(handler.env_vars.is_empty());
137        assert!(handler.args.is_empty());
138    }
139
140    #[test]
141    fn test_auth_handler_env_var() {
142        let method = AuthMethod::EnvVar {
143            id: "openai".to_string(),
144            name: "OpenAI Key".to_string(),
145            description: None,
146            var_name: "OPENAI_API_KEY".to_string(),
147            link: None,
148        };
149        let handler = AuthHandler::new(&method).unwrap();
150        assert!(handler.env_vars.is_empty());
151        assert!(handler.args.is_empty());
152    }
153
154    #[test]
155    fn test_auth_handler_env_var_invalid_name() {
156        let method = AuthMethod::EnvVar {
157            id: "test".to_string(),
158            name: "Test".to_string(),
159            description: None,
160            var_name: "INVALID-VAR-NAME".to_string(), // Hyphens not allowed
161            link: None,
162        };
163        AuthHandler::new(&method).unwrap_err();
164    }
165
166    #[test]
167    fn test_auth_handler_terminal() {
168        let mut env = HashMap::new();
169        env.insert("VAR1".to_string(), "value1".to_string());
170
171        let method = AuthMethod::Terminal {
172            id: "terminal".to_string(),
173            name: "Terminal".to_string(),
174            description: None,
175            args: vec!["--login".to_string()],
176            env,
177        };
178        let handler = AuthHandler::new(&method).unwrap();
179        assert_eq!(handler.args.len(), 1);
180        assert_eq!(handler.args[0], "--login");
181        assert_eq!(handler.env_vars.get("VAR1").unwrap(), "value1");
182    }
183
184    #[test]
185    fn test_auth_handler_with_env() {
186        let method = AuthMethod::Agent {
187            id: "agent".to_string(),
188            name: "Agent".to_string(),
189            description: None,
190        };
191        let handler = AuthHandler::new(&method).unwrap().with_env("MY_VAR", "my_value");
192        assert_eq!(handler.env_vars.get("MY_VAR").unwrap(), "my_value");
193    }
194
195    #[test]
196    fn test_auth_handler_with_arg() {
197        let method = AuthMethod::Agent {
198            id: "agent".to_string(),
199            name: "Agent".to_string(),
200            description: None,
201        };
202        let handler = AuthHandler::new(&method).unwrap().with_arg("--flag");
203        assert_eq!(handler.args.len(), 1);
204        assert_eq!(handler.args[0], "--flag");
205    }
206
207    #[test]
208    fn test_auth_handler_merge() {
209        let method1 = AuthMethod::Agent {
210            id: "agent".to_string(),
211            name: "Agent".to_string(),
212            description: None,
213        };
214        let method2 = AuthMethod::Terminal {
215            id: "terminal".to_string(),
216            name: "Terminal".to_string(),
217            description: None,
218            args: vec!["--login".to_string()],
219            env: {
220                let mut m = HashMap::new();
221                m.insert("VAR".to_string(), "val".to_string());
222                m
223            },
224        };
225
226        let mut handler1 = AuthHandler::new(&method1).unwrap();
227        let handler2 = AuthHandler::new(&method2).unwrap();
228
229        handler1.merge(&handler2);
230        assert_eq!(handler1.args.len(), 1);
231        assert_eq!(handler1.env_vars.get("VAR").unwrap(), "val");
232    }
233}
vtcode_core/auth/auth_handler.rs

vtcode_core/auth/
auth_handler.rs
