feat(core/providers): support command-backed auth

This commit is contained in:
James Peter
2026-04-03 13:16:12 +10:00
parent 1715a804df
commit 6dc620cbb9
6 changed files with 473 additions and 19 deletions
+16 -2
View File
@@ -389,8 +389,9 @@ pub(crate) async fn stream_chat_completions(
loop { loop {
attempt += 1; attempt += 1;
let auth = auth_manager.as_ref().and_then(|m| m.auth()); let base_auth = auth_manager.as_ref().and_then(|m| m.auth());
let mut req_builder = provider.create_request_builder(client, &auth).await?; let auth = provider.effective_auth(&base_auth).await?;
let mut req_builder = provider.create_request_builder_with_auth(client, &auth).await?;
if let Some(auth) = auth.as_ref() { if let Some(auth) = auth.as_ref() {
if auth.mode.is_chatgpt() { if auth.mode.is_chatgpt() {
@@ -454,6 +455,19 @@ pub(crate) async fn stream_chat_completions(
} }
Ok(res) => { Ok(res) => {
let status = res.status(); let status = res.status();
if status == StatusCode::UNAUTHORIZED && provider.has_command_auth() {
provider.invalidate_cached_auth_token();
if attempt > max_retries {
return Err(CodexErr::RetryLimit(RetryLimitReachedError {
status,
request_id: None,
retryable: true,
}));
}
let delay = backoff(attempt);
tokio::time::sleep(delay).await;
continue;
}
if !(status == StatusCode::TOO_MANY_REQUESTS || status.is_server_error()) { if !(status == StatusCode::TOO_MANY_REQUESTS || status.is_server_error()) {
let body = (res.text().await).unwrap_or_default(); let body = (res.text().await).unwrap_or_default();
if let Ok(logger) = debug_logger.lock() { if let Ok(logger) = debug_logger.lock() {
+31 -8
View File
@@ -780,7 +780,8 @@ impl ModelClient {
} }
} }
let auth = auth_manager.as_ref().and_then(|m| m.auth()); let base_auth = auth_manager.as_ref().and_then(|m| m.auth());
let auth = self.provider.effective_auth(&base_auth).await?;
let endpoint = self.provider.get_full_url(&auth); let endpoint = self.provider.get_full_url(&auth);
let url = reqwest::Url::parse(&endpoint).map_err(|err| { let url = reqwest::Url::parse(&endpoint).map_err(|err| {
@@ -798,7 +799,12 @@ impl ModelClient {
}; };
let mut req_builder = self let mut req_builder = self
.provider .provider
.create_request_builder_for_url(&self.client, &auth, reqwest::Method::GET, url) .create_request_builder_for_url_with_auth(
&self.client,
&auth,
reqwest::Method::GET,
url,
)
.await?; .await?;
let has_beta_header = req_builder let has_beta_header = req_builder
@@ -1226,7 +1232,8 @@ impl ModelClient {
let mut auth_refresh_error: Option<RefreshTokenError> = None; let mut auth_refresh_error: Option<RefreshTokenError> = None;
// Always fetch the latest auth in case a prior attempt refreshed the token. // Always fetch the latest auth in case a prior attempt refreshed the token.
let auth = auth_manager.as_ref().and_then(|m| m.auth()); let base_auth = auth_manager.as_ref().and_then(|m| m.auth());
let auth = self.provider.effective_auth(&base_auth).await?;
trace!( trace!(
"POST to {}: {}", "POST to {}: {}",
@@ -1236,7 +1243,7 @@ impl ModelClient {
let mut req_builder = self let mut req_builder = self
.provider .provider
.create_request_builder(&self.client, &auth) .create_request_builder_with_auth(&self.client, &auth)
.await?; .await?;
let has_beta_header = req_builder let has_beta_header = req_builder
@@ -1425,7 +1432,9 @@ impl ModelClient {
.and_then(|raw| parse_retry_after_header(raw, now)); .and_then(|raw| parse_retry_after_header(raw, now));
if status == StatusCode::UNAUTHORIZED { if status == StatusCode::UNAUTHORIZED {
if let Some(manager) = auth_manager.as_ref() { if self.provider.has_command_auth() {
self.provider.invalidate_cached_auth_token();
} else if let Some(manager) = auth_manager.as_ref() {
match manager.refresh_token_classified().await { match manager.refresh_token_classified().await {
Ok(Some(_)) => {} Ok(Some(_)) => {}
Ok(None) => { Ok(None) => {
@@ -1437,7 +1446,7 @@ impl ModelClient {
auth_refresh_error = Some(err); auth_refresh_error = Some(err);
} }
} }
} else { } else if auth.is_none() {
auth_refresh_error = Some(RefreshTokenError::permanent( auth_refresh_error = Some(RefreshTokenError::permanent(
"Authentication manager unavailable; please log in again.", "Authentication manager unavailable; please log in again.",
)); ));
@@ -1849,10 +1858,11 @@ impl ModelClient {
let mut request_id = String::new(); let mut request_id = String::new();
loop { loop {
let auth = auth_manager.as_ref().and_then(|m| m.auth()); let base_auth = auth_manager.as_ref().and_then(|m| m.auth());
let auth = self.provider.effective_auth(&base_auth).await?;
let mut request = self let mut request = self
.provider .provider
.create_compact_request_builder(&self.client, &auth) .create_compact_request_builder_with_auth(&self.client, &auth)
.await?; .await?;
// Ensure Responses API beta header is present for compact calls. Mirror the // Ensure Responses API beta header is present for compact calls. Mirror the
@@ -2991,6 +3001,7 @@ mod tests {
env_key: None, env_key: None,
env_key_instructions: None, env_key_instructions: None,
experimental_bearer_token: None, experimental_bearer_token: None,
auth: None,
wire_api: WireApi::Responses, wire_api: WireApi::Responses,
query_params: None, query_params: None,
http_headers: None, http_headers: None,
@@ -3039,6 +3050,7 @@ mod tests {
env_key: None, env_key: None,
env_key_instructions: None, env_key_instructions: None,
experimental_bearer_token: None, experimental_bearer_token: None,
auth: None,
wire_api: WireApi::Responses, wire_api: WireApi::Responses,
query_params: None, query_params: None,
http_headers: None, http_headers: None,
@@ -3089,6 +3101,7 @@ mod tests {
env_key: None, env_key: None,
env_key_instructions: None, env_key_instructions: None,
experimental_bearer_token: None, experimental_bearer_token: None,
auth: None,
wire_api: WireApi::Responses, wire_api: WireApi::Responses,
query_params: None, query_params: None,
http_headers: Some(headers), http_headers: Some(headers),
@@ -3236,6 +3249,7 @@ mod tests {
env_key: Some("TEST_API_KEY".to_string()), env_key: Some("TEST_API_KEY".to_string()),
env_key_instructions: None, env_key_instructions: None,
experimental_bearer_token: None, experimental_bearer_token: None,
auth: None,
wire_api: WireApi::Responses, wire_api: WireApi::Responses,
query_params: None, query_params: None,
http_headers: None, http_headers: None,
@@ -3302,6 +3316,7 @@ mod tests {
env_key: Some("TEST_API_KEY".to_string()), env_key: Some("TEST_API_KEY".to_string()),
env_key_instructions: None, env_key_instructions: None,
experimental_bearer_token: None, experimental_bearer_token: None,
auth: None,
wire_api: WireApi::Responses, wire_api: WireApi::Responses,
query_params: None, query_params: None,
http_headers: None, http_headers: None,
@@ -3354,6 +3369,7 @@ mod tests {
env_key: Some("TEST_API_KEY".to_string()), env_key: Some("TEST_API_KEY".to_string()),
env_key_instructions: None, env_key_instructions: None,
experimental_bearer_token: None, experimental_bearer_token: None,
auth: None,
wire_api: WireApi::Responses, wire_api: WireApi::Responses,
query_params: None, query_params: None,
http_headers: None, http_headers: None,
@@ -3449,6 +3465,7 @@ mod tests {
env_key: Some("TEST_API_KEY".to_string()), env_key: Some("TEST_API_KEY".to_string()),
env_key_instructions: None, env_key_instructions: None,
experimental_bearer_token: None, experimental_bearer_token: None,
auth: None,
wire_api: WireApi::Responses, wire_api: WireApi::Responses,
query_params: None, query_params: None,
http_headers: None, http_headers: None,
@@ -3556,6 +3573,7 @@ mod tests {
env_key: Some("TEST_API_KEY".to_string()), env_key: Some("TEST_API_KEY".to_string()),
env_key_instructions: None, env_key_instructions: None,
experimental_bearer_token: None, experimental_bearer_token: None,
auth: None,
wire_api: WireApi::Responses, wire_api: WireApi::Responses,
query_params: None, query_params: None,
http_headers: None, http_headers: None,
@@ -3796,6 +3814,7 @@ mod tests {
env_key: Some("TEST_API_KEY".to_string()), env_key: Some("TEST_API_KEY".to_string()),
env_key_instructions: None, env_key_instructions: None,
experimental_bearer_token: None, experimental_bearer_token: None,
auth: None,
wire_api: WireApi::Responses, wire_api: WireApi::Responses,
query_params: None, query_params: None,
http_headers: None, http_headers: None,
@@ -3827,6 +3846,7 @@ mod tests {
env_key: Some("TEST_API_KEY".to_string()), env_key: Some("TEST_API_KEY".to_string()),
env_key_instructions: None, env_key_instructions: None,
experimental_bearer_token: None, experimental_bearer_token: None,
auth: None,
wire_api: WireApi::Responses, wire_api: WireApi::Responses,
query_params: None, query_params: None,
http_headers: None, http_headers: None,
@@ -3861,6 +3881,7 @@ mod tests {
env_key: Some("TEST_API_KEY".to_string()), env_key: Some("TEST_API_KEY".to_string()),
env_key_instructions: None, env_key_instructions: None,
experimental_bearer_token: None, experimental_bearer_token: None,
auth: None,
wire_api: WireApi::Responses, wire_api: WireApi::Responses,
query_params: None, query_params: None,
http_headers: None, http_headers: None,
@@ -3892,6 +3913,7 @@ mod tests {
env_key: Some("TEST_API_KEY".to_string()), env_key: Some("TEST_API_KEY".to_string()),
env_key_instructions: None, env_key_instructions: None,
experimental_bearer_token: None, experimental_bearer_token: None,
auth: None,
wire_api: WireApi::Responses, wire_api: WireApi::Responses,
query_params: None, query_params: None,
http_headers: None, http_headers: None,
@@ -3923,6 +3945,7 @@ mod tests {
env_key: Some("TEST_API_KEY".to_string()), env_key: Some("TEST_API_KEY".to_string()),
env_key_instructions: None, env_key_instructions: None,
experimental_bearer_token: None, experimental_bearer_token: None,
auth: None,
wire_api: WireApi::Responses, wire_api: WireApi::Responses,
query_params: None, query_params: None,
http_headers: None, http_headers: None,
+33
View File
@@ -65,6 +65,36 @@ mod validation;
use defaults::{default_responses_originator, default_review_model, default_true_local}; use defaults::{default_responses_originator, default_review_model, default_true_local};
const OPENAI_BASE_URL_ENV_VAR: &str = "OPENAI_BASE_URL"; const OPENAI_BASE_URL_ENV_VAR: &str = "OPENAI_BASE_URL";
const RESERVED_MODEL_PROVIDER_IDS: [&str; 2] = ["openai", "oss"];
fn validate_reserved_model_provider_ids(
model_providers: &HashMap<String, ModelProviderInfo>,
) -> Result<(), String> {
let mut conflicts = model_providers
.keys()
.filter(|key| RESERVED_MODEL_PROVIDER_IDS.contains(&key.as_str()))
.map(|key| format!("`{key}`"))
.collect::<Vec<_>>();
conflicts.sort_unstable();
if conflicts.is_empty() {
Ok(())
} else {
Err(format!(
"model_providers contains reserved built-in provider IDs: {}. Built-in providers cannot be overridden. Rename your custom provider.",
conflicts.join(", ")
))
}
}
fn validate_model_providers(model_providers: &HashMap<String, ModelProviderInfo>) -> Result<(), String> {
validate_reserved_model_provider_ids(model_providers)?;
for (key, provider) in model_providers {
provider
.validate()
.map_err(|message| format!("model_providers.{key}: {message}"))?;
}
Ok(())
}
pub use builder::ConfigBuilder; pub use builder::ConfigBuilder;
pub use defaults::set_default_originator; pub use defaults::set_default_originator;
@@ -1099,6 +1129,8 @@ impl Config {
} }
let effective_openai_base_url = openai_base_url.or(openai_base_url_from_env); let effective_openai_base_url = openai_base_url.or(openai_base_url_from_env);
validate_model_providers(&cfg.model_providers).map_err(std::io::Error::other)?;
let mut model_providers = built_in_model_providers(effective_openai_base_url); let mut model_providers = built_in_model_providers(effective_openai_base_url);
// Merge user-defined providers into the built-in list. // Merge user-defined providers into the built-in list.
for (key, provider) in cfg.model_providers.into_iter() { for (key, provider) in cfg.model_providers.into_iter() {
@@ -2372,6 +2404,7 @@ model_verbosity = "high"
wire_api: crate::WireApi::Chat, wire_api: crate::WireApi::Chat,
env_key_instructions: None, env_key_instructions: None,
experimental_bearer_token: None, experimental_bearer_token: None,
auth: None,
query_params: None, query_params: None,
http_headers: None, http_headers: None,
env_http_headers: None, env_http_headers: None,
+316 -9
View File
@@ -8,12 +8,21 @@
use crate::CodexAuth; use crate::CodexAuth;
use crate::error::CodexErr; use crate::error::CodexErr;
use code_protocol::config_types::ModelProviderAuthInfo;
use serde::Deserialize; use serde::Deserialize;
use serde::Serialize; use serde::Serialize;
use serde_json::Value; use serde_json::Value;
use std::collections::{BTreeMap, HashMap}; use std::collections::{BTreeMap, HashMap};
use std::env::VarError; use std::env::VarError;
use std::io;
use std::path::Path;
use std::path::PathBuf;
use std::process::Stdio;
use std::sync::Mutex;
use std::sync::OnceLock;
use std::time::Duration; use std::time::Duration;
use std::time::Instant;
use tokio::process::Command;
use crate::error::EnvVarError; use crate::error::EnvVarError;
const DEFAULT_STREAM_IDLE_TIMEOUT_MS: u64 = 300_000; const DEFAULT_STREAM_IDLE_TIMEOUT_MS: u64 = 300_000;
const DEFAULT_STREAM_MAX_RETRIES: u64 = 5; const DEFAULT_STREAM_MAX_RETRIES: u64 = 5;
@@ -23,6 +32,35 @@ const MAX_STREAM_MAX_RETRIES: u64 = 100;
/// Hard cap for user-configured `request_max_retries`. /// Hard cap for user-configured `request_max_retries`.
const MAX_REQUEST_MAX_RETRIES: u64 = 100; const MAX_REQUEST_MAX_RETRIES: u64 = 100;
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
struct ProviderAuthCacheKey {
command: String,
args: Vec<String>,
cwd: PathBuf,
}
impl From<&ModelProviderAuthInfo> for ProviderAuthCacheKey {
fn from(value: &ModelProviderAuthInfo) -> Self {
Self {
command: value.command.clone(),
args: value.args.clone(),
cwd: value.cwd.as_path().to_path_buf(),
}
}
}
#[derive(Debug, Clone)]
struct CachedProviderAuthToken {
access_token: String,
fetched_at: Instant,
}
fn provider_auth_cache() -> &'static Mutex<HashMap<ProviderAuthCacheKey, CachedProviderAuthToken>> {
static CACHE: OnceLock<Mutex<HashMap<ProviderAuthCacheKey, CachedProviderAuthToken>>> =
OnceLock::new();
CACHE.get_or_init(|| Mutex::new(HashMap::new()))
}
/// Wire protocol that the provider speaks. Most third-party services only /// Wire protocol that the provider speaks. Most third-party services only
/// implement the classic OpenAI Chat Completions JSON schema, whereas OpenAI /// implement the classic OpenAI Chat Completions JSON schema, whereas OpenAI
/// itself (and a handful of others) additionally expose the more modern /// itself (and a handful of others) additionally expose the more modern
@@ -63,6 +101,9 @@ pub struct ModelProviderInfo {
/// this may be necessary when using this programmatically. /// this may be necessary when using this programmatically.
pub experimental_bearer_token: Option<String>, pub experimental_bearer_token: Option<String>,
/// Command-backed bearer-token configuration for this provider.
pub auth: Option<ModelProviderAuthInfo>,
/// Which wire protocol this provider expects. /// Which wire protocol this provider expects.
#[serde(default)] #[serde(default)]
pub wire_api: WireApi, pub wire_api: WireApi,
@@ -189,6 +230,49 @@ pub struct OpenRouterMaxPrice {
} }
impl ModelProviderInfo { impl ModelProviderInfo {
pub(crate) fn validate(&self) -> std::result::Result<(), String> {
let Some(auth) = self.auth.as_ref() else {
return Ok(());
};
if auth.command.trim().is_empty() {
return Err("provider auth.command must not be empty".to_string());
}
let mut conflicts = Vec::new();
if self.env_key.is_some() {
conflicts.push("env_key");
}
if self.experimental_bearer_token.is_some() {
conflicts.push("experimental_bearer_token");
}
if self.requires_openai_auth {
conflicts.push("requires_openai_auth");
}
if conflicts.is_empty() {
Ok(())
} else {
Err(format!(
"provider auth cannot be combined with {}",
conflicts.join(", ")
))
}
}
pub(crate) fn has_command_auth(&self) -> bool {
self.auth.is_some()
}
pub(crate) fn invalidate_cached_auth_token(&self) {
if let Some(auth) = self.auth.as_ref() {
provider_auth_cache()
.lock()
.unwrap()
.remove(&ProviderAuthCacheKey::from(auth));
}
}
/// Construct a `POST` RequestBuilder for the given URL using the provided /// Construct a `POST` RequestBuilder for the given URL using the provided
/// reqwest Client applying: /// reqwest Client applying:
/// • provider-specific headers (static + env based) /// • provider-specific headers (static + env based)
@@ -202,13 +286,22 @@ impl ModelProviderInfo {
client: &'a reqwest::Client, client: &'a reqwest::Client,
auth: &Option<CodexAuth>, auth: &Option<CodexAuth>,
) -> crate::error::Result<reqwest::RequestBuilder> { ) -> crate::error::Result<reqwest::RequestBuilder> {
let effective_auth = self.effective_auth(auth)?; let effective_auth = self.effective_auth(auth).await?;
let url = self.get_full_url(&effective_auth); self.create_request_builder_with_auth(client, &effective_auth)
.await
}
pub async fn create_request_builder_with_auth<'a>(
&'a self,
client: &'a reqwest::Client,
auth: &Option<CodexAuth>,
) -> crate::error::Result<reqwest::RequestBuilder> {
let url = self.get_full_url(auth);
let mut builder = client.post(&url); let mut builder = client.post(&url);
if let Some(auth) = effective_auth.as_ref() { if let Some(auth) = auth.as_ref() {
builder = builder.bearer_auth(auth.get_token().await?); builder = builder.bearer_auth(auth.get_token().await?);
} }
@@ -224,11 +317,23 @@ impl ModelProviderInfo {
method: reqwest::Method, method: reqwest::Method,
url: reqwest::Url, url: reqwest::Url,
) -> crate::error::Result<reqwest::RequestBuilder> { ) -> crate::error::Result<reqwest::RequestBuilder> {
let effective_auth = self.effective_auth(auth)?; let effective_auth = self.effective_auth(auth).await?;
self.create_request_builder_for_url_with_auth(client, &effective_auth, method, url)
.await
}
pub async fn create_request_builder_for_url_with_auth<'a>(
&'a self,
client: &'a reqwest::Client,
auth: &Option<CodexAuth>,
method: reqwest::Method,
url: reqwest::Url,
) -> crate::error::Result<reqwest::RequestBuilder> {
let mut builder = client.request(method, url); let mut builder = client.request(method, url);
if let Some(auth) = effective_auth.as_ref() { if let Some(auth) = auth.as_ref() {
builder = builder.bearer_auth(auth.get_token().await?); builder = builder.bearer_auth(auth.get_token().await?);
} }
@@ -239,28 +344,37 @@ impl ModelProviderInfo {
&'a self, &'a self,
client: &'a reqwest::Client, client: &'a reqwest::Client,
auth: &Option<CodexAuth>, auth: &Option<CodexAuth>,
) -> crate::error::Result<reqwest::RequestBuilder> {
let effective_auth = self.effective_auth(auth).await?;
self.create_compact_request_builder_with_auth(client, &effective_auth)
.await
}
pub async fn create_compact_request_builder_with_auth<'a>(
&'a self,
client: &'a reqwest::Client,
auth: &Option<CodexAuth>,
) -> crate::error::Result<reqwest::RequestBuilder> { ) -> crate::error::Result<reqwest::RequestBuilder> {
if !matches!(self.wire_api, WireApi::Responses | WireApi::ResponsesWebsocket) { if !matches!(self.wire_api, WireApi::Responses | WireApi::ResponsesWebsocket) {
return Err(CodexErr::UnsupportedOperation( return Err(CodexErr::UnsupportedOperation(
"Compaction endpoint requires Responses API providers".to_string(), "Compaction endpoint requires Responses API providers".to_string(),
)); ));
} }
let effective_auth = self.effective_auth(auth)?; let url = self.get_compact_url(auth).ok_or_else(|| {
let url = self.get_compact_url(&effective_auth).ok_or_else(|| {
CodexErr::UnsupportedOperation( CodexErr::UnsupportedOperation(
"Compaction endpoint requires Responses API providers".to_string(), "Compaction endpoint requires Responses API providers".to_string(),
) )
})?; })?;
let mut builder = client.post(url); let mut builder = client.post(url);
if let Some(auth) = effective_auth.as_ref() { if let Some(auth) = auth.as_ref() {
builder = builder.bearer_auth(auth.get_token().await?); builder = builder.bearer_auth(auth.get_token().await?);
} }
Ok(self.apply_http_headers(builder)) Ok(self.apply_http_headers(builder))
} }
fn effective_auth( pub(crate) async fn effective_auth(
&self, &self,
auth: &Option<CodexAuth>, auth: &Option<CodexAuth>,
) -> crate::error::Result<Option<CodexAuth>> { ) -> crate::error::Result<Option<CodexAuth>> {
@@ -273,6 +387,13 @@ impl ModelProviderInfo {
return Ok(Some(CodexAuth::from_api_key(token))); return Ok(Some(CodexAuth::from_api_key(token)));
} }
if let Some(provider_auth) = self.auth.as_ref() {
let token = resolve_provider_auth_token(provider_auth)
.await
.map_err(|err| CodexErr::UnsupportedOperation(err.to_string()))?;
return Ok(Some(CodexAuth::from_api_key(&token)));
}
match self.api_key() { match self.api_key() {
Ok(Some(key)) => Ok(Some(CodexAuth::from_api_key(&key))), Ok(Some(key)) => Ok(Some(CodexAuth::from_api_key(&key))),
Ok(None) => Ok(auth.clone()), Ok(None) => Ok(auth.clone()),
@@ -459,6 +580,100 @@ impl ModelProviderInfo {
} }
} }
async fn resolve_provider_auth_token(config: &ModelProviderAuthInfo) -> io::Result<String> {
let cache_key = ProviderAuthCacheKey::from(config);
if let Some(cached_token) = provider_auth_cache().lock().unwrap().get(&cache_key).cloned() {
let should_use_cached_token = match config.refresh_interval() {
Some(refresh_interval) => cached_token.fetched_at.elapsed() < refresh_interval,
None => true,
};
if should_use_cached_token {
return Ok(cached_token.access_token);
}
}
let access_token = run_provider_auth_command(config).await?;
provider_auth_cache().lock().unwrap().insert(
cache_key,
CachedProviderAuthToken {
access_token: access_token.clone(),
fetched_at: Instant::now(),
},
);
Ok(access_token)
}
async fn run_provider_auth_command(config: &ModelProviderAuthInfo) -> io::Result<String> {
let program = resolve_provider_auth_program(&config.command, config.cwd.as_path())?;
let mut command = Command::new(&program);
command
.args(&config.args)
.current_dir(config.cwd.as_path())
.stdin(Stdio::null())
.stdout(Stdio::piped())
.stderr(Stdio::piped())
.kill_on_drop(true);
let output = tokio::time::timeout(config.timeout(), command.output())
.await
.map_err(|_| {
io::Error::other(format!(
"provider auth command `{}` timed out after {} ms",
config.command,
config.timeout_ms.get()
))
})?
.map_err(|err| {
io::Error::other(format!(
"provider auth command `{}` failed to start: {err}",
config.command
))
})?;
if !output.status.success() {
let status = output.status;
let stderr = String::from_utf8_lossy(&output.stderr).trim().to_string();
let stderr_suffix = if stderr.is_empty() {
String::new()
} else {
format!(": {stderr}")
};
return Err(io::Error::other(format!(
"provider auth command `{}` exited with status {status}{stderr_suffix}",
config.command
)));
}
let stdout = String::from_utf8(output.stdout).map_err(|_| {
io::Error::other(format!(
"provider auth command `{}` wrote non-UTF-8 data to stdout",
config.command
))
})?;
let access_token = stdout.trim().to_string();
if access_token.is_empty() {
return Err(io::Error::other(format!(
"provider auth command `{}` produced an empty token",
config.command
)));
}
Ok(access_token)
}
fn resolve_provider_auth_program(command: &str, cwd: &Path) -> io::Result<PathBuf> {
let path = Path::new(command);
if path.is_absolute() {
return Ok(path.to_path_buf());
}
if path.components().count() > 1 {
return Ok(cwd.join(path));
}
Ok(PathBuf::from(command))
}
const DEFAULT_OLLAMA_PORT: u32 = 11434; const DEFAULT_OLLAMA_PORT: u32 = 11434;
pub const BUILT_IN_OSS_MODEL_PROVIDER_ID: &str = "oss"; pub const BUILT_IN_OSS_MODEL_PROVIDER_ID: &str = "oss";
@@ -500,6 +715,7 @@ pub fn built_in_model_providers(
env_key: None, env_key: None,
env_key_instructions: None, env_key_instructions: None,
experimental_bearer_token: None, experimental_bearer_token: None,
auth: None,
wire_api: wire_api_override_from_env("OPENAI_WIRE_API") wire_api: wire_api_override_from_env("OPENAI_WIRE_API")
.unwrap_or(WireApi::Responses), .unwrap_or(WireApi::Responses),
query_params: None, query_params: None,
@@ -567,6 +783,7 @@ pub fn create_oss_provider_with_base_url(base_url: &str) -> ModelProviderInfo {
env_key: None, env_key: None,
env_key_instructions: None, env_key_instructions: None,
experimental_bearer_token: None, experimental_bearer_token: None,
auth: None,
wire_api: WireApi::Chat, wire_api: WireApi::Chat,
query_params: None, query_params: None,
http_headers: None, http_headers: None,
@@ -594,7 +811,11 @@ fn matches_azure_responses_base_url(base_url: &str) -> bool {
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use super::*; use super::*;
use code_utils_absolute_path::AbsolutePathBuf;
use code_utils_absolute_path::AbsolutePathBufGuard;
use pretty_assertions::assert_eq; use pretty_assertions::assert_eq;
use std::num::NonZeroU64;
use tempfile::tempdir;
#[test] #[test]
fn test_deserialize_ollama_model_provider_toml() { fn test_deserialize_ollama_model_provider_toml() {
@@ -608,6 +829,7 @@ base_url = "http://localhost:11434/v1"
env_key: None, env_key: None,
env_key_instructions: None, env_key_instructions: None,
experimental_bearer_token: None, experimental_bearer_token: None,
auth: None,
wire_api: WireApi::Chat, wire_api: WireApi::Chat,
query_params: None, query_params: None,
http_headers: None, http_headers: None,
@@ -637,6 +859,7 @@ query_params = { api-version = "2025-04-01-preview" }
env_key: Some("AZURE_OPENAI_API_KEY".into()), env_key: Some("AZURE_OPENAI_API_KEY".into()),
env_key_instructions: None, env_key_instructions: None,
experimental_bearer_token: None, experimental_bearer_token: None,
auth: None,
wire_api: WireApi::Chat, wire_api: WireApi::Chat,
query_params: Some(maplit::hashmap! { query_params: Some(maplit::hashmap! {
"api-version".to_string() => "2025-04-01-preview".to_string(), "api-version".to_string() => "2025-04-01-preview".to_string(),
@@ -669,6 +892,7 @@ env_http_headers = { "X-Example-Env-Header" = "EXAMPLE_ENV_VAR" }
env_key: Some("API_KEY".into()), env_key: Some("API_KEY".into()),
env_key_instructions: None, env_key_instructions: None,
experimental_bearer_token: None, experimental_bearer_token: None,
auth: None,
wire_api: WireApi::Chat, wire_api: WireApi::Chat,
query_params: None, query_params: None,
http_headers: Some(maplit::hashmap! { http_headers: Some(maplit::hashmap! {
@@ -697,6 +921,7 @@ env_http_headers = { "X-Example-Env-Header" = "EXAMPLE_ENV_VAR" }
env_key: None, env_key: None,
env_key_instructions: None, env_key_instructions: None,
experimental_bearer_token: None, experimental_bearer_token: None,
auth: None,
wire_api: WireApi::Responses, wire_api: WireApi::Responses,
query_params: None, query_params: None,
http_headers: None, http_headers: None,
@@ -731,6 +956,7 @@ env_http_headers = { "X-Example-Env-Header" = "EXAMPLE_ENV_VAR" }
env_key: None, env_key: None,
env_key_instructions: None, env_key_instructions: None,
experimental_bearer_token: None, experimental_bearer_token: None,
auth: None,
wire_api: WireApi::Responses, wire_api: WireApi::Responses,
query_params: None, query_params: None,
http_headers: None, http_headers: None,
@@ -768,4 +994,85 @@ env_http_headers = { "X-Example-Env-Header" = "EXAMPLE_ENV_VAR" }
assert_eq!(version, code_version::wire_compatible_version()); assert_eq!(version, code_version::wire_compatible_version());
} }
#[test]
fn test_deserialize_provider_auth_config_defaults() {
let base_dir = tempdir().unwrap();
let provider_toml = r#"
name = "Corp"
[auth]
command = "./scripts/print-token"
args = ["--format=text"]
"#;
let provider: ModelProviderInfo = {
let _guard = AbsolutePathBufGuard::new(base_dir.path());
toml::from_str(provider_toml).unwrap()
};
assert_eq!(
provider.auth,
Some(ModelProviderAuthInfo {
command: "./scripts/print-token".to_string(),
args: vec!["--format=text".to_string()],
timeout_ms: NonZeroU64::new(5_000).unwrap(),
refresh_interval_ms: 300_000,
cwd: AbsolutePathBuf::resolve_path_against_base(".", base_dir.path()).unwrap(),
})
);
}
#[test]
fn test_deserialize_provider_auth_config_allows_zero_refresh_interval() {
let base_dir = tempdir().unwrap();
let provider_toml = r#"
name = "Corp"
[auth]
command = "./scripts/print-token"
refresh_interval_ms = 0
"#;
let provider: ModelProviderInfo = {
let _guard = AbsolutePathBufGuard::new(base_dir.path());
toml::from_str(provider_toml).unwrap()
};
let auth = provider.auth.expect("auth config should deserialize");
assert_eq!(auth.refresh_interval_ms, 0);
assert_eq!(auth.refresh_interval(), None);
}
#[test]
fn provider_auth_validation_rejects_conflicting_fields() {
let provider = ModelProviderInfo {
name: "Corp".into(),
base_url: Some("https://example.com".into()),
env_key: Some("CORP_API_KEY".into()),
env_key_instructions: None,
experimental_bearer_token: None,
auth: Some(ModelProviderAuthInfo {
command: "./print-token".to_string(),
args: Vec::new(),
timeout_ms: NonZeroU64::new(5_000).unwrap(),
refresh_interval_ms: 300_000,
cwd: AbsolutePathBuf::current_dir().unwrap(),
}),
wire_api: WireApi::Responses,
query_params: None,
http_headers: None,
env_http_headers: None,
request_max_retries: None,
stream_max_retries: None,
stream_idle_timeout_ms: None,
requires_openai_auth: false,
openrouter: None,
};
assert_eq!(
provider.validate(),
Err("provider auth cannot be combined with env_key".to_string())
);
}
} }
@@ -50,6 +50,7 @@ fn provider_for(base_url: String) -> ModelProviderInfo {
env_key: None, env_key: None,
env_key_instructions: None, env_key_instructions: None,
experimental_bearer_token: None, experimental_bearer_token: None,
auth: None,
wire_api: WireApi::Responses, wire_api: WireApi::Responses,
query_params: None, query_params: None,
http_headers: None, http_headers: None,
+76
View File
@@ -1,6 +1,9 @@
use code_utils_absolute_path::AbsolutePathBuf;
use schemars::JsonSchema; use schemars::JsonSchema;
use serde::Deserialize; use serde::Deserialize;
use serde::Serialize; use serde::Serialize;
use std::num::NonZeroU64;
use std::time::Duration;
use strum_macros::Display; use strum_macros::Display;
use strum_macros::EnumIter; use strum_macros::EnumIter;
use ts_rs::TS; use ts_rs::TS;
@@ -130,6 +133,79 @@ pub enum ForcedLoginMethod {
Api, Api,
} }
const DEFAULT_PROVIDER_AUTH_TIMEOUT_MS: u64 = 5_000;
const DEFAULT_PROVIDER_AUTH_REFRESH_INTERVAL_MS: u64 = 300_000;
/// Configuration for obtaining a provider bearer token from a command.
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, JsonSchema)]
#[schemars(deny_unknown_fields)]
pub struct ModelProviderAuthInfo {
/// Command to execute. Bare names are resolved via `PATH`; paths are resolved against `cwd`.
pub command: String,
/// Command arguments.
#[serde(default)]
pub args: Vec<String>,
/// Maximum time to wait for the token command to exit successfully.
#[serde(default = "default_provider_auth_timeout_ms")]
pub timeout_ms: NonZeroU64,
/// Maximum age for the cached token before rerunning the command.
/// Set to `0` to disable proactive refresh and only rerun after a 401 retry path.
#[serde(default = "default_provider_auth_refresh_interval_ms")]
pub refresh_interval_ms: u64,
/// Working directory used when running the token command.
#[serde(default = "default_provider_auth_cwd")]
#[schemars(skip_serializing_if = "is_default_provider_auth_cwd")]
pub cwd: AbsolutePathBuf,
}
impl ModelProviderAuthInfo {
pub fn timeout(&self) -> Duration {
Duration::from_millis(self.timeout_ms.get())
}
pub fn refresh_interval(&self) -> Option<Duration> {
NonZeroU64::new(self.refresh_interval_ms).map(|value| Duration::from_millis(value.get()))
}
}
fn default_provider_auth_timeout_ms() -> NonZeroU64 {
non_zero_u64(
DEFAULT_PROVIDER_AUTH_TIMEOUT_MS,
"model_providers.<id>.auth.timeout_ms",
)
}
fn default_provider_auth_refresh_interval_ms() -> u64 {
DEFAULT_PROVIDER_AUTH_REFRESH_INTERVAL_MS
}
fn non_zero_u64(value: u64, field_name: &str) -> NonZeroU64 {
match NonZeroU64::new(value) {
Some(value) => value,
None => panic!("{field_name} must be non-zero"),
}
}
fn default_provider_auth_cwd() -> AbsolutePathBuf {
let deserializer = serde::de::value::StrDeserializer::<serde::de::value::Error>::new(".");
if let Ok(cwd) = AbsolutePathBuf::deserialize(deserializer) {
return cwd;
}
match AbsolutePathBuf::current_dir() {
Ok(cwd) => cwd,
Err(err) => panic!("provider auth cwd must resolve: {err}"),
}
}
fn is_default_provider_auth_cwd(path: &AbsolutePathBuf) -> bool {
path == &default_provider_auth_cwd()
}
/// Represents the trust level for a project directory. /// Represents the trust level for a project directory.
/// This determines the approval policy and sandbox mode applied. /// This determines the approval policy and sandbox mode applied.
#[derive(Debug, Serialize, Deserialize, Clone, Copy, PartialEq, Eq, Display, JsonSchema, TS)] #[derive(Debug, Serialize, Deserialize, Clone, Copy, PartialEq, Eq, Display, JsonSchema, TS)]