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 {
70                    env_vars: HashMap::new(),
71                    args: Vec::new(),
72                })
73            }
74
75            AuthMethod::EnvVar { var_name, .. } => {
76                // For env var auth, the client is responsible for setting the variable
77                // We just validate the variable name here
78                if var_name.is_empty() {
79                    anyhow::bail!("Environment variable name cannot be empty");
80                }
81                if !var_name
82                    .chars()
83                    .all(|c: char| c.is_alphanumeric() || c == '_')
84                {
85                    anyhow::bail!(
86                        "Invalid environment variable name: '{}'. Must contain only alphanumeric characters and underscores.",
87                        var_name
88                    );
89                }
90
91                Ok(Self {
92                    env_vars: HashMap::new(),
93                    args: Vec::new(),
94                })
95            }
96
97            AuthMethod::Terminal { args, env, .. } => {
98                // Terminal auth: pass args and env to the agent process
99                Ok(Self {
100                    env_vars: env.clone(),
101                    args: args.clone(),
102                })
103            }
104
105            // Legacy methods
106            AuthMethod::ApiKey => Ok(Self {
107                env_vars: HashMap::new(),
108                args: Vec::new(),
109            }),
110
111            AuthMethod::OAuth2 => Ok(Self {
112                env_vars: HashMap::new(),
113                args: Vec::new(),
114            }),
115
116            AuthMethod::Bearer => Ok(Self {
117                env_vars: HashMap::new(),
118                args: Vec::new(),
119            }),
120
121            AuthMethod::Custom(_) => Ok(Self {
122                env_vars: HashMap::new(),
123                args: Vec::new(),
124            }),
125        }
126    }
127
128    /// Set an environment variable for this auth handler
129    pub fn with_env(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
130        self.env_vars.insert(key.into(), value.into());
131        self
132    }
133
134    /// Add an argument for this auth handler
135    pub fn with_arg(mut self, arg: impl Into<String>) -> Self {
136        self.args.push(arg.into());
137        self
138    }
139
140    /// Merge another auth handler's configuration into this one
141    pub fn merge(&mut self, other: &AuthHandler) {
142        for (key, value) in &other.env_vars {
143            self.env_vars.insert(key.clone(), value.clone());
144        }
145        self.args.extend(other.args.iter().cloned());
146    }
147}
148
149#[cfg(test)]
150mod tests {
151    use super::*;
152
153    #[test]
154    fn test_auth_handler_agent() {
155        let method = AuthMethod::Agent {
156            id: "agent".to_string(),
157            name: "Agent".to_string(),
158            description: None,
159        };
160        let handler = AuthHandler::new(&method).unwrap();
161        assert!(handler.env_vars.is_empty());
162        assert!(handler.args.is_empty());
163    }
164
165    #[test]
166    fn test_auth_handler_env_var() {
167        let method = AuthMethod::EnvVar {
168            id: "openai".to_string(),
169            name: "OpenAI Key".to_string(),
170            description: None,
171            var_name: "OPENAI_API_KEY".to_string(),
172            link: None,
173        };
174        let handler = AuthHandler::new(&method).unwrap();
175        assert!(handler.env_vars.is_empty());
176        assert!(handler.args.is_empty());
177    }
178
179    #[test]
180    fn test_auth_handler_env_var_invalid_name() {
181        let method = AuthMethod::EnvVar {
182            id: "test".to_string(),
183            name: "Test".to_string(),
184            description: None,
185            var_name: "INVALID-VAR-NAME".to_string(), // Hyphens not allowed
186            link: None,
187        };
188        AuthHandler::new(&method).unwrap_err();
189    }
190
191    #[test]
192    fn test_auth_handler_terminal() {
193        let mut env = HashMap::new();
194        env.insert("VAR1".to_string(), "value1".to_string());
195
196        let method = AuthMethod::Terminal {
197            id: "terminal".to_string(),
198            name: "Terminal".to_string(),
199            description: None,
200            args: vec!["--login".to_string()],
201            env,
202        };
203        let handler = AuthHandler::new(&method).unwrap();
204        assert_eq!(handler.args.len(), 1);
205        assert_eq!(handler.args[0], "--login");
206        assert_eq!(handler.env_vars.get("VAR1").unwrap(), "value1");
207    }
208
209    #[test]
210    fn test_auth_handler_with_env() {
211        let method = AuthMethod::Agent {
212            id: "agent".to_string(),
213            name: "Agent".to_string(),
214            description: None,
215        };
216        let handler = AuthHandler::new(&method)
217            .unwrap()
218            .with_env("MY_VAR", "my_value");
219        assert_eq!(handler.env_vars.get("MY_VAR").unwrap(), "my_value");
220    }
221
222    #[test]
223    fn test_auth_handler_with_arg() {
224        let method = AuthMethod::Agent {
225            id: "agent".to_string(),
226            name: "Agent".to_string(),
227            description: None,
228        };
229        let handler = AuthHandler::new(&method).unwrap().with_arg("--flag");
230        assert_eq!(handler.args.len(), 1);
231        assert_eq!(handler.args[0], "--flag");
232    }
233
234    #[test]
235    fn test_auth_handler_merge() {
236        let method1 = AuthMethod::Agent {
237            id: "agent".to_string(),
238            name: "Agent".to_string(),
239            description: None,
240        };
241        let method2 = AuthMethod::Terminal {
242            id: "terminal".to_string(),
243            name: "Terminal".to_string(),
244            description: None,
245            args: vec!["--login".to_string()],
246            env: {
247                let mut m = HashMap::new();
248                m.insert("VAR".to_string(), "val".to_string());
249                m
250            },
251        };
252
253        let mut handler1 = AuthHandler::new(&method1).unwrap();
254        let handler2 = AuthHandler::new(&method2).unwrap();
255
256        handler1.merge(&handler2);
257        assert_eq!(handler1.args.len(), 1);
258        assert_eq!(handler1.env_vars.get("VAR").unwrap(), "val");
259    }
260}
vtcode_core/auth/auth_handler.rs

vtcode_core/auth/
auth_handler.rs
