codetether_agent/secrets/
mod.rs1use anyhow::{Context, Result};
7use std::collections::HashMap;
8use std::sync::Arc;
9use tokio::sync::RwLock;
10use vaultrs::client::{VaultClient, VaultClientSettingsBuilder};
11use vaultrs::kv2;
12
13#[allow(dead_code)]
15const DEFAULT_SECRETS_PATH: &str = "secret/data/codetether/providers";
16
17#[derive(Clone)]
19pub struct SecretsManager {
20 client: Option<Arc<VaultClient>>,
21 pub cache: Arc<RwLock<HashMap<String, String>>>,
23 mount: String,
24 path: String,
25}
26
27impl Default for SecretsManager {
28 fn default() -> Self {
29 Self {
30 client: None,
31 cache: Arc::new(RwLock::new(HashMap::new())),
32 mount: "secret".to_string(),
33 path: "codetether/providers".to_string(),
34 }
35 }
36}
37
38impl SecretsManager {
39 pub async fn new(config: &VaultConfig) -> Result<Self> {
41 let settings = VaultClientSettingsBuilder::default()
42 .address(&config.address)
43 .token(&config.token)
44 .build()
45 .context("Failed to build Vault client settings")?;
46
47 let client = VaultClient::new(settings).context("Failed to create Vault client")?;
48
49 Ok(Self {
50 client: Some(Arc::new(client)),
51 cache: Arc::new(RwLock::new(HashMap::new())),
52 mount: config.mount.clone().unwrap_or_else(|| "secret".to_string()),
53 path: config
54 .path
55 .clone()
56 .unwrap_or_else(|| "codetether/providers".to_string()),
57 })
58 }
59
60 pub async fn from_env() -> Result<Self> {
62 let address = std::env::var("VAULT_ADDR").context("VAULT_ADDR not set")?;
63 let token = std::env::var("VAULT_TOKEN").context("VAULT_TOKEN not set")?;
64 let mount = std::env::var("VAULT_MOUNT").ok();
65 let path = std::env::var("VAULT_SECRETS_PATH").ok();
66
67 let config = VaultConfig {
68 address,
69 token,
70 mount,
71 path,
72 };
73
74 Self::new(&config).await
75 }
76
77 pub fn is_connected(&self) -> bool {
79 self.client.is_some()
80 }
81
82 pub async fn get_api_key(&self, provider_id: &str) -> Result<Option<String>> {
84 {
86 let cache = self.cache.read().await;
87 if let Some(key) = cache.get(provider_id) {
88 return Ok(Some(key.clone()));
89 }
90 }
91
92 let client = match &self.client {
94 Some(c) => c,
95 None => return Ok(None),
96 };
97
98 let secret_path = format!("{}/{}", self.path, provider_id);
99
100 match kv2::read::<ProviderSecrets>(client.as_ref(), &self.mount, &secret_path).await {
101 Ok(secret) => {
102 if let Some(ref api_key) = secret.api_key {
104 let mut cache = self.cache.write().await;
105 cache.insert(provider_id.to_string(), api_key.clone());
106 }
107 Ok(secret.api_key)
108 }
109 Err(vaultrs::error::ClientError::APIError { code: 404, .. }) => Ok(None),
110 Err(e) => {
111 tracing::warn!("Failed to fetch secret for {}: {}", provider_id, e);
112 Ok(None)
113 }
114 }
115 }
116
117 pub async fn get_provider_secrets(&self, provider_id: &str) -> Result<Option<ProviderSecrets>> {
119 let client = match &self.client {
120 Some(c) => c,
121 None => return Ok(None),
122 };
123
124 let secret_path = format!("{}/{}", self.path, provider_id);
125
126 match kv2::read::<ProviderSecrets>(client.as_ref(), &self.mount, &secret_path).await {
127 Ok(secret) => Ok(Some(secret)),
128 Err(vaultrs::error::ClientError::APIError { code: 404, .. }) => Ok(None),
129 Err(e) => {
130 tracing::warn!("Failed to fetch secrets for {}: {}", provider_id, e);
131 Ok(None)
132 }
133 }
134 }
135
136 pub async fn set_provider_secrets(
138 &self,
139 provider_id: &str,
140 secrets: &ProviderSecrets,
141 ) -> Result<()> {
142 let client = match &self.client {
143 Some(c) => c,
144 None => anyhow::bail!("Vault client not configured"),
145 };
146
147 let secret_path = format!("{}/{}", self.path, provider_id);
148 kv2::set(client.as_ref(), &self.mount, &secret_path, secrets)
149 .await
150 .with_context(|| format!("Failed to write provider secrets for {}", provider_id))?;
151
152 let mut cache = self.cache.write().await;
154 if let Some(api_key) = secrets.api_key.clone() {
155 cache.insert(provider_id.to_string(), api_key);
156 } else {
157 cache.remove(provider_id);
158 }
159
160 Ok(())
161 }
162
163 pub async fn has_api_key(&self, provider_id: &str) -> bool {
165 matches!(self.get_api_key(provider_id).await, Ok(Some(_)))
166 }
167
168 pub async fn list_configured_providers(&self) -> Result<Vec<String>> {
170 let client = match &self.client {
171 Some(c) => c,
172 None => return Ok(Vec::new()),
173 };
174
175 match kv2::list(client.as_ref(), &self.mount, &self.path).await {
176 Ok(keys) => Ok(keys),
177 Err(vaultrs::error::ClientError::APIError { code: 404, .. }) => Ok(Vec::new()),
178 Err(e) => {
179 tracing::warn!("Failed to list providers: {}", e);
180 Ok(Vec::new())
181 }
182 }
183 }
184
185 pub async fn clear_cache(&self) {
187 let mut cache = self.cache.write().await;
188 cache.clear();
189 }
190}
191
192#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
194pub struct VaultConfig {
195 pub address: String,
197
198 pub token: String,
200
201 #[serde(default)]
203 pub mount: Option<String>,
204
205 #[serde(default)]
207 pub path: Option<String>,
208}
209
210impl Default for VaultConfig {
211 fn default() -> Self {
212 Self {
213 address: String::new(),
214 token: String::new(),
215 mount: Some("secret".to_string()),
216 path: Some("codetether/providers".to_string()),
217 }
218 }
219}
220
221#[derive(Clone, serde::Serialize, serde::Deserialize)]
223pub struct ProviderSecrets {
224 #[serde(default)]
226 pub api_key: Option<String>,
227
228 #[serde(default)]
230 pub base_url: Option<String>,
231
232 #[serde(default)]
234 pub organization: Option<String>,
235
236 #[serde(default)]
238 pub headers: Option<HashMap<String, String>>,
239
240 #[serde(flatten)]
242 pub extra: HashMap<String, serde_json::Value>,
243}
244
245impl std::fmt::Debug for ProviderSecrets {
246 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
247 f.debug_struct("ProviderSecrets")
248 .field("api_key", &self.api_key.as_ref().map(|_| "<REDACTED>"))
249 .field("api_key_len", &self.api_key.as_ref().map(|k| k.len()))
250 .field("base_url", &self.base_url)
251 .field("organization", &self.organization)
252 .field("headers_present", &self.headers.is_some())
253 .field("extra_fields", &self.extra.len())
254 .finish()
255 }
256}
257
258impl ProviderSecrets {
259 pub fn has_valid_api_key(&self) -> bool {
261 self.api_key
262 .as_ref()
263 .map(|k| !k.is_empty())
264 .unwrap_or(false)
265 }
266
267 pub fn api_key_len(&self) -> Option<usize> {
269 self.api_key.as_ref().map(|k| k.len())
270 }
271}
272
273static SECRETS_MANAGER: tokio::sync::OnceCell<SecretsManager> = tokio::sync::OnceCell::const_new();
275
276pub async fn init_secrets_manager(config: &VaultConfig) -> Result<()> {
278 let manager = SecretsManager::new(config).await?;
279 SECRETS_MANAGER
280 .set(manager)
281 .map_err(|_| anyhow::anyhow!("Secrets manager already initialized"))?;
282 Ok(())
283}
284
285pub fn init_from_manager(manager: SecretsManager) -> Result<()> {
287 SECRETS_MANAGER
288 .set(manager)
289 .map_err(|_| anyhow::anyhow!("Secrets manager already initialized"))?;
290 Ok(())
291}
292
293pub fn secrets_manager() -> Option<&'static SecretsManager> {
295 SECRETS_MANAGER.get()
296}
297
298pub async fn get_api_key(provider_id: &str) -> Option<String> {
300 match SECRETS_MANAGER.get() {
301 Some(manager) => manager.get_api_key(provider_id).await.ok().flatten(),
302 None => None,
303 }
304}
305
306pub async fn has_api_key(provider_id: &str) -> bool {
308 match SECRETS_MANAGER.get() {
309 Some(manager) => manager.has_api_key(provider_id).await,
310 None => false,
311 }
312}
313
314pub async fn get_provider_secrets(provider_id: &str) -> Option<ProviderSecrets> {
316 match SECRETS_MANAGER.get() {
317 Some(manager) => manager
318 .get_provider_secrets(provider_id)
319 .await
320 .ok()
321 .flatten(),
322 None => None,
323 }
324}
325
326pub async fn set_provider_secrets(provider_id: &str, secrets: &ProviderSecrets) -> Result<()> {
328 match SECRETS_MANAGER.get() {
329 Some(manager) => manager.set_provider_secrets(provider_id, secrets).await,
330 None => anyhow::bail!("Secrets manager not initialized"),
331 }
332}