From 26419d7d968398f7e0dfbdbbccb3da8fa1e593c6 Mon Sep 17 00:00:00 2001 From: Igor Yamolov Date: Fri, 10 Apr 2026 01:10:49 +0200 Subject: [PATCH] feat(auth): add API key helper command support Allow users to specify a shell command that generates API keys dynamically, for environments where keys are ephemeral, one-use, or rotated periodically. - Add `ApiKeyProvider` enum (`StaticKey` / `HelperCommand`) to model both static and command-based API key sources - Helper commands are configured via env var (`{API_KEY_VAR}_HELPER` convention), `api_key_helper_var` in provider config, or interactively through the provider login UI - Commands are executed asynchronously with configurable timeout (`FORGE_API_KEY_HELPER_TIMEOUT`, default 30s) and `kill_on_drop` - Output format supports optional TTL: `\n---\nTTL: ` or `Expires: ` - Only the command is persisted to credentials file; the key is always obtained fresh by executing the command on load - Backward-compatible serde: old `"sk-123"` format still deserializes correctly via `#[serde(untagged)]` Co-Authored-By: Claude Code --- crates/forge_api/src/forge_api.rs | 2 +- crates/forge_app/src/command_generator.rs | 2 +- .../src/dto/openai/transformers/pipeline.rs | 2 +- crates/forge_config/src/config.rs | 27 ++ crates/forge_domain/src/auth/auth_context.rs | 26 +- crates/forge_domain/src/auth/credentials.rs | 295 +++++++++++++++++- crates/forge_domain/src/auth/new_types.rs | 1 + crates/forge_domain/src/node.rs | 2 +- crates/forge_domain/src/provider.rs | 8 +- crates/forge_infra/src/auth/api_key_helper.rs | 205 ++++++++++++ crates/forge_infra/src/auth/mod.rs | 1 + crates/forge_infra/src/auth/strategy.rs | 81 ++++- crates/forge_infra/src/lib.rs | 1 + crates/forge_main/src/ui.rs | 98 ++++-- crates/forge_repo/src/provider/anthropic.rs | 8 +- crates/forge_repo/src/provider/bedrock.rs | 8 +- crates/forge_repo/src/provider/google.rs | 2 +- crates/forge_repo/src/provider/openai.rs | 4 +- .../provider/openai_responses/repository.rs | 6 +- .../forge_repo/src/provider/provider_repo.rs | 92 ++++-- crates/forge_services/src/app_config.rs | 4 +- crates/forge_services/src/context_engine.rs | 4 +- crates/forge_services/src/provider_auth.rs | 4 +- crates/forge_services/src/provider_service.rs | 4 +- forge.schema.json | 7 + 25 files changed, 799 insertions(+), 95 deletions(-) create mode 100644 crates/forge_infra/src/auth/api_key_helper.rs diff --git a/crates/forge_api/src/forge_api.rs b/crates/forge_api/src/forge_api.rs index 8569b21d6b..f90e374d4a 100644 --- a/crates/forge_api/src/forge_api.rs +++ b/crates/forge_api/src/forge_api.rs @@ -274,7 +274,7 @@ impl< .credential .as_ref() .and_then(|c| match &c.auth_details { - forge_domain::AuthDetails::ApiKey(key) => Some(key.as_str()), + forge_domain::AuthDetails::ApiKey(provider) => Some(provider.api_key().as_str()), _ => None, }) { diff --git a/crates/forge_app/src/command_generator.rs b/crates/forge_app/src/command_generator.rs index 688d3b6f65..4ae0e33564 100644 --- a/crates/forge_app/src/command_generator.rs +++ b/crates/forge_app/src/command_generator.rs @@ -225,7 +225,7 @@ mod tests { url_params: vec![], credential: Some(AuthCredential { id: ProviderId::OPENAI, - auth_details: AuthDetails::ApiKey("test-key".to_string().into()), + auth_details: AuthDetails::static_api_key("test-key".to_string().into()), url_params: Default::default(), }), custom_headers: None, diff --git a/crates/forge_app/src/dto/openai/transformers/pipeline.rs b/crates/forge_app/src/dto/openai/transformers/pipeline.rs index ad1693cbaa..4f6e4b2bc6 100644 --- a/crates/forge_app/src/dto/openai/transformers/pipeline.rs +++ b/crates/forge_app/src/dto/openai/transformers/pipeline.rs @@ -128,7 +128,7 @@ mod tests { fn make_credential(provider_id: ProviderId, key: &str) -> Option { Some(forge_domain::AuthCredential { id: provider_id, - auth_details: forge_domain::AuthDetails::ApiKey(forge_domain::ApiKey::from( + auth_details: forge_domain::AuthDetails::static_api_key(forge_domain::ApiKey::from( key.to_string(), )), url_params: HashMap::new(), diff --git a/crates/forge_config/src/config.rs b/crates/forge_config/src/config.rs index 6b9baaa213..344690b73e 100644 --- a/crates/forge_config/src/config.rs +++ b/crates/forge_config/src/config.rs @@ -72,6 +72,11 @@ pub struct ProviderEntry { /// Environment variable holding the API key for this provider. #[serde(default, skip_serializing_if = "Option::is_none")] pub api_key_var: Option, + /// Shell command that produces an API key on stdout. When set, the + /// command is executed instead of reading a static key from an environment + /// variable. Falls back to `{api_key_var}_HELPER` env var when absent. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub api_key_helper: Option, /// URL template for chat completions; may contain `{{VAR}}` placeholders /// that are substituted from the credential's url params. pub url: String, @@ -353,4 +358,26 @@ mod tests { assert_eq!(actual.temperature, fixture.temperature); } + + #[test] + fn test_provider_entry_api_key_helper_round_trip() { + let fixture = ForgeConfig { + providers: vec![ProviderEntry { + id: "test_provider".to_string(), + url: "https://api.example.com/v1/chat".to_string(), + api_key_helper: Some("vault read -field=token secret/key".to_string()), + ..Default::default() + }], + ..Default::default() + }; + + let toml = toml_edit::ser::to_string_pretty(&fixture).unwrap(); + let actual = ConfigReader::default().read_toml(&toml).build().unwrap(); + + assert_eq!(actual.providers.len(), 1); + assert_eq!( + actual.providers[0].api_key_helper, + Some("vault read -field=token secret/key".to_string()) + ); + } } diff --git a/crates/forge_domain/src/auth/auth_context.rs b/crates/forge_domain/src/auth/auth_context.rs index 418c058577..e0e8581588 100644 --- a/crates/forge_domain/src/auth/auth_context.rs +++ b/crates/forge_domain/src/auth/auth_context.rs @@ -26,6 +26,9 @@ pub struct ApiKeyRequest { pub struct ApiKeyResponse { pub api_key: ApiKey, pub url_params: HashMap, + /// When set, the API key was produced by this shell command and the + /// credential should be stored as a [`HelperCommand`](super::ApiKeyProvider::HelperCommand). + pub helper_command: Option, } // Authorization Code Flow @@ -95,7 +98,7 @@ pub enum AuthContextResponse { } impl AuthContextResponse { - /// Creates an API key authentication context + /// Creates an API key authentication context with a static key. pub fn api_key( request: ApiKeyRequest, api_key: impl ToString, @@ -109,6 +112,27 @@ impl AuthContextResponse { .into_iter() .map(|(k, v)| (k.into(), v.into())) .collect(), + helper_command: None, + }, + }) + } + + /// Creates an API key authentication context backed by a helper command. + pub fn api_key_with_helper( + request: ApiKeyRequest, + api_key: impl ToString, + url_params: HashMap, + command: String, + ) -> Self { + Self::ApiKey(AuthContext { + request, + response: ApiKeyResponse { + api_key: api_key.to_string().into(), + url_params: url_params + .into_iter() + .map(|(k, v)| (k.into(), v.into())) + .collect(), + helper_command: Some(command), }, }) } diff --git a/crates/forge_domain/src/auth/credentials.rs b/crates/forge_domain/src/auth/credentials.rs index 4fcf6e8507..623ed77dfa 100644 --- a/crates/forge_domain/src/auth/credentials.rs +++ b/crates/forge_domain/src/auth/credentials.rs @@ -6,6 +6,57 @@ use serde::{Deserialize, Serialize}; use crate::{AccessToken, ApiKey, OAuthConfig, ProviderId, RefreshToken, URLParam, URLParamValue}; +/// Strategy for providing API keys to a credential. +/// +/// Uses untagged serde representation so that a bare string (legacy format) +/// deserializes as [`StaticKey`](Self::StaticKey), preserving backward +/// compatibility with existing `~/.forge/.credentials.json` files. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(untagged)] +pub enum ApiKeyProvider { + /// A static, user-supplied API key. + StaticKey(ApiKey), + /// A shell command executed via `sh -c` whose trimmed stdout is used as the + /// API key. Only the `command` is persisted; `last_key` and `expires_at` + /// are populated at runtime by executing the command. + HelperCommand { + command: String, + #[serde(skip, default)] + last_key: ApiKey, + #[serde(skip, default)] + expires_at: Option>, + }, +} + +impl ApiKeyProvider { + /// Returns the current API key value. + /// + /// For [`StaticKey`](Self::StaticKey) this is the user-supplied key. For + /// [`HelperCommand`](Self::HelperCommand) this is the last key obtained by + /// executing the command (empty until the first refresh). + pub fn api_key(&self) -> &ApiKey { + match self { + Self::StaticKey(key) => key, + Self::HelperCommand { last_key, .. } => last_key, + } + } + + /// Returns `true` when the key should be refreshed before use. + /// + /// Static keys never expire. Helper-command keys expire based on the + /// `expires_at` field (populated from the command's TTL/Expires metadata). + /// When `expires_at` is `None` the key is treated as single-use and + /// refreshed on every call. + pub fn needs_refresh(&self, buffer: chrono::Duration) -> bool { + match self { + Self::StaticKey(_) => false, + Self::HelperCommand { expires_at: Some(exp), .. } => Utc::now() + buffer >= *exp, + Self::HelperCommand { expires_at: None, .. } => true, + } + } +} + +/// Stored authentication credential for a provider. #[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Setters)] pub struct AuthCredential { pub id: ProviderId, @@ -14,13 +65,15 @@ pub struct AuthCredential { pub url_params: HashMap, } impl AuthCredential { + /// Creates a credential with a static API key. pub fn new_api_key(id: ProviderId, api_key: ApiKey) -> Self { Self { id, - auth_details: AuthDetails::ApiKey(api_key), + auth_details: AuthDetails::static_api_key(api_key), url_params: HashMap::new(), } } + /// Creates a credential with OAuth tokens. pub fn new_oauth(id: ProviderId, tokens: OAuthTokens, config: OAuthConfig) -> Self { Self { id, @@ -28,6 +81,7 @@ impl AuthCredential { url_params: HashMap::new(), } } + /// Creates a credential with OAuth tokens and an API key. pub fn new_oauth_with_api_key( id: ProviderId, tokens: OAuthTokens, @@ -41,6 +95,7 @@ impl AuthCredential { } } + /// Creates a credential with a Google Application Default Credentials token. pub fn new_google_adc(id: ProviderId, access_token: ApiKey) -> Self { Self { id, @@ -52,7 +107,7 @@ impl AuthCredential { /// Checks if the credential needs to be refreshed. pub fn needs_refresh(&self, buffer: chrono::Duration) -> bool { match &self.auth_details { - AuthDetails::ApiKey(_) => false, + AuthDetails::ApiKey(provider) => provider.needs_refresh(buffer), // Google ADC tokens are short-lived (1 hour) and should always be checked/refreshed AuthDetails::GoogleAdc(_) => true, AuthDetails::OAuth { tokens, .. } | AuthDetails::OAuthWithApiKey { tokens, .. } => { @@ -76,7 +131,7 @@ impl AuthCredential { #[serde(rename_all = "snake_case")] pub enum AuthDetails { #[serde(alias = "ApiKey")] - ApiKey(ApiKey), + ApiKey(ApiKeyProvider), #[serde(alias = "GoogleAdc")] GoogleAdc(ApiKey), #[serde(alias = "OAuth")] @@ -93,9 +148,24 @@ pub enum AuthDetails { } impl AuthDetails { + /// Creates a static API key auth details. + pub fn static_api_key(key: ApiKey) -> Self { + Self::ApiKey(ApiKeyProvider::StaticKey(key)) + } + + /// Creates an API key auth details backed by a helper command. + pub fn api_key_from_helper( + command: String, + last_key: ApiKey, + expires_at: Option>, + ) -> Self { + Self::ApiKey(ApiKeyProvider::HelperCommand { command, last_key, expires_at }) + } + + /// Returns the API key if these auth details contain one. pub fn api_key(&self) -> Option<&ApiKey> { match self { - AuthDetails::ApiKey(api_key) => Some(api_key), + AuthDetails::ApiKey(provider) => Some(provider.api_key()), AuthDetails::GoogleAdc(api_key) => Some(api_key), AuthDetails::OAuth { .. } => None, AuthDetails::OAuthWithApiKey { .. } => None, @@ -135,3 +205,220 @@ impl OAuthTokens { Utc::now() >= self.expires_at } } + +#[cfg(test)] +mod tests { + use super::*; + + mod api_key_provider { + use super::*; + + mod static_key { + use pretty_assertions::assert_eq; + + use super::*; + + #[test] + fn api_key_returns_the_key() { + let fixture = ApiKeyProvider::StaticKey(ApiKey::from("sk-test".to_string())); + let actual = fixture.api_key(); + let expected = &ApiKey::from("sk-test".to_string()); + assert_eq!(actual, expected); + } + + #[test] + fn serde_roundtrip() { + let fixture = + ApiKeyProvider::StaticKey(ApiKey::from("sk-test".to_string())); + let json = serde_json::to_string(&fixture).unwrap(); + let actual: ApiKeyProvider = serde_json::from_str(&json).unwrap(); + assert_eq!(actual, fixture); + } + + #[test] + fn serializes_as_bare_string() { + let fixture = + ApiKeyProvider::StaticKey(ApiKey::from("sk-test".to_string())); + let actual = serde_json::to_string(&fixture).unwrap(); + let expected = r#""sk-test""#; + assert_eq!(actual, expected); + } + + #[test] + fn deserializes_from_bare_string() { + let actual: ApiKeyProvider = serde_json::from_str(r#""sk-old-key""#).unwrap(); + let expected = ApiKeyProvider::StaticKey(ApiKey::from("sk-old-key".to_string())); + assert_eq!(actual, expected); + } + } + + mod helper_command { + use pretty_assertions::assert_eq; + + use super::*; + + #[test] + fn api_key_returns_last_key() { + let fixture = ApiKeyProvider::HelperCommand { + command: "echo key".to_string(), + last_key: ApiKey::from("dynamic-key".to_string()), + expires_at: None, + }; + let actual = fixture.api_key(); + let expected = &ApiKey::from("dynamic-key".to_string()); + assert_eq!(actual, expected); + } + + #[test] + fn serializes_only_command() { + let fixture = ApiKeyProvider::HelperCommand { + command: "vault read -field=token".to_string(), + last_key: ApiKey::from("resolved".to_string()), + expires_at: None, + }; + let actual = serde_json::to_string(&fixture).unwrap(); + let expected = r#"{"command":"vault read -field=token"}"#; + assert_eq!(actual, expected); + } + + #[test] + fn deserializes_with_empty_last_key() { + let json = r#"{"command":"vault read -field=token"}"#; + let actual: ApiKeyProvider = serde_json::from_str(json).unwrap(); + let expected = ApiKeyProvider::HelperCommand { + command: "vault read -field=token".to_string(), + last_key: ApiKey::default(), + expires_at: None, + }; + assert_eq!(actual, expected); + } + + #[test] + fn deserialized_needs_refresh() { + let json = r#"{"command":"echo fresh-key"}"#; + let fixture: ApiKeyProvider = serde_json::from_str(json).unwrap(); + let actual = fixture.needs_refresh(chrono::Duration::minutes(5)); + assert!(actual); + } + } + } + + mod needs_refresh { + use super::*; + + mod helper_command { + use pretty_assertions::assert_eq; + + use super::*; + + #[test] + fn without_expires_at_returns_true() { + let fixture = AuthCredential { + auth_details: AuthDetails::api_key_from_helper( + "echo key".to_string(), + ApiKey::from("key".to_string()), + None, + ), + ..AuthCredential::new_api_key( + ProviderId::from("test".to_string()), + ApiKey::from("key".to_string()), + ) + }; + let actual = fixture.needs_refresh(chrono::Duration::minutes(5)); + let expected = true; + assert_eq!(actual, expected); + } + + #[test] + fn with_future_expires_at_returns_false() { + let fixture = AuthCredential { + auth_details: AuthDetails::api_key_from_helper( + "echo key".to_string(), + ApiKey::from("key".to_string()), + Some(Utc::now() + chrono::Duration::hours(1)), + ), + ..AuthCredential::new_api_key( + ProviderId::from("test".to_string()), + ApiKey::from("key".to_string()), + ) + }; + let actual = fixture.needs_refresh(chrono::Duration::minutes(5)); + let expected = false; + assert_eq!(actual, expected); + } + + #[test] + fn with_past_expires_at_returns_true() { + let fixture = AuthCredential { + auth_details: AuthDetails::api_key_from_helper( + "echo key".to_string(), + ApiKey::from("key".to_string()), + Some(Utc::now() - chrono::Duration::minutes(1)), + ), + ..AuthCredential::new_api_key( + ProviderId::from("test".to_string()), + ApiKey::from("key".to_string()), + ) + }; + let actual = fixture.needs_refresh(chrono::Duration::minutes(5)); + let expected = true; + assert_eq!(actual, expected); + } + } + + mod static_key { + use pretty_assertions::assert_eq; + + use super::*; + + #[test] + fn returns_false() { + let fixture = AuthCredential::new_api_key( + ProviderId::from("test".to_string()), + ApiKey::from("key".to_string()), + ); + let actual = fixture.needs_refresh(chrono::Duration::minutes(5)); + let expected = false; + assert_eq!(actual, expected); + } + } + } + + mod backward_compat { + use pretty_assertions::assert_eq; + + use super::*; + + #[test] + fn legacy_credential_json_deserializes() { + let fixture = r#"{ + "id": "anthropic", + "auth_details": {"api_key": "sk-legacy-key"} + }"#; + let actual: AuthCredential = serde_json::from_str(fixture).unwrap(); + let expected = AuthCredential::new_api_key( + ProviderId::from("anthropic".to_string()), + ApiKey::from("sk-legacy-key".to_string()), + ); + assert_eq!(actual, expected); + } + + #[test] + fn helper_credential_serializes_as_expected() { + let fixture = vec![AuthCredential { + id: ProviderId::from("xai".to_string()), + auth_details: AuthDetails::api_key_from_helper( + "printf 'sk-test\\n---\\nTTL: 300'".to_string(), + ApiKey::from("sk-test".to_string()), + None, + ), + url_params: HashMap::new(), + }]; + let actual = serde_json::to_string_pretty(&fixture).unwrap(); + // command persisted, last_key and expires_at skipped + assert!(actual.contains(r#""command""#), "should contain command: {actual}"); + assert!(!actual.contains("last_key"), "should NOT contain last_key: {actual}"); + assert!(!actual.contains("expires_at"), "should NOT contain expires_at: {actual}"); + } + } +} diff --git a/crates/forge_domain/src/auth/new_types.rs b/crates/forge_domain/src/auth/new_types.rs index 2d1e10ffee..20e54549cf 100644 --- a/crates/forge_domain/src/auth/new_types.rs +++ b/crates/forge_domain/src/auth/new_types.rs @@ -2,6 +2,7 @@ use serde::{Deserialize, Serialize}; #[derive( Clone, Serialize, Deserialize, derive_more::From, derive_more::Deref, PartialEq, Eq, Hash, Debug, + Default, )] #[serde(transparent)] pub struct ApiKey(String); diff --git a/crates/forge_domain/src/node.rs b/crates/forge_domain/src/node.rs index 5f172185b8..aaca589f51 100644 --- a/crates/forge_domain/src/node.rs +++ b/crates/forge_domain/src/node.rs @@ -94,7 +94,7 @@ pub struct WorkspaceAuth { impl From for crate::AuthDetails { fn from(auth: WorkspaceAuth) -> Self { - crate::AuthDetails::ApiKey(auth.token) + crate::AuthDetails::static_api_key(auth.token) } } diff --git a/crates/forge_domain/src/provider.rs b/crates/forge_domain/src/provider.rs index 371b74b31e..eefb3bcbf9 100644 --- a/crates/forge_domain/src/provider.rs +++ b/crates/forge_domain/src/provider.rs @@ -252,7 +252,7 @@ impl Provider { self.credential .as_ref() .and_then(|c| match &c.auth_details { - AuthDetails::ApiKey(key) => Some(key), + AuthDetails::ApiKey(provider) => Some(provider.api_key()), _ => None, }) } @@ -354,7 +354,7 @@ mod test_helpers { fn make_credential(provider_id: ProviderId, key: &str) -> Option { Some(AuthCredential { id: provider_id, - auth_details: AuthDetails::ApiKey(ApiKey::from(key.to_string())), + auth_details: AuthDetails::static_api_key(ApiKey::from(key.to_string())), url_params: HashMap::new(), }) } @@ -608,7 +608,7 @@ mod tests { .unwrap(), credential: Some(AuthCredential { id: ProviderId::IO_INTELLIGENCE, - auth_details: AuthDetails::ApiKey(ApiKey::from(fixture.to_string())), + auth_details: AuthDetails::static_api_key(ApiKey::from(fixture.to_string())), url_params: HashMap::new(), }), auth_methods: vec![crate::AuthMethod::ApiKey], @@ -632,7 +632,7 @@ mod tests { url: Url::from_str("https://api.x.ai/v1/chat/completions").unwrap(), credential: Some(AuthCredential { id: ProviderId::XAI, - auth_details: AuthDetails::ApiKey(ApiKey::from(fixture.to_string())), + auth_details: AuthDetails::static_api_key(ApiKey::from(fixture.to_string())), url_params: HashMap::new(), }), auth_methods: vec![crate::AuthMethod::ApiKey], diff --git a/crates/forge_infra/src/auth/api_key_helper.rs b/crates/forge_infra/src/auth/api_key_helper.rs new file mode 100644 index 0000000000..ed708b4c1c --- /dev/null +++ b/crates/forge_infra/src/auth/api_key_helper.rs @@ -0,0 +1,205 @@ +use std::time::Duration; + +use chrono::{DateTime, Utc}; +use forge_domain::{ApiKey, ApiKeyProvider}; + +/// Default timeout for helper command execution. +const DEFAULT_TIMEOUT_SECS: u64 = 30; + +/// Returns the configured helper command timeout, read from the +/// `FORGE_API_KEY_HELPER_TIMEOUT` environment variable. Falls back to +/// [`DEFAULT_TIMEOUT_SECS`] when the variable is absent or unparseable. +fn helper_timeout() -> Duration { + std::env::var("FORGE_API_KEY_HELPER_TIMEOUT") + .ok() + .and_then(|v| v.parse::().ok()) + .map(Duration::from_secs) + .unwrap_or(Duration::from_secs(DEFAULT_TIMEOUT_SECS)) +} + +/// Executes an [`ApiKeyProvider`] to obtain a fresh key. +/// +/// - `StaticKey` — returns the provider unchanged (no-op). +/// - `HelperCommand` — runs the shell command via `sh -c`, parses stdout, and +/// returns an updated provider with the new key and optional expiry. The +/// command is killed if it exceeds the configured timeout. +pub async fn execute(provider: &ApiKeyProvider) -> anyhow::Result { + match provider { + ApiKeyProvider::StaticKey(_) => Ok(provider.clone()), + ApiKeyProvider::HelperCommand { command, .. } => { + let timeout = helper_timeout(); + + let output = tokio::time::timeout( + timeout, + tokio::process::Command::new("sh") + .arg("-c") + .arg(command) + .kill_on_drop(true) + .output(), + ) + .await + .map_err(|_| { + anyhow::anyhow!( + "Auth helper '{command}' timed out after {}s", + timeout.as_secs() + ) + })? + .map_err(|e| anyhow::anyhow!("Failed to execute auth helper '{command}': {e}"))?; + + if !output.status.success() { + anyhow::bail!( + "Auth helper '{command}' exited with status {}", + output.status + ); + } + + let stdout = String::from_utf8(output.stdout) + .map_err(|e| anyhow::anyhow!("Auth helper output is not valid UTF-8: {e}"))?; + + // Normalize CRLF to LF for cross-platform compatibility + let stdout = stdout.replace("\r\n", "\n"); + + let (key, expires_at) = parse_output(&stdout)?; + Ok(ApiKeyProvider::HelperCommand { + command: command.clone(), + last_key: key, + expires_at, + }) + } + } +} + +/// Parses helper command output into an API key and optional expiry. +/// +/// Format: `` or `\n---\nTTL: ` or +/// `\n---\nExpires: `. +fn parse_output(output: &str) -> anyhow::Result<(ApiKey, Option>)> { + let (key_part, metadata) = match output.split_once("\n---\n") { + Some((key, rest)) => (key, Some(rest)), + None => (output, None), + }; + + let key = key_part.trim().to_string(); + if key.is_empty() { + anyhow::bail!("Auth helper produced empty output"); + } + + let expires_at = if let Some(meta) = metadata { + let meta = meta.trim(); + if let Some(secs) = meta.strip_prefix("TTL:") { + let ttl: u64 = secs + .trim() + .parse() + .map_err(|e| anyhow::anyhow!("Invalid TTL value: {e}"))?; + Some(Utc::now() + chrono::Duration::seconds(ttl as i64)) + } else if let Some(ts) = meta.strip_prefix("Expires:") { + let timestamp: i64 = ts + .trim() + .parse() + .map_err(|e| anyhow::anyhow!("Invalid Expires timestamp: {e}"))?; + Some( + DateTime::from_timestamp(timestamp, 0) + .ok_or_else(|| anyhow::anyhow!("Invalid unix timestamp: {timestamp}"))?, + ) + } else { + None + } + } else { + None + }; + + Ok((ApiKey::from(key), expires_at)) +} + +#[cfg(test)] +mod tests { + use super::*; + + mod parse_output { + use super::*; + + #[test] + fn key_only() { + let (key, expires_at) = parse_output("sk-test-key\n").unwrap(); + assert_eq!(key.as_ref(), "sk-test-key"); + assert!(expires_at.is_none()); + } + + #[test] + fn key_with_ttl() { + let (key, expires_at) = parse_output("sk-test-key\n---\nTTL: 3600\n").unwrap(); + assert_eq!(key.as_ref(), "sk-test-key"); + let exp = expires_at.unwrap(); + let expected = Utc::now() + chrono::Duration::seconds(3600); + assert!((exp - expected).num_seconds().abs() < 5); + } + + #[test] + fn key_with_expires() { + let future_ts = Utc::now().timestamp() + 7200; + let input = format!("sk-test-key\n---\nExpires: {future_ts}\n"); + let (key, expires_at) = parse_output(&input).unwrap(); + assert_eq!(key.as_ref(), "sk-test-key"); + let exp = expires_at.unwrap(); + let expected = DateTime::from_timestamp(future_ts, 0).unwrap(); + assert!((exp - expected).num_seconds().abs() < 5); + } + + #[test] + fn empty_output_returns_error() { + assert!(parse_output(" \n").is_err()); + } + + #[test] + fn unknown_metadata_ignored() { + let (key, expires_at) = parse_output("sk-key\n---\nFoo: bar\n").unwrap(); + assert_eq!(key.as_ref(), "sk-key"); + assert!(expires_at.is_none()); + } + + #[test] + fn crlf_line_endings() { + let (key, expires_at) = + parse_output("sk-test-key\r\n---\r\nTTL: 3600\r\n").unwrap(); + assert_eq!(key.as_ref(), "sk-test-key"); + assert!(expires_at.is_some()); + } + } + + mod execute { + use super::*; + + #[tokio::test] + async fn static_key_returns_unchanged() { + let provider = ApiKeyProvider::StaticKey(ApiKey::from("sk-static".to_string())); + let result = execute(&provider).await.unwrap(); + assert_eq!(result, provider); + } + + #[tokio::test] + async fn helper_command_returns_key() { + let provider = ApiKeyProvider::HelperCommand { + command: "echo sk-from-helper".to_string(), + last_key: ApiKey::from("old".to_string()), + expires_at: None, + }; + let result = execute(&provider).await.unwrap(); + match &result { + ApiKeyProvider::HelperCommand { last_key, .. } => { + assert_eq!(last_key.as_ref(), "sk-from-helper"); + } + _ => panic!("Expected HelperCommand"), + } + } + + #[tokio::test] + async fn failing_command_returns_error() { + let provider = ApiKeyProvider::HelperCommand { + command: "false".to_string(), + last_key: ApiKey::from("old".to_string()), + expires_at: None, + }; + assert!(execute(&provider).await.is_err()); + } + } +} diff --git a/crates/forge_infra/src/auth/mod.rs b/crates/forge_infra/src/auth/mod.rs index 1edf7f7320..964def4f99 100644 --- a/crates/forge_infra/src/auth/mod.rs +++ b/crates/forge_infra/src/auth/mod.rs @@ -1,3 +1,4 @@ +pub mod api_key_helper; mod mcp_credentials; mod mcp_token_storage; diff --git a/crates/forge_infra/src/auth/strategy.rs b/crates/forge_infra/src/auth/strategy.rs index 2e223c98a4..6dd9b78f75 100644 --- a/crates/forge_infra/src/auth/strategy.rs +++ b/crates/forge_infra/src/auth/strategy.rs @@ -2,8 +2,9 @@ use std::time::Duration; use forge_app::{AuthStrategy, OAuthHttpProvider, StrategyFactory}; use forge_domain::{ - ApiKey, ApiKeyRequest, AuthContextRequest, AuthContextResponse, AuthCredential, CodeRequest, - DeviceCodeRequest, OAuthConfig, OAuthTokenResponse, OAuthTokens, ProviderId, URLParamSpec, + ApiKey, ApiKeyRequest, AuthContextRequest, AuthContextResponse, AuthCredential, AuthDetails, + CodeRequest, DeviceCodeRequest, OAuthConfig, OAuthTokenResponse, OAuthTokens, ProviderId, + URLParamSpec, }; use google_cloud_auth::credentials::Builder; use oauth2::basic::BasicClient; @@ -42,17 +43,39 @@ impl AuthStrategy for ApiKeyStrategy { context_response: AuthContextResponse, ) -> anyhow::Result { match context_response { - AuthContextResponse::ApiKey(ctx) => Ok(AuthCredential::new_api_key( - self.provider_id.clone(), - ctx.response.api_key, - ) - .url_params(ctx.response.url_params)), + AuthContextResponse::ApiKey(ctx) => { + let auth_details = if let Some(command) = ctx.response.helper_command { + // Execute the helper to validate and obtain the initial key + let initial = forge_domain::ApiKeyProvider::HelperCommand { + command, + last_key: forge_domain::ApiKey::default(), + expires_at: None, + }; + let provider = crate::auth::api_key_helper::execute(&initial).await?; + AuthDetails::ApiKey(provider) + } else { + AuthDetails::static_api_key(ctx.response.api_key) + }; + Ok(AuthCredential { + id: self.provider_id.clone(), + auth_details, + url_params: ctx.response.url_params, + }) + } _ => Err(AuthError::InvalidContext("Expected ApiKey context".to_string()).into()), } } async fn refresh(&self, credential: &AuthCredential) -> anyhow::Result { - // API keys don't expire - return as-is + // Check if credential has a helper command provider that needs refreshing + if let AuthDetails::ApiKey(provider) = &credential.auth_details { + let refreshed_provider = crate::auth::api_key_helper::execute(provider).await?; + return Ok(AuthCredential { + id: credential.id.clone(), + auth_details: AuthDetails::ApiKey(refreshed_provider), + url_params: credential.url_params.clone(), + }); + } Ok(credential.clone()) } } @@ -1127,7 +1150,7 @@ impl StrategyFactory for ForgeAuthStrategyFactory { mod tests { use std::collections::HashMap; - use forge_domain::URLParam; + use forge_domain::{ApiKeyProvider, URLParam}; use pretty_assertions::assert_eq; use super::*; @@ -1361,4 +1384,44 @@ mod tests { let expected = fixture_url_params; assert_eq!(actual.url_params, expected); } + + #[tokio::test] + async fn test_api_key_strategy_refresh_with_helper_command() { + let strategy = ApiKeyStrategy::new(ProviderId::OPENAI, vec![]); + let fixture = AuthCredential::new_api_key( + ProviderId::OPENAI, + ApiKey::from("old-key".to_string()), + ); + // Replace auth_details with a HelperCommand + let fixture = AuthCredential { + auth_details: AuthDetails::api_key_from_helper( + "echo refreshed-key".to_string(), + ApiKey::from("old-key".to_string()), + None, + ), + ..fixture + }; + + let actual = strategy.refresh(&fixture).await.unwrap(); + + match &actual.auth_details { + AuthDetails::ApiKey(ApiKeyProvider::HelperCommand { last_key, .. }) => { + assert_eq!(last_key.as_ref(), "refreshed-key"); + } + other => panic!("Expected HelperCommand, got {other:?}"), + } + } + + #[tokio::test] + async fn test_api_key_strategy_refresh_static_key_unchanged() { + let strategy = ApiKeyStrategy::new(ProviderId::OPENAI, vec![]); + let fixture = AuthCredential::new_api_key( + ProviderId::OPENAI, + ApiKey::from("static-key".to_string()), + ); + + let actual = strategy.refresh(&fixture).await.unwrap(); + + assert_eq!(actual, fixture); + } } diff --git a/crates/forge_infra/src/lib.rs b/crates/forge_infra/src/lib.rs index a6a726d477..dc86d010bd 100644 --- a/crates/forge_infra/src/lib.rs +++ b/crates/forge_infra/src/lib.rs @@ -18,6 +18,7 @@ mod mcp_client; mod mcp_server; mod walker; +pub use auth::api_key_helper; pub use console::StdConsoleWriter; pub use env::ForgeEnvironmentInfra; pub use executor::ForgeCommandExecutorService; diff --git a/crates/forge_main/src/ui.rs b/crates/forge_main/src/ui.rs index fde42605d5..5a4101675a 100644 --- a/crates/forge_main/src/ui.rs +++ b/crates/forge_main/src/ui.rs @@ -2354,43 +2354,81 @@ impl A + Send + Sync> UI }) .collect::>>()?; - // Check if API key is already provided - // For Google ADC, we use a marker to skip prompting - // For other providers, we use the existing key as a default value (autofill) - let api_key_str = if let Some(default_key) = &request.api_key { - let key_str = default_key.as_ref(); - - // Skip prompting only for Google ADC marker - if key_str == "google_adc_marker" { - key_str.to_string() - } else { - // For other providers, show the existing key as default (autofill) + // Choose between static API key and helper command + // Skip for Google ADC marker (auto-completed) + let is_adc_marker = request + .api_key + .as_ref() + .map_or(false, |k| k.as_ref() == "google_adc_marker"); + + if is_adc_marker { + // Google ADC: skip prompting entirely + let response = AuthContextResponse::api_key( + request.clone(), + "google_adc_marker", + url_params, + ); + self.api + .complete_provider_auth(provider_id, response, Duration::from_secs(0)) + .await?; + return Ok(()); + } + + let auth_type_options = vec![ + "Static API Key".to_string(), + "Helper Command (script that generates a key)".to_string(), + ]; + let use_helper = ForgeWidget::select("Authentication type:", auth_type_options) + .prompt()? + .as_deref() + == Some("Helper Command (script that generates a key)"); + + if use_helper { + // Helper command flow + let command = ForgeWidget::input( + "Enter helper command (e.g. vault read -field=token secret/key)", + ) + .prompt()? + .context("Helper command input cancelled")?; + anyhow::ensure!(!command.trim().is_empty(), "Helper command cannot be empty"); + let command = command.trim().to_string(); + + // Send with a placeholder key — the strategy will validate by + // executing the command during complete_provider_auth + let response = AuthContextResponse::api_key_with_helper( + request.clone(), + "", + url_params, + command, + ); + self.spinner.start(Some("Validating helper command..."))?; + self.api + .complete_provider_auth(provider_id, response, Duration::from_secs(0)) + .await?; + self.spinner.stop(None)?; + } else { + // Static API key flow + let api_key_str = if let Some(default_key) = &request.api_key { + let key_str = default_key.as_ref(); let input = ForgeWidget::input(format!("Enter your {provider_id} API key")) .with_default(key_str); let api_key = input.prompt()?.context("API key input cancelled")?; let api_key_str = api_key.trim(); anyhow::ensure!(!api_key_str.is_empty(), "API key cannot be empty"); api_key_str.to_string() - } - } else { - // Prompt for API key input (no existing key) - let input = ForgeWidget::input(format!("Enter your {provider_id} API key")); - let api_key = input.prompt()?.context("API key input cancelled")?; - let api_key_str = api_key.trim(); - anyhow::ensure!(!api_key_str.is_empty(), "API key cannot be empty"); - api_key_str.to_string() - }; - - // Update the context with collected data - let response = AuthContextResponse::api_key(request.clone(), &api_key_str, url_params); + } else { + let input = ForgeWidget::input(format!("Enter your {provider_id} API key")); + let api_key = input.prompt()?.context("API key input cancelled")?; + let api_key_str = api_key.trim(); + anyhow::ensure!(!api_key_str.is_empty(), "API key cannot be empty"); + api_key_str.to_string() + }; - self.api - .complete_provider_auth( - provider_id, - response, - Duration::from_secs(0), // No timeout needed since we have the data - ) - .await?; + let response = AuthContextResponse::api_key(request.clone(), &api_key_str, url_params); + self.api + .complete_provider_auth(provider_id, response, Duration::from_secs(0)) + .await?; + } Ok(()) } diff --git a/crates/forge_repo/src/provider/anthropic.rs b/crates/forge_repo/src/provider/anthropic.rs index 3292f5ab9f..b560e10725 100644 --- a/crates/forge_repo/src/provider/anthropic.rs +++ b/crates/forge_repo/src/provider/anthropic.rs @@ -47,7 +47,7 @@ impl Anthropic { .credential .as_ref() .map(|c| match &c.auth_details { - forge_domain::AuthDetails::ApiKey(key) => key.as_str(), + forge_domain::AuthDetails::ApiKey(provider) => provider.api_key().as_str(), forge_domain::AuthDetails::OAuthWithApiKey { api_key, .. } => api_key.as_str(), forge_domain::AuthDetails::OAuth { tokens, .. } => tokens.access_token.as_str(), forge_domain::AuthDetails::GoogleAdc(api_key) => api_key.as_str(), @@ -437,7 +437,7 @@ mod tests { url: chat_url, credential: Some(forge_domain::AuthCredential { id: forge_app::domain::ProviderId::ANTHROPIC, - auth_details: forge_domain::AuthDetails::ApiKey(forge_domain::ApiKey::from( + auth_details: forge_domain::AuthDetails::static_api_key(forge_domain::ApiKey::from( "sk-test-key".to_string(), )), url_params: std::collections::HashMap::new(), @@ -505,7 +505,7 @@ mod tests { url: chat_url, credential: Some(forge_domain::AuthCredential { id: forge_app::domain::ProviderId::ANTHROPIC, - auth_details: forge_domain::AuthDetails::ApiKey(forge_domain::ApiKey::from( + auth_details: forge_domain::AuthDetails::static_api_key(forge_domain::ApiKey::from( "sk-some-key".to_string(), )), url_params: std::collections::HashMap::new(), @@ -645,7 +645,7 @@ mod tests { url: chat_url, credential: Some(forge_domain::AuthCredential { id: forge_app::domain::ProviderId::ANTHROPIC, - auth_details: forge_domain::AuthDetails::ApiKey(forge_domain::ApiKey::from( + auth_details: forge_domain::AuthDetails::static_api_key(forge_domain::ApiKey::from( "sk-test-key".to_string(), )), url_params: std::collections::HashMap::new(), diff --git a/crates/forge_repo/src/provider/bedrock.rs b/crates/forge_repo/src/provider/bedrock.rs index c5e9653167..340875efe5 100644 --- a/crates/forge_repo/src/provider/bedrock.rs +++ b/crates/forge_repo/src/provider/bedrock.rs @@ -42,7 +42,7 @@ impl BedrockProvider { // Validate API key (bearer token) match &credential.auth_details { - AuthDetails::ApiKey(key) if !key.is_empty() => {} + AuthDetails::ApiKey(provider) if !provider.api_key().is_empty() => {} _ => anyhow::bail!("Bearer token is required in API key field"), } @@ -76,8 +76,8 @@ impl BedrockProvider { .credential .as_ref() .and_then(|c| match &c.auth_details { - AuthDetails::ApiKey(key) if !key.is_empty() => { - Some(key.as_ref().to_string()) + AuthDetails::ApiKey(provider) if !provider.api_key().is_empty() => { + Some(provider.api_key().as_ref().to_string()) } _ => None, }) @@ -997,7 +997,7 @@ mod tests { url_params: vec![], credential: Some(AuthCredential { id: ProviderId::from("bedrock".to_string()), - auth_details: AuthDetails::ApiKey(ApiKey::from(token.to_string())), + auth_details: AuthDetails::static_api_key(ApiKey::from(token.to_string())), url_params, }), custom_headers: None, diff --git a/crates/forge_repo/src/provider/google.rs b/crates/forge_repo/src/provider/google.rs index 390f1cd7f6..e1552e6546 100644 --- a/crates/forge_repo/src/provider/google.rs +++ b/crates/forge_repo/src/provider/google.rs @@ -166,7 +166,7 @@ impl GoogleResponseRepository { // For Vertex AI, the Google ADC token is stored as ApiKey // For OAuth, extract the access token let (token, use_api_key_header) = match creds { - forge_domain::AuthDetails::ApiKey(api_key) => (api_key.as_str().to_string(), true), + forge_domain::AuthDetails::ApiKey(provider) => (provider.api_key().as_str().to_string(), true), forge_domain::AuthDetails::GoogleAdc(token) => (token.as_str().to_string(), false), forge_domain::AuthDetails::OAuth { tokens, .. } => { (tokens.access_token.as_str().to_string(), false) diff --git a/crates/forge_repo/src/provider/openai.rs b/crates/forge_repo/src/provider/openai.rs index 31eccd8592..eca3b1939e 100644 --- a/crates/forge_repo/src/provider/openai.rs +++ b/crates/forge_repo/src/provider/openai.rs @@ -58,7 +58,7 @@ impl OpenAIProvider { .credential .as_ref() .map(|c| match &c.auth_details { - forge_domain::AuthDetails::ApiKey(key) => key.as_str(), + forge_domain::AuthDetails::ApiKey(provider) => provider.api_key().as_str(), forge_domain::AuthDetails::OAuthWithApiKey { api_key, .. } => api_key.as_str(), forge_domain::AuthDetails::OAuth { tokens, .. } => tokens.access_token.as_str(), forge_domain::AuthDetails::GoogleAdc(token) => token.as_str(), @@ -383,7 +383,7 @@ mod tests { fn make_credential(provider_id: ProviderId, key: &str) -> Option { Some(forge_domain::AuthCredential { id: provider_id, - auth_details: forge_domain::AuthDetails::ApiKey(forge_domain::ApiKey::from( + auth_details: forge_domain::AuthDetails::static_api_key(forge_domain::ApiKey::from( key.to_string(), )), url_params: HashMap::new(), diff --git a/crates/forge_repo/src/provider/openai_responses/repository.rs b/crates/forge_repo/src/provider/openai_responses/repository.rs index e311d13c20..91836f7761 100644 --- a/crates/forge_repo/src/provider/openai_responses/repository.rs +++ b/crates/forge_repo/src/provider/openai_responses/repository.rs @@ -77,7 +77,7 @@ impl OpenAIResponsesProvider { .credential .as_ref() .map(|c| match &c.auth_details { - forge_domain::AuthDetails::ApiKey(key) => key.as_str(), + forge_domain::AuthDetails::ApiKey(provider) => provider.api_key().as_str(), forge_domain::AuthDetails::OAuthWithApiKey { api_key, .. } => api_key.as_str(), forge_domain::AuthDetails::OAuth { tokens, .. } => tokens.access_token.as_str(), forge_domain::AuthDetails::GoogleAdc(token) => token.as_str(), @@ -448,7 +448,7 @@ mod tests { fn make_credential(provider_id: ProviderId, key: &str) -> Option { Some(forge_domain::AuthCredential { id: provider_id, - auth_details: forge_domain::AuthDetails::ApiKey(forge_domain::ApiKey::from( + auth_details: forge_domain::AuthDetails::static_api_key(forge_domain::ApiKey::from( key.to_string(), )), url_params: HashMap::new(), @@ -1155,7 +1155,7 @@ mod tests { url: Url::parse("https://api.openai.com/v1").unwrap(), credential: Some(forge_domain::AuthCredential { id: ProviderId::OPENAI, - auth_details: forge_domain::AuthDetails::ApiKey(forge_domain::ApiKey::from( + auth_details: forge_domain::AuthDetails::static_api_key(forge_domain::ApiKey::from( "test-key".to_string(), )), url_params, diff --git a/crates/forge_repo/src/provider/provider_repo.rs b/crates/forge_repo/src/provider/provider_repo.rs index 74c71e9c08..e89a2fa8b6 100644 --- a/crates/forge_repo/src/provider/provider_repo.rs +++ b/crates/forge_repo/src/provider/provider_repo.rs @@ -4,8 +4,8 @@ use bytes::Bytes; use forge_app::domain::{ProviderId, ProviderResponse}; use forge_app::{EnvironmentInfra, FileReaderInfra, FileWriterInfra, HttpInfra}; use forge_domain::{ - AnyProvider, ApiKey, AuthCredential, AuthDetails, Error, MigrationResult, Provider, - ProviderRepository, ProviderType, URLParam, URLParamSpec, URLParamValue, + AnyProvider, ApiKey, ApiKeyProvider, AuthCredential, AuthDetails, Error, MigrationResult, + Provider, ProviderRepository, ProviderType, URLParam, URLParamSpec, URLParamValue, }; use merge::Merge; use serde::Deserialize; @@ -62,6 +62,9 @@ struct ProviderConfig { #[merge(strategy = overwrite)] api_key_vars: Option, #[serde(default)] + #[merge(strategy = overwrite)] + api_key_helper: Option, + #[serde(default)] #[merge(strategy = merge::vec::append)] url_param_vars: Vec, #[serde(default)] @@ -152,6 +155,7 @@ impl From for ProviderConfig { id: ProviderId::from(entry.id), provider_type, api_key_vars: entry.api_key_var, + api_key_helper: entry.api_key_helper, url_param_vars: entry.url_param_vars.into_iter().map(Into::into).collect(), response_type, url: entry.url, @@ -316,7 +320,7 @@ impl< } // Try to create credential from environment variables - if let Ok(credential) = self.create_credential_from_env(&config) { + if let Ok(credential) = self.create_credential_from_env(&config).await { migrated_providers.push(config.id); credentials.push(credential); } @@ -332,20 +336,10 @@ impl< } /// Creates a credential from environment variables for a given config - fn create_credential_from_env( + async fn create_credential_from_env( &self, config: &ProviderConfig, ) -> anyhow::Result { - // Check API key environment variable (if specified) - let api_key = if let Some(api_key_var) = &config.api_key_vars { - self.infra - .get_env_var(api_key_var) - .ok_or_else(|| Error::env_var_not_found(config.id.clone(), api_key_var))? - } else { - // For context engine, we don't use env vars for API key - String::new() - }; - // Check URL parameter environment variables let mut url_params = std::collections::HashMap::new(); @@ -358,10 +352,39 @@ impl< } } - // Create AuthCredential + // Check for helper command: direct config or {api_key_var}_HELPER env var + let helper_command = config.api_key_helper.clone().or_else(|| { + let helper_var = config.api_key_vars.as_ref().map(|v| format!("{v}_HELPER"))?; + self.infra.get_env_var(&helper_var) + }); + + if let Some(command) = helper_command { + let initial = ApiKeyProvider::HelperCommand { + command, + last_key: ApiKey::from(String::new()), + expires_at: None, + }; + let provider = forge_infra::api_key_helper::execute(&initial).await?; + return Ok(AuthCredential { + id: config.id.clone(), + auth_details: AuthDetails::ApiKey(provider), + url_params, + }); + } + + // Fall back to static API key + let api_key = if let Some(api_key_var) = &config.api_key_vars { + self.infra + .get_env_var(api_key_var) + .ok_or_else(|| Error::env_var_not_found(config.id.clone(), api_key_var))? + } else { + // For context engine, we don't use env vars for API key + String::new() + }; + Ok(AuthCredential { id: config.id.clone(), - auth_details: AuthDetails::ApiKey(ApiKey::from(api_key)), + auth_details: AuthDetails::static_api_key(ApiKey::from(api_key)), url_params, }) } @@ -382,8 +405,8 @@ impl< // Google ADC tokens expire quickly, so we refresh them on every load if (credential.id == forge_domain::ProviderId::VERTEX_AI || credential.id == forge_domain::ProviderId::VERTEX_AI_ANTHROPIC) - && let forge_domain::AuthDetails::ApiKey(ref api_key) = credential.auth_details - && api_key.as_ref() == "google_adc_marker" + && let forge_domain::AuthDetails::ApiKey(ref provider) = credential.auth_details + && provider.api_key().as_ref() == "google_adc_marker" { // Refresh the Google ADC credential, preserving url_params match self.refresh_google_adc_credential(&credential).await { @@ -398,6 +421,33 @@ impl< } } + // Refresh helper-command credentials on load — the last_key is not + // persisted so it must be obtained by executing the command. + if let forge_domain::AuthDetails::ApiKey(ref provider) = credential.auth_details { + if provider.needs_refresh(chrono::Duration::zero()) { + match forge_infra::api_key_helper::execute(provider).await { + Ok(refreshed) => { + credential.auth_details = forge_domain::AuthDetails::ApiKey(refreshed); + tracing::info!( + provider = %config.id, + "Successfully refreshed API key from helper command" + ); + } + Err(e) => { + tracing::error!( + provider = %config.id, + error = %e, + "Failed to refresh API key from helper command" + ); + return Err(e.context(format!( + "Failed to execute API key helper for provider {}", + config.id + ))); + } + } + } + } + // Handle models - keep as templates let models = config.models.as_ref().map(|m| match m { Models::Url(model_url_template) => forge_domain::ModelSource::Url( @@ -1031,7 +1081,7 @@ mod env_tests { .find(|c| c.id == ProviderId::OPENAI_COMPATIBLE) .unwrap(); match &openai_compat_cred.auth_details { - AuthDetails::ApiKey(key) => assert_eq!(key.as_str(), "test-openai-key"), + AuthDetails::ApiKey(provider) => assert_eq!(provider.api_key().as_str(), "test-openai-key"), _ => panic!("Expected API key"), } @@ -1051,7 +1101,7 @@ mod env_tests { .find(|c| c.id == ProviderId::ANTHROPIC) .unwrap(); match &anthropic_cred.auth_details { - AuthDetails::ApiKey(key) => assert_eq!(key.as_str(), "test-anthropic-key"), + AuthDetails::ApiKey(provider) => assert_eq!(provider.api_key().as_str(), "test-anthropic-key"), _ => panic!("Expected API key"), } } @@ -1198,7 +1248,7 @@ mod env_tests { .credential .as_ref() .and_then(|c| match &c.auth_details { - forge_domain::AuthDetails::ApiKey(key) => Some(key.to_string()), + forge_domain::AuthDetails::ApiKey(provider) => Some(provider.api_key().to_string()), _ => None, }), Some("test-key-123".to_string()) diff --git a/crates/forge_services/src/app_config.rs b/crates/forge_services/src/app_config.rs index 45551d63f7..7f6d68bddc 100644 --- a/crates/forge_services/src/app_config.rs +++ b/crates/forge_services/src/app_config.rs @@ -132,7 +132,7 @@ mod tests { url: Url::parse("https://api.openai.com").unwrap(), credential: Some(forge_domain::AuthCredential { id: ProviderId::OPENAI, - auth_details: forge_domain::AuthDetails::ApiKey( + auth_details: forge_domain::AuthDetails::static_api_key( forge_domain::ApiKey::from("test-key".to_string()), ), url_params: HashMap::new(), @@ -160,7 +160,7 @@ mod tests { url_params: vec![], credential: Some(forge_domain::AuthCredential { id: ProviderId::ANTHROPIC, - auth_details: forge_domain::AuthDetails::ApiKey( + auth_details: forge_domain::AuthDetails::static_api_key( forge_domain::ApiKey::from("test-key".to_string()), ), url_params: HashMap::new(), diff --git a/crates/forge_services/src/context_engine.rs b/crates/forge_services/src/context_engine.rs index 713ca56bb4..3017c6b423 100644 --- a/crates/forge_services/src/context_engine.rs +++ b/crates/forge_services/src/context_engine.rs @@ -103,7 +103,7 @@ impl< .context("No authentication credentials found. Please authenticate first.")?; match &credential.auth_details { - AuthDetails::ApiKey(token) => { + AuthDetails::ApiKey(provider) => { // Extract user_id from URL params let user_id_str = credential .url_params @@ -113,7 +113,7 @@ impl< })?; let user_id = UserId::from_string(user_id_str.as_str())?; - Ok((token.clone(), user_id)) + Ok((provider.api_key().clone(), user_id)) } _ => anyhow::bail!("ForgeServices credential must be an API key"), } diff --git a/crates/forge_services/src/provider_auth.rs b/crates/forge_services/src/provider_auth.rs index 72b8693632..bcdd724a46 100644 --- a/crates/forge_services/src/provider_auth.rs +++ b/crates/forge_services/src/provider_auth.rs @@ -149,7 +149,8 @@ where // Iterate through auth methods and try to refresh for auth_method in &provider.auth_methods { match auth_method { - AuthMethod::OAuthDevice(_) + AuthMethod::ApiKey + | AuthMethod::OAuthDevice(_) | AuthMethod::OAuthCode(_) | AuthMethod::CodexDevice(_) | AuthMethod::GoogleAdc => { @@ -197,7 +198,6 @@ where } } } - _ => {} } } } diff --git a/crates/forge_services/src/provider_service.rs b/crates/forge_services/src/provider_service.rs index 39748dd3d6..4e12f63722 100644 --- a/crates/forge_services/src/provider_service.rs +++ b/crates/forge_services/src/provider_service.rs @@ -214,7 +214,7 @@ mod tests { url_params: vec![], credential: Some(AuthCredential { id: ProviderId::OPENAI, - auth_details: AuthDetails::ApiKey(forge_domain::ApiKey::from( + auth_details: AuthDetails::static_api_key(forge_domain::ApiKey::from( "test-key".to_string(), )), url_params: HashMap::new(), @@ -238,7 +238,7 @@ mod tests { url_params: vec![], credential: Some(AuthCredential { id: ProviderId::OPENAI, - auth_details: AuthDetails::ApiKey(forge_domain::ApiKey::from( + auth_details: AuthDetails::static_api_key(forge_domain::ApiKey::from( "test-key".to_string(), )), url_params: HashMap::new(), diff --git a/forge.schema.json b/forge.schema.json index 43cc190609..241ef8303e 100644 --- a/forge.schema.json +++ b/forge.schema.json @@ -614,6 +614,13 @@ "description": "A single provider entry defined inline in `forge.toml`.\n\nInline providers are merged with the built-in provider list; entries with\nthe same `id` override the corresponding built-in entry field-by-field,\nwhile entries with a new `id` are appended to the list.", "type": "object", "properties": { + "api_key_helper": { + "description": "Shell command that produces an API key on stdout. When set, the\ncommand is executed instead of reading a static key from an environment\nvariable. Falls back to `{api_key_var}_HELPER` env var when absent.", + "type": [ + "string", + "null" + ] + }, "api_key_var": { "description": "Environment variable holding the API key for this provider.", "type": [