mirror of
https://github.com/hanzoai/dev.git
synced 2026-10-25 03:32:32 +00:00
feat(core/providers): support command-backed auth
This commit is contained in:
@@ -389,8 +389,9 @@ pub(crate) async fn stream_chat_completions(
|
||||
loop {
|
||||
attempt += 1;
|
||||
|
||||
let auth = auth_manager.as_ref().and_then(|m| m.auth());
|
||||
let mut req_builder = provider.create_request_builder(client, &auth).await?;
|
||||
let base_auth = auth_manager.as_ref().and_then(|m| m.auth());
|
||||
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 auth.mode.is_chatgpt() {
|
||||
@@ -454,6 +455,19 @@ pub(crate) async fn stream_chat_completions(
|
||||
}
|
||||
Ok(res) => {
|
||||
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()) {
|
||||
let body = (res.text().await).unwrap_or_default();
|
||||
if let Ok(logger) = debug_logger.lock() {
|
||||
|
||||
@@ -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 url = reqwest::Url::parse(&endpoint).map_err(|err| {
|
||||
@@ -798,7 +799,12 @@ impl ModelClient {
|
||||
};
|
||||
let mut req_builder = self
|
||||
.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?;
|
||||
|
||||
let has_beta_header = req_builder
|
||||
@@ -1226,7 +1232,8 @@ impl ModelClient {
|
||||
let mut auth_refresh_error: Option<RefreshTokenError> = None;
|
||||
|
||||
// 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!(
|
||||
"POST to {}: {}",
|
||||
@@ -1236,7 +1243,7 @@ impl ModelClient {
|
||||
|
||||
let mut req_builder = self
|
||||
.provider
|
||||
.create_request_builder(&self.client, &auth)
|
||||
.create_request_builder_with_auth(&self.client, &auth)
|
||||
.await?;
|
||||
|
||||
let has_beta_header = req_builder
|
||||
@@ -1425,7 +1432,9 @@ impl ModelClient {
|
||||
.and_then(|raw| parse_retry_after_header(raw, now));
|
||||
|
||||
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 {
|
||||
Ok(Some(_)) => {}
|
||||
Ok(None) => {
|
||||
@@ -1437,7 +1446,7 @@ impl ModelClient {
|
||||
auth_refresh_error = Some(err);
|
||||
}
|
||||
}
|
||||
} else {
|
||||
} else if auth.is_none() {
|
||||
auth_refresh_error = Some(RefreshTokenError::permanent(
|
||||
"Authentication manager unavailable; please log in again.",
|
||||
));
|
||||
@@ -1849,10 +1858,11 @@ impl ModelClient {
|
||||
let mut request_id = String::new();
|
||||
|
||||
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
|
||||
.provider
|
||||
.create_compact_request_builder(&self.client, &auth)
|
||||
.create_compact_request_builder_with_auth(&self.client, &auth)
|
||||
.await?;
|
||||
|
||||
// Ensure Responses API beta header is present for compact calls. Mirror the
|
||||
@@ -2991,6 +3001,7 @@ mod tests {
|
||||
env_key: None,
|
||||
env_key_instructions: None,
|
||||
experimental_bearer_token: None,
|
||||
auth: None,
|
||||
wire_api: WireApi::Responses,
|
||||
query_params: None,
|
||||
http_headers: None,
|
||||
@@ -3039,6 +3050,7 @@ mod tests {
|
||||
env_key: None,
|
||||
env_key_instructions: None,
|
||||
experimental_bearer_token: None,
|
||||
auth: None,
|
||||
wire_api: WireApi::Responses,
|
||||
query_params: None,
|
||||
http_headers: None,
|
||||
@@ -3089,6 +3101,7 @@ mod tests {
|
||||
env_key: None,
|
||||
env_key_instructions: None,
|
||||
experimental_bearer_token: None,
|
||||
auth: None,
|
||||
wire_api: WireApi::Responses,
|
||||
query_params: None,
|
||||
http_headers: Some(headers),
|
||||
@@ -3236,6 +3249,7 @@ mod tests {
|
||||
env_key: Some("TEST_API_KEY".to_string()),
|
||||
env_key_instructions: None,
|
||||
experimental_bearer_token: None,
|
||||
auth: None,
|
||||
wire_api: WireApi::Responses,
|
||||
query_params: None,
|
||||
http_headers: None,
|
||||
@@ -3302,6 +3316,7 @@ mod tests {
|
||||
env_key: Some("TEST_API_KEY".to_string()),
|
||||
env_key_instructions: None,
|
||||
experimental_bearer_token: None,
|
||||
auth: None,
|
||||
wire_api: WireApi::Responses,
|
||||
query_params: None,
|
||||
http_headers: None,
|
||||
@@ -3354,6 +3369,7 @@ mod tests {
|
||||
env_key: Some("TEST_API_KEY".to_string()),
|
||||
env_key_instructions: None,
|
||||
experimental_bearer_token: None,
|
||||
auth: None,
|
||||
wire_api: WireApi::Responses,
|
||||
query_params: None,
|
||||
http_headers: None,
|
||||
@@ -3449,6 +3465,7 @@ mod tests {
|
||||
env_key: Some("TEST_API_KEY".to_string()),
|
||||
env_key_instructions: None,
|
||||
experimental_bearer_token: None,
|
||||
auth: None,
|
||||
wire_api: WireApi::Responses,
|
||||
query_params: None,
|
||||
http_headers: None,
|
||||
@@ -3556,6 +3573,7 @@ mod tests {
|
||||
env_key: Some("TEST_API_KEY".to_string()),
|
||||
env_key_instructions: None,
|
||||
experimental_bearer_token: None,
|
||||
auth: None,
|
||||
wire_api: WireApi::Responses,
|
||||
query_params: None,
|
||||
http_headers: None,
|
||||
@@ -3796,6 +3814,7 @@ mod tests {
|
||||
env_key: Some("TEST_API_KEY".to_string()),
|
||||
env_key_instructions: None,
|
||||
experimental_bearer_token: None,
|
||||
auth: None,
|
||||
wire_api: WireApi::Responses,
|
||||
query_params: None,
|
||||
http_headers: None,
|
||||
@@ -3827,6 +3846,7 @@ mod tests {
|
||||
env_key: Some("TEST_API_KEY".to_string()),
|
||||
env_key_instructions: None,
|
||||
experimental_bearer_token: None,
|
||||
auth: None,
|
||||
wire_api: WireApi::Responses,
|
||||
query_params: None,
|
||||
http_headers: None,
|
||||
@@ -3861,6 +3881,7 @@ mod tests {
|
||||
env_key: Some("TEST_API_KEY".to_string()),
|
||||
env_key_instructions: None,
|
||||
experimental_bearer_token: None,
|
||||
auth: None,
|
||||
wire_api: WireApi::Responses,
|
||||
query_params: None,
|
||||
http_headers: None,
|
||||
@@ -3892,6 +3913,7 @@ mod tests {
|
||||
env_key: Some("TEST_API_KEY".to_string()),
|
||||
env_key_instructions: None,
|
||||
experimental_bearer_token: None,
|
||||
auth: None,
|
||||
wire_api: WireApi::Responses,
|
||||
query_params: None,
|
||||
http_headers: None,
|
||||
@@ -3923,6 +3945,7 @@ mod tests {
|
||||
env_key: Some("TEST_API_KEY".to_string()),
|
||||
env_key_instructions: None,
|
||||
experimental_bearer_token: None,
|
||||
auth: None,
|
||||
wire_api: WireApi::Responses,
|
||||
query_params: None,
|
||||
http_headers: None,
|
||||
|
||||
@@ -65,6 +65,36 @@ mod validation;
|
||||
use defaults::{default_responses_originator, default_review_model, default_true_local};
|
||||
|
||||
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 defaults::set_default_originator;
|
||||
@@ -1099,6 +1129,8 @@ impl Config {
|
||||
}
|
||||
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);
|
||||
// Merge user-defined providers into the built-in list.
|
||||
for (key, provider) in cfg.model_providers.into_iter() {
|
||||
@@ -2372,6 +2404,7 @@ model_verbosity = "high"
|
||||
wire_api: crate::WireApi::Chat,
|
||||
env_key_instructions: None,
|
||||
experimental_bearer_token: None,
|
||||
auth: None,
|
||||
query_params: None,
|
||||
http_headers: None,
|
||||
env_http_headers: None,
|
||||
|
||||
@@ -8,12 +8,21 @@
|
||||
|
||||
use crate::CodexAuth;
|
||||
use crate::error::CodexErr;
|
||||
use code_protocol::config_types::ModelProviderAuthInfo;
|
||||
use serde::Deserialize;
|
||||
use serde::Serialize;
|
||||
use serde_json::Value;
|
||||
use std::collections::{BTreeMap, HashMap};
|
||||
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::Instant;
|
||||
use tokio::process::Command;
|
||||
use crate::error::EnvVarError;
|
||||
const DEFAULT_STREAM_IDLE_TIMEOUT_MS: u64 = 300_000;
|
||||
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`.
|
||||
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
|
||||
/// implement the classic OpenAI Chat Completions JSON schema, whereas OpenAI
|
||||
/// 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.
|
||||
pub experimental_bearer_token: Option<String>,
|
||||
|
||||
/// Command-backed bearer-token configuration for this provider.
|
||||
pub auth: Option<ModelProviderAuthInfo>,
|
||||
|
||||
/// Which wire protocol this provider expects.
|
||||
#[serde(default)]
|
||||
pub wire_api: WireApi,
|
||||
@@ -189,6 +230,49 @@ pub struct OpenRouterMaxPrice {
|
||||
}
|
||||
|
||||
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
|
||||
/// reqwest Client applying:
|
||||
/// • provider-specific headers (static + env based)
|
||||
@@ -202,13 +286,22 @@ impl ModelProviderInfo {
|
||||
client: &'a reqwest::Client,
|
||||
auth: &Option<CodexAuth>,
|
||||
) -> 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);
|
||||
|
||||
if let Some(auth) = effective_auth.as_ref() {
|
||||
if let Some(auth) = auth.as_ref() {
|
||||
builder = builder.bearer_auth(auth.get_token().await?);
|
||||
}
|
||||
|
||||
@@ -224,11 +317,23 @@ impl ModelProviderInfo {
|
||||
method: reqwest::Method,
|
||||
url: reqwest::Url,
|
||||
) -> 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);
|
||||
|
||||
if let Some(auth) = effective_auth.as_ref() {
|
||||
if let Some(auth) = auth.as_ref() {
|
||||
builder = builder.bearer_auth(auth.get_token().await?);
|
||||
}
|
||||
|
||||
@@ -239,28 +344,37 @@ impl ModelProviderInfo {
|
||||
&'a self,
|
||||
client: &'a reqwest::Client,
|
||||
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> {
|
||||
if !matches!(self.wire_api, WireApi::Responses | WireApi::ResponsesWebsocket) {
|
||||
return Err(CodexErr::UnsupportedOperation(
|
||||
"Compaction endpoint requires Responses API providers".to_string(),
|
||||
));
|
||||
}
|
||||
let effective_auth = self.effective_auth(auth)?;
|
||||
let url = self.get_compact_url(&effective_auth).ok_or_else(|| {
|
||||
let url = self.get_compact_url(auth).ok_or_else(|| {
|
||||
CodexErr::UnsupportedOperation(
|
||||
"Compaction endpoint requires Responses API providers".to_string(),
|
||||
)
|
||||
})?;
|
||||
|
||||
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?);
|
||||
}
|
||||
|
||||
Ok(self.apply_http_headers(builder))
|
||||
}
|
||||
|
||||
fn effective_auth(
|
||||
pub(crate) async fn effective_auth(
|
||||
&self,
|
||||
auth: &Option<CodexAuth>,
|
||||
) -> crate::error::Result<Option<CodexAuth>> {
|
||||
@@ -273,6 +387,13 @@ impl ModelProviderInfo {
|
||||
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() {
|
||||
Ok(Some(key)) => Ok(Some(CodexAuth::from_api_key(&key))),
|
||||
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;
|
||||
|
||||
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_instructions: None,
|
||||
experimental_bearer_token: None,
|
||||
auth: None,
|
||||
wire_api: wire_api_override_from_env("OPENAI_WIRE_API")
|
||||
.unwrap_or(WireApi::Responses),
|
||||
query_params: None,
|
||||
@@ -567,6 +783,7 @@ pub fn create_oss_provider_with_base_url(base_url: &str) -> ModelProviderInfo {
|
||||
env_key: None,
|
||||
env_key_instructions: None,
|
||||
experimental_bearer_token: None,
|
||||
auth: None,
|
||||
wire_api: WireApi::Chat,
|
||||
query_params: None,
|
||||
http_headers: None,
|
||||
@@ -594,7 +811,11 @@ fn matches_azure_responses_base_url(base_url: &str) -> bool {
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use code_utils_absolute_path::AbsolutePathBuf;
|
||||
use code_utils_absolute_path::AbsolutePathBufGuard;
|
||||
use pretty_assertions::assert_eq;
|
||||
use std::num::NonZeroU64;
|
||||
use tempfile::tempdir;
|
||||
|
||||
#[test]
|
||||
fn test_deserialize_ollama_model_provider_toml() {
|
||||
@@ -608,6 +829,7 @@ base_url = "http://localhost:11434/v1"
|
||||
env_key: None,
|
||||
env_key_instructions: None,
|
||||
experimental_bearer_token: None,
|
||||
auth: None,
|
||||
wire_api: WireApi::Chat,
|
||||
query_params: 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_instructions: None,
|
||||
experimental_bearer_token: None,
|
||||
auth: None,
|
||||
wire_api: WireApi::Chat,
|
||||
query_params: Some(maplit::hashmap! {
|
||||
"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_instructions: None,
|
||||
experimental_bearer_token: None,
|
||||
auth: None,
|
||||
wire_api: WireApi::Chat,
|
||||
query_params: None,
|
||||
http_headers: Some(maplit::hashmap! {
|
||||
@@ -697,6 +921,7 @@ env_http_headers = { "X-Example-Env-Header" = "EXAMPLE_ENV_VAR" }
|
||||
env_key: None,
|
||||
env_key_instructions: None,
|
||||
experimental_bearer_token: None,
|
||||
auth: None,
|
||||
wire_api: WireApi::Responses,
|
||||
query_params: None,
|
||||
http_headers: None,
|
||||
@@ -731,6 +956,7 @@ env_http_headers = { "X-Example-Env-Header" = "EXAMPLE_ENV_VAR" }
|
||||
env_key: None,
|
||||
env_key_instructions: None,
|
||||
experimental_bearer_token: None,
|
||||
auth: None,
|
||||
wire_api: WireApi::Responses,
|
||||
query_params: 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());
|
||||
}
|
||||
|
||||
#[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_instructions: None,
|
||||
experimental_bearer_token: None,
|
||||
auth: None,
|
||||
wire_api: WireApi::Responses,
|
||||
query_params: None,
|
||||
http_headers: None,
|
||||
|
||||
@@ -1,6 +1,9 @@
|
||||
use code_utils_absolute_path::AbsolutePathBuf;
|
||||
use schemars::JsonSchema;
|
||||
use serde::Deserialize;
|
||||
use serde::Serialize;
|
||||
use std::num::NonZeroU64;
|
||||
use std::time::Duration;
|
||||
use strum_macros::Display;
|
||||
use strum_macros::EnumIter;
|
||||
use ts_rs::TS;
|
||||
@@ -130,6 +133,79 @@ pub enum ForcedLoginMethod {
|
||||
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.
|
||||
/// This determines the approval policy and sandbox mode applied.
|
||||
#[derive(Debug, Serialize, Deserialize, Clone, Copy, PartialEq, Eq, Display, JsonSchema, TS)]
|
||||
|
||||
Reference in New Issue
Block a user