core: support dynamic auth tokens for model providers

This commit is contained in:
Michael Bolin
2026-03-30 16:51:30 -07:00
parent 4127419694
commit 770ed064cf
16 changed files with 778 additions and 35 deletions
+53
View File
@@ -816,10 +816,63 @@
},
"type": "object"
},
"ModelProviderAuthInfo": {
"additionalProperties": false,
"description": "Configuration for obtaining a provider bearer token from a command.",
"properties": {
"args": {
"default": [],
"description": "Command arguments.",
"items": {
"type": "string"
},
"type": "array"
},
"command": {
"description": "Command to execute. Bare names are resolved via `PATH`; paths are resolved against `cwd`.",
"type": "string"
},
"cwd": {
"allOf": [
{
"$ref": "#/definitions/AbsolutePathBuf"
}
],
"default": "/Users/mbolin/code/codex/codex-rs",
"description": "Working directory used when running the token command."
},
"refresh_interval_ms": {
"default": 300000,
"description": "Maximum age for the cached token before rerunning the command.",
"format": "uint64",
"minimum": 1.0,
"type": "integer"
},
"timeout_ms": {
"default": 5000,
"description": "Maximum time to wait for the token command to exit successfully.",
"format": "uint64",
"minimum": 1.0,
"type": "integer"
}
},
"required": [
"command"
],
"type": "object"
},
"ModelProviderInfo": {
"additionalProperties": false,
"description": "Serializable representation of a provider definition.",
"properties": {
"auth": {
"allOf": [
{
"$ref": "#/definitions/ModelProviderAuthInfo"
}
],
"description": "Command-backed bearer-token configuration for this provider."
},
"base_url": {
"description": "Base URL for the provider's OpenAI-compatible API.",
"type": "string"
+1
View File
@@ -64,6 +64,7 @@ mod tests {
env_key: Some("sk-should-not-leak".to_string()),
env_key_instructions: None,
experimental_bearer_token: None,
auth: None,
wire_api: crate::model_provider_info::WireApi::Responses,
query_params: None,
http_headers: None,
+6
View File
@@ -104,6 +104,7 @@ use crate::error::Result;
use crate::flags::CODEX_RS_SSE_FIXTURE;
use crate::model_provider_info::ModelProviderInfo;
use crate::model_provider_info::WireApi;
use crate::provider_auth::scoped_auth_manager_for_provider;
use crate::response_debug_context::extract_response_debug_context;
use crate::response_debug_context::extract_response_debug_context_from_api_error;
use crate::response_debug_context::telemetry_api_error_message;
@@ -261,6 +262,7 @@ impl ModelClient {
include_timing_metrics: bool,
beta_features_header: Option<String>,
) -> Self {
let auth_manager = scoped_auth_manager_for_provider(auth_manager, &provider);
let codex_api_key_env_enabled = auth_manager
.as_ref()
.is_some_and(|manager| manager.codex_api_key_env_enabled());
@@ -294,6 +296,10 @@ impl ModelClient {
}
}
pub(crate) fn auth_manager(&self) -> Option<Arc<AuthManager>> {
self.state.auth_manager.clone()
}
fn take_cached_websocket_session(&self) -> WebsocketSession {
let mut cached_websocket_session = self
.state
+21
View File
@@ -243,6 +243,26 @@ web_search = false
);
}
#[test]
fn rejects_provider_auth_with_env_key() {
let err = toml::from_str::<ConfigToml>(
r#"
[model_providers.corp]
name = "Corp"
env_key = "CORP_TOKEN"
[model_providers.corp.auth]
command = "print-token"
"#,
)
.unwrap_err();
assert!(
err.to_string()
.contains("model_providers.corp: provider auth cannot be combined with env_key")
);
}
#[test]
fn config_toml_deserializes_model_availability_nux() {
let toml = r#"
@@ -4315,6 +4335,7 @@ model_verbosity = "high"
wire_api: crate::WireApi::Responses,
env_key_instructions: None,
experimental_bearer_token: None,
auth: None,
query_params: None,
http_headers: None,
env_http_headers: None,
+14 -2
View File
@@ -1837,6 +1837,18 @@ Built-in providers cannot be overridden. Rename your custom provider (for exampl
}
}
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(())
}
fn deserialize_model_providers<'de, D>(
deserializer: D,
) -> Result<HashMap<String, ModelProviderInfo>, D::Error>
@@ -1844,7 +1856,7 @@ where
D: serde::Deserializer<'de>,
{
let model_providers = HashMap::<String, ModelProviderInfo>::deserialize(deserializer)?;
validate_reserved_model_provider_ids(&model_providers).map_err(serde::de::Error::custom)?;
validate_model_providers(&model_providers).map_err(serde::de::Error::custom)?;
Ok(model_providers)
}
@@ -1969,7 +1981,7 @@ impl Config {
codex_home: PathBuf,
config_layer_stack: ConfigLayerStack,
) -> std::io::Result<Self> {
validate_reserved_model_provider_ids(&cfg.model_providers)
validate_model_providers(&cfg.model_providers)
.map_err(|message| std::io::Error::new(std::io::ErrorKind::InvalidInput, message))?;
// Ensure that every field of ConfigRequirements is applied to the final
// Config.
+2
View File
@@ -65,6 +65,7 @@ pub mod utils;
pub use utils::path_utils;
pub mod personality_migration;
pub mod plugins;
mod provider_auth;
pub(crate) mod mentions {
pub(crate) use crate::plugins::build_connector_slug_counts;
pub(crate) use crate::plugins::build_skill_name_counts;
@@ -107,6 +108,7 @@ pub use client::X_RESPONSESAPI_INCLUDE_TIMING_METRICS_HEADER;
pub use model_provider_info::DEFAULT_LMSTUDIO_PORT;
pub use model_provider_info::DEFAULT_OLLAMA_PORT;
pub use model_provider_info::LMSTUDIO_OSS_PROVIDER_ID;
pub use model_provider_info::ModelProviderAuthInfo;
pub use model_provider_info::ModelProviderInfo;
pub use model_provider_info::OLLAMA_OSS_PROVIDER_ID;
pub use model_provider_info::OPENAI_PROVIDER_ID;
+110
View File
@@ -9,6 +9,7 @@ use crate::auth::AuthMode;
use crate::error::EnvVarError;
use codex_api::Provider as ApiProvider;
use codex_api::provider::RetryConfig as ApiRetryConfig;
use codex_utils_absolute_path::AbsolutePathBuf;
use http::HeaderMap;
use http::header::HeaderName;
use http::header::HeaderValue;
@@ -17,8 +18,11 @@ use serde::Deserialize;
use serde::Serialize;
use std::collections::HashMap;
use std::fmt;
use std::num::NonZeroU64;
use std::time::Duration;
const DEFAULT_PROVIDER_AUTH_TIMEOUT_MS: u64 = 5_000;
const DEFAULT_PROVIDER_AUTH_REFRESH_INTERVAL_MS: u64 = 300_000;
const DEFAULT_STREAM_IDLE_TIMEOUT_MS: u64 = 300_000;
const DEFAULT_STREAM_MAX_RETRIES: u64 = 5;
const DEFAULT_REQUEST_MAX_RETRIES: u64 = 4;
@@ -66,6 +70,73 @@ impl<'de> Deserialize<'de> for WireApi {
}
}
/// 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.
#[serde(default = "default_provider_auth_refresh_interval_ms")]
pub refresh_interval_ms: NonZeroU64,
/// Working directory used when running the token command.
#[serde(default = "default_provider_auth_cwd")]
pub cwd: AbsolutePathBuf,
}
impl ModelProviderAuthInfo {
pub(crate) fn timeout(&self) -> Duration {
Duration::from_millis(self.timeout_ms.get())
}
pub(crate) fn refresh_interval(&self) -> Duration {
Duration::from_millis(self.refresh_interval_ms.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() -> NonZeroU64 {
non_zero_u64(
DEFAULT_PROVIDER_AUTH_REFRESH_INTERVAL_MS,
"model_providers.<id>.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}"),
}
}
/// Serializable representation of a provider definition.
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, JsonSchema)]
#[schemars(deny_unknown_fields)]
@@ -86,6 +157,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,
@@ -130,6 +204,36 @@ pub struct 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(", ")
))
}
}
fn build_header_map(&self) -> crate::error::Result<HeaderMap> {
let capacity = self.http_headers.as_ref().map_or(0, HashMap::len)
+ self.env_http_headers.as_ref().map_or(0, HashMap::len);
@@ -246,6 +350,7 @@ impl ModelProviderInfo {
env_key: None,
env_key_instructions: None,
experimental_bearer_token: None,
auth: None,
wire_api: WireApi::Responses,
query_params: None,
http_headers: Some(
@@ -277,6 +382,10 @@ impl ModelProviderInfo {
pub fn is_openai(&self) -> bool {
self.name == OPENAI_PROVIDER_NAME
}
pub(crate) fn has_command_auth(&self) -> bool {
self.auth.is_some()
}
}
pub const DEFAULT_LMSTUDIO_PORT: u16 = 1234;
@@ -338,6 +447,7 @@ pub fn create_oss_provider_with_base_url(base_url: &str, wire_api: WireApi) -> M
env_key: None,
env_key_instructions: None,
experimental_bearer_token: None,
auth: None,
wire_api,
query_params: None,
http_headers: None,
@@ -1,5 +1,8 @@
use super::*;
use codex_utils_absolute_path::AbsolutePathBuf;
use codex_utils_absolute_path::AbsolutePathBufGuard;
use pretty_assertions::assert_eq;
use tempfile::tempdir;
#[test]
fn test_deserialize_ollama_model_provider_toml() {
@@ -13,6 +16,7 @@ base_url = "http://localhost:11434/v1"
env_key: None,
env_key_instructions: None,
experimental_bearer_token: None,
auth: None,
wire_api: WireApi::Responses,
query_params: None,
http_headers: None,
@@ -43,6 +47,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::Responses,
query_params: Some(maplit::hashmap! {
"api-version".to_string() => "2025-04-01-preview".to_string(),
@@ -76,6 +81,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::Responses,
query_params: None,
http_headers: Some(maplit::hashmap! {
@@ -121,3 +127,29 @@ supports_websockets = true
let provider: ModelProviderInfo = toml::from_str(provider_toml).unwrap();
assert_eq!(provider.websocket_connect_timeout_ms, Some(15_000));
}
#[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 _guard = AbsolutePathBufGuard::new(base_dir.path());
let provider: ModelProviderInfo = 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: default_provider_auth_timeout_ms(),
refresh_interval_ms: default_provider_auth_refresh_interval_ms(),
cwd: AbsolutePathBuf::resolve_path_against_base(".", base_dir.path()).unwrap(),
})
);
}
+71 -32
View File
@@ -4,6 +4,8 @@ use crate::api_bridge::map_api_error;
use crate::auth::AuthManager;
use crate::auth::AuthMode;
use crate::auth::CodexAuth;
use crate::auth::RefreshTokenError;
use crate::auth::UnauthorizedRecovery;
use crate::auth_env_telemetry::AuthEnvTelemetry;
use crate::auth_env_telemetry::collect_auth_env_telemetry;
use crate::config::Config;
@@ -14,6 +16,7 @@ use crate::model_provider_info::ModelProviderInfo;
use crate::models_manager::collaboration_mode_presets::CollaborationModesConfig;
use crate::models_manager::collaboration_mode_presets::builtin_collaboration_mode_presets;
use crate::models_manager::model_info;
use crate::provider_auth::scoped_auth_manager;
use crate::response_debug_context::extract_response_debug_context;
use crate::response_debug_context::telemetry_transport_error_message;
use crate::util::FeedbackRequestTags;
@@ -212,6 +215,7 @@ impl ModelsManager {
collaboration_modes_config: CollaborationModesConfig,
provider: ModelProviderInfo,
) -> Self {
let auth_manager = scoped_auth_manager(auth_manager, &provider);
let cache_path = codex_home.join(MODEL_CACHE_FILE);
let cache_manager = ModelsCacheManager::new(cache_path, DEFAULT_MODEL_CACHE_TTL);
let catalog_mode = if model_catalog.is_some() {
@@ -396,7 +400,9 @@ impl ModelsManager {
return Ok(());
}
if self.auth_manager.auth_mode() != Some(AuthMode::Chatgpt) {
if self.auth_manager.auth_mode() != Some(AuthMode::Chatgpt)
&& !self.provider.has_command_auth()
{
if matches!(
refresh_strategy,
RefreshStrategy::Offline | RefreshStrategy::OnlineIfUncached
@@ -431,39 +437,72 @@ impl ModelsManager {
async fn fetch_and_update_models(&self) -> CoreResult<()> {
let _timer =
codex_otel::start_global_timer("codex.remote_models.fetch_update.duration_ms", &[]);
let auth = self.auth_manager.auth().await;
let auth_mode = auth.as_ref().map(CodexAuth::auth_mode);
let api_provider = self.provider.to_api_provider(auth_mode)?;
let api_auth = auth_provider_from_auth(auth.clone(), &self.provider)?;
let auth_env = collect_auth_env_telemetry(
&self.provider,
self.auth_manager.codex_api_key_env_enabled(),
);
let transport = ReqwestTransport::new(build_reqwest_client());
let request_telemetry: Arc<dyn RequestTelemetry> = Arc::new(ModelsRequestTelemetry {
auth_mode: auth_mode.map(|mode| TelemetryAuthMode::from(mode).to_string()),
auth_header_attached: api_auth.auth_header_attached(),
auth_header_name: api_auth.auth_header_name(),
auth_env,
});
let client = ModelsClient::new(transport, api_provider, api_auth)
.with_telemetry(Some(request_telemetry));
let client_version = crate::models_manager::client_version_to_whole();
let (models, etag) = timeout(
MODELS_REFRESH_TIMEOUT,
client.list_models(&client_version, HeaderMap::new()),
)
.await
.map_err(|_| CodexErr::Timeout)?
.map_err(map_api_error)?;
let mut auth_recovery = self.auth_manager.unauthorized_recovery();
self.apply_remote_models(models.clone()).await;
*self.etag.write().await = etag.clone();
self.cache_manager
.persist_cache(&models, etag, client_version)
.await;
Ok(())
loop {
let auth = self.auth_manager.auth().await;
let auth_mode = auth.as_ref().map(CodexAuth::auth_mode);
let api_provider = self.provider.to_api_provider(auth_mode)?;
let api_auth = auth_provider_from_auth(auth.clone(), &self.provider)?;
let auth_env = collect_auth_env_telemetry(
&self.provider,
self.auth_manager.codex_api_key_env_enabled(),
);
let transport = ReqwestTransport::new(build_reqwest_client());
let request_telemetry: Arc<dyn RequestTelemetry> = Arc::new(ModelsRequestTelemetry {
auth_mode: auth_mode.map(|mode| TelemetryAuthMode::from(mode).to_string()),
auth_header_attached: api_auth.auth_header_attached(),
auth_header_name: api_auth.auth_header_name(),
auth_env,
});
let client = ModelsClient::new(transport, api_provider, api_auth)
.with_telemetry(Some(request_telemetry));
let result = timeout(
MODELS_REFRESH_TIMEOUT,
client.list_models(&client_version, HeaderMap::new()),
)
.await
.map_err(|_| CodexErr::Timeout)?
.map_err(map_api_error);
match result {
Ok((models, etag)) => {
self.apply_remote_models(models.clone()).await;
*self.etag.write().await = etag.clone();
self.cache_manager
.persist_cache(&models, etag, client_version)
.await;
return Ok(());
}
Err(
err @ CodexErr::UnexpectedStatus(crate::error::UnexpectedResponseError {
status,
..
}),
) if status == http::StatusCode::UNAUTHORIZED => {
if !Self::recover_after_unauthorized(&mut auth_recovery).await? {
return Err(err);
}
}
Err(err) => return Err(err),
}
}
}
async fn recover_after_unauthorized(
auth_recovery: &mut UnauthorizedRecovery,
) -> CoreResult<bool> {
if !auth_recovery.has_next() {
return Ok(false);
}
match auth_recovery.next().await {
Ok(_) => Ok(true),
Err(RefreshTokenError::Permanent(failed)) => Err(CodexErr::RefreshTokenFailed(failed)),
Err(RefreshTokenError::Transient(other)) => Err(CodexErr::Io(other)),
}
}
async fn get_etag(&self) -> Option<String> {
@@ -1,7 +1,9 @@
use super::*;
use crate::AuthManager;
use crate::CodexAuth;
use crate::auth::AuthCredentialsStoreMode;
use crate::config::ConfigBuilder;
use crate::model_provider_info::ModelProviderAuthInfo;
use crate::model_provider_info::WireApi;
use base64::Engine as _;
use chrono::Utc;
@@ -13,8 +15,10 @@ use http::StatusCode;
use pretty_assertions::assert_eq;
use serde_json::json;
use std::collections::BTreeMap;
use std::num::NonZeroU64;
use std::sync::Arc;
use std::sync::Mutex;
use tempfile::TempDir;
use tempfile::tempdir;
use tracing::Event;
use tracing::Subscriber;
@@ -24,7 +28,12 @@ use tracing_subscriber::layer::Context;
use tracing_subscriber::layer::SubscriberExt;
use tracing_subscriber::registry::LookupSpan;
use tracing_subscriber::util::SubscriberInitExt;
use wiremock::Mock;
use wiremock::MockServer;
use wiremock::ResponseTemplate;
use wiremock::matchers::header_regex;
use wiremock::matchers::method;
use wiremock::matchers::path;
fn remote_model(slug: &str, display: &str, priority: i32) -> ModelInfo {
remote_model_with_visibility(slug, display, priority, "list")
@@ -79,6 +88,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,
@@ -92,6 +102,90 @@ fn provider_for(base_url: String) -> ModelProviderInfo {
}
}
struct ProviderAuthScript {
tempdir: TempDir,
command: String,
args: Vec<String>,
}
impl ProviderAuthScript {
fn new(tokens: &[&str]) -> std::io::Result<Self> {
let tempdir = tempfile::tempdir()?;
let tokens_file = tempdir.path().join("tokens.txt");
std::fs::write(&tokens_file, format!("{}\n", tokens.join("\n")))?;
#[cfg(unix)]
let (command, args) = {
let script_path = tempdir.path().join("print-token.sh");
std::fs::write(
&script_path,
r#"#!/bin/sh
first_line=$(sed -n '1p' tokens.txt)
printf '%s\n' "$first_line"
tail -n +2 tokens.txt > tokens.next
mv tokens.next tokens.txt
"#,
)?;
let mut permissions = std::fs::metadata(&script_path)?.permissions();
{
use std::os::unix::fs::PermissionsExt;
permissions.set_mode(0o755);
}
std::fs::set_permissions(&script_path, permissions)?;
("./print-token.sh".to_string(), Vec::new())
};
#[cfg(windows)]
let (command, args) = {
let script_path = tempdir.path().join("print-token.ps1");
std::fs::write(
&script_path,
r#"$lines = Get-Content -Path tokens.txt
if ($lines.Count -eq 0) { exit 1 }
Write-Output $lines[0]
$lines | Select-Object -Skip 1 | Set-Content -Path tokens.txt
"#,
)?;
(
"powershell".to_string(),
vec![
"-NoProfile".to_string(),
"-ExecutionPolicy".to_string(),
"Bypass".to_string(),
"-File".to_string(),
".\\print-token.ps1".to_string(),
],
)
};
Ok(Self {
tempdir,
command,
args,
})
}
fn auth_config(&self) -> ModelProviderAuthInfo {
ModelProviderAuthInfo {
command: self.command.clone(),
args: self.args.clone(),
timeout_ms: non_zero_u64(/*value*/ 1_000),
refresh_interval_ms: non_zero_u64(/*value*/ 60_000),
cwd: match codex_utils_absolute_path::AbsolutePathBuf::try_from(self.tempdir.path()) {
Ok(cwd) => cwd,
Err(err) => panic!("tempdir should be absolute: {err}"),
},
}
}
}
fn non_zero_u64(value: u64) -> NonZeroU64 {
match NonZeroU64::new(value) {
Some(value) => value,
None => panic!("expected non-zero value: {value}"),
}
}
#[derive(Default)]
struct TagCollectorVisitor {
tags: BTreeMap<String, String>,
@@ -310,6 +404,57 @@ async fn refresh_available_models_sorts_by_priority() {
);
}
#[tokio::test]
async fn refresh_available_models_refreshes_provider_auth_after_401() {
let server = MockServer::start().await;
let auth_script = ProviderAuthScript::new(&["first-token", "second-token"]).unwrap();
let remote_models = vec![remote_model(
"provider-model",
"Provider",
/*priority*/ 0,
)];
Mock::given(method("GET"))
.and(path("/models"))
.and(header_regex("Authorization", "Bearer first-token"))
.respond_with(ResponseTemplate::new(401).set_body_string("unauthorized"))
.expect(1)
.mount(&server)
.await;
Mock::given(method("GET"))
.and(path("/models"))
.and(header_regex("Authorization", "Bearer second-token"))
.respond_with(
ResponseTemplate::new(200)
.insert_header("content-type", "application/json")
.set_body_json(ModelsResponse {
models: remote_models.clone(),
}),
)
.expect(1)
.mount(&server)
.await;
let codex_home = tempdir().expect("temp dir");
let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::from_api_key("unused"));
let provider = ModelProviderInfo {
auth: Some(auth_script.auth_config()),
..provider_for(server.uri())
};
let manager = ModelsManager::with_provider_for_tests(
codex_home.path().to_path_buf(),
auth_manager,
provider,
);
manager
.refresh_available_models(RefreshStrategy::Online)
.await
.expect("refresh succeeds");
assert_models_contain(&manager.get_remote_models().await, &remote_models);
}
#[tokio::test]
async fn refresh_available_models_uses_cache_when_fresh() {
let server = MockServer::start().await;
+6 -1
View File
@@ -452,7 +452,12 @@ async fn prepare_realtime_start(
params: ConversationStartParams,
) -> CodexResult<PreparedRealtimeConversationStart> {
let provider = sess.provider().await;
let auth = sess.services.auth_manager.auth().await;
let auth_manager = sess
.services
.model_client
.auth_manager()
.unwrap_or_else(|| Arc::clone(&sess.services.auth_manager));
let auth = auth_manager.auth().await;
let realtime_api_key = realtime_api_key(auth.as_ref(), &provider)?;
let mut api_provider = provider.to_api_provider(Some(crate::auth::AuthMode::ApiKey))?;
let config = sess.get_config().await;
+3
View File
@@ -46,6 +46,7 @@ async fn responses_stream_includes_subagent_header_on_review() {
env_key: None,
env_key_instructions: None,
experimental_bearer_token: None,
auth: None,
wire_api: WireApi::Responses,
query_params: None,
http_headers: None,
@@ -158,6 +159,7 @@ async fn responses_stream_includes_subagent_header_on_other() {
env_key: None,
env_key_instructions: None,
experimental_bearer_token: None,
auth: None,
wire_api: WireApi::Responses,
query_params: None,
http_headers: None,
@@ -265,6 +267,7 @@ async fn responses_respects_model_info_overrides_from_config() {
env_key: None,
env_key_instructions: None,
experimental_bearer_token: None,
auth: None,
wire_api: WireApi::Responses,
query_params: None,
http_headers: None,
+311
View File
@@ -1,5 +1,7 @@
use codex_core::AuthManager;
use codex_core::CodexAuth;
use codex_core::ModelClient;
use codex_core::ModelProviderAuthInfo;
use codex_core::ModelProviderInfo;
use codex_core::NewThread;
use codex_core::Prompt;
@@ -64,6 +66,7 @@ use futures::StreamExt;
use pretty_assertions::assert_eq;
use serde_json::json;
use std::io::Write;
use std::num::NonZeroU64;
use std::sync::Arc;
use tempfile::TempDir;
use uuid::Uuid;
@@ -143,6 +146,90 @@ fn write_auth_json(
fake_jwt
}
struct ProviderAuthCommandFixture {
tempdir: TempDir,
command: String,
args: Vec<String>,
}
impl ProviderAuthCommandFixture {
fn new(tokens: &[&str]) -> std::io::Result<Self> {
let tempdir = tempfile::tempdir()?;
let tokens_file = tempdir.path().join("tokens.txt");
std::fs::write(&tokens_file, format!("{}\n", tokens.join("\n")))?;
#[cfg(unix)]
let (command, args) = {
let script_path = tempdir.path().join("print-token.sh");
std::fs::write(
&script_path,
r#"#!/bin/sh
first_line=$(sed -n '1p' tokens.txt)
printf '%s\n' "$first_line"
tail -n +2 tokens.txt > tokens.next
mv tokens.next tokens.txt
"#,
)?;
let mut permissions = std::fs::metadata(&script_path)?.permissions();
{
use std::os::unix::fs::PermissionsExt;
permissions.set_mode(0o755);
}
std::fs::set_permissions(&script_path, permissions)?;
("./print-token.sh".to_string(), Vec::new())
};
#[cfg(windows)]
let (command, args) = {
let script_path = tempdir.path().join("print-token.ps1");
std::fs::write(
&script_path,
r#"$lines = Get-Content -Path tokens.txt
if ($lines.Count -eq 0) { exit 1 }
Write-Output $lines[0]
$lines | Select-Object -Skip 1 | Set-Content -Path tokens.txt
"#,
)?;
(
"powershell".to_string(),
vec![
"-NoProfile".to_string(),
"-ExecutionPolicy".to_string(),
"Bypass".to_string(),
"-File".to_string(),
".\\print-token.ps1".to_string(),
],
)
};
Ok(Self {
tempdir,
command,
args,
})
}
fn auth(&self) -> ModelProviderAuthInfo {
ModelProviderAuthInfo {
command: self.command.clone(),
args: self.args.clone(),
timeout_ms: non_zero_u64(/*value*/ 1_000),
refresh_interval_ms: non_zero_u64(/*value*/ 60_000),
cwd: match codex_utils_absolute_path::AbsolutePathBuf::try_from(self.tempdir.path()) {
Ok(cwd) => cwd,
Err(err) => panic!("tempdir should be absolute: {err}"),
},
}
}
}
fn non_zero_u64(value: u64) -> NonZeroU64 {
match NonZeroU64::new(value) {
Some(value) => value,
None => panic!("expected non-zero value: {value}"),
}
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn resume_includes_initial_messages_and_sends_prior_items() {
skip_if_no_network!();
@@ -659,6 +746,227 @@ async fn includes_conversation_id_and_model_headers_in_request() {
assert_eq!(request_authorization, "Bearer Test API Key");
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn provider_auth_command_supplies_bearer_token() {
skip_if_no_network!();
let server = MockServer::start().await;
let resp_mock = mount_sse_once(
&server,
sse(vec![ev_response_created("resp1"), ev_completed("resp1")]),
)
.await;
let auth_fixture = ProviderAuthCommandFixture::new(&["command-token"]).unwrap();
let provider = ModelProviderInfo {
name: "corp".into(),
base_url: Some(format!("{}/v1", server.uri())),
env_key: None,
env_key_instructions: None,
experimental_bearer_token: None,
auth: Some(auth_fixture.auth()),
wire_api: WireApi::Responses,
query_params: None,
http_headers: None,
env_http_headers: None,
request_max_retries: Some(0),
stream_max_retries: Some(0),
stream_idle_timeout_ms: Some(5_000),
websocket_connect_timeout_ms: None,
requires_openai_auth: false,
supports_websockets: false,
};
let codex_home = TempDir::new().unwrap();
let mut config = load_default_config_for_test(&codex_home).await;
config.model_provider_id = provider.name.clone();
config.model_provider = provider.clone();
let effort = config.model_reasoning_effort;
let summary = config.model_reasoning_summary;
let model = codex_core::test_support::get_model_offline(config.model.as_deref());
config.model = Some(model.clone());
let config = Arc::new(config);
let model_info =
codex_core::test_support::construct_model_info_offline(model.as_str(), &config);
let conversation_id = ThreadId::new();
let session_telemetry = SessionTelemetry::new(
conversation_id,
model.as_str(),
model_info.slug.as_str(),
/*account_id*/ None,
Some("test@test.com".to_string()),
/*auth_mode*/ None,
"test_originator".to_string(),
/*log_user_prompts*/ false,
"test".to_string(),
SessionSource::Exec,
);
let client = ModelClient::new(
Some(AuthManager::from_auth_for_testing(CodexAuth::from_api_key(
"unused-api-key",
))),
conversation_id,
provider,
SessionSource::Exec,
config.model_verbosity,
/*enable_request_compression*/ false,
/*include_timing_metrics*/ false,
/*beta_features_header*/ None,
);
let mut client_session = client.new_session();
let mut prompt = Prompt::default();
prompt.input.push(ResponseItem::Message {
id: None,
role: "user".to_string(),
content: vec![ContentItem::InputText {
text: "hello".to_string(),
}],
end_turn: None,
phase: None,
});
let mut stream = client_session
.stream(
&prompt,
&model_info,
&session_telemetry,
effort,
summary.unwrap_or(ReasoningSummary::Auto),
/*service_tier*/ None,
/*turn_metadata_header*/ None,
)
.await
.expect("responses stream to start");
while let Some(event) = stream.next().await {
if let Ok(ResponseEvent::Completed { .. }) = event {
break;
}
}
let request = resp_mock.single_request();
let request_authorization = request
.header("authorization")
.expect("authorization header");
assert_eq!(request_authorization, "Bearer command-token");
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn provider_auth_command_refreshes_after_401() {
skip_if_no_network!();
let server = MockServer::start().await;
let auth_fixture = ProviderAuthCommandFixture::new(&["first-token", "second-token"]).unwrap();
Mock::given(method("POST"))
.and(path("/v1/responses"))
.and(header_regex("Authorization", "Bearer first-token"))
.respond_with(ResponseTemplate::new(401).set_body_string("unauthorized"))
.expect(1)
.mount(&server)
.await;
Mock::given(method("POST"))
.and(path("/v1/responses"))
.and(header_regex("Authorization", "Bearer second-token"))
.respond_with(
ResponseTemplate::new(200)
.insert_header("content-type", "text/event-stream")
.set_body_raw(
sse(vec![ev_response_created("resp1"), ev_completed("resp1")]),
"text/event-stream",
),
)
.expect(1)
.mount(&server)
.await;
let provider = ModelProviderInfo {
name: "corp".into(),
base_url: Some(format!("{}/v1", server.uri())),
env_key: None,
env_key_instructions: None,
experimental_bearer_token: None,
auth: Some(auth_fixture.auth()),
wire_api: WireApi::Responses,
query_params: None,
http_headers: None,
env_http_headers: None,
request_max_retries: Some(0),
stream_max_retries: Some(0),
stream_idle_timeout_ms: Some(5_000),
websocket_connect_timeout_ms: None,
requires_openai_auth: false,
supports_websockets: false,
};
let codex_home = TempDir::new().unwrap();
let mut config = load_default_config_for_test(&codex_home).await;
config.model_provider_id = provider.name.clone();
config.model_provider = provider.clone();
let effort = config.model_reasoning_effort;
let summary = config.model_reasoning_summary;
let model = codex_core::test_support::get_model_offline(config.model.as_deref());
config.model = Some(model.clone());
let config = Arc::new(config);
let model_info =
codex_core::test_support::construct_model_info_offline(model.as_str(), &config);
let conversation_id = ThreadId::new();
let session_telemetry = SessionTelemetry::new(
conversation_id,
model.as_str(),
model_info.slug.as_str(),
/*account_id*/ None,
Some("test@test.com".to_string()),
/*auth_mode*/ None,
"test_originator".to_string(),
/*log_user_prompts*/ false,
"test".to_string(),
SessionSource::Exec,
);
let client = ModelClient::new(
Some(AuthManager::from_auth_for_testing(CodexAuth::from_api_key(
"unused-api-key",
))),
conversation_id,
provider,
SessionSource::Exec,
config.model_verbosity,
/*enable_request_compression*/ false,
/*include_timing_metrics*/ false,
/*beta_features_header*/ None,
);
let mut client_session = client.new_session();
let mut prompt = Prompt::default();
prompt.input.push(ResponseItem::Message {
id: None,
role: "user".to_string(),
content: vec![ContentItem::InputText {
text: "hello".to_string(),
}],
end_turn: None,
phase: None,
});
let mut stream = client_session
.stream(
&prompt,
&model_info,
&session_telemetry,
effort,
summary.unwrap_or(ReasoningSummary::Auto),
/*service_tier*/ None,
/*turn_metadata_header*/ None,
)
.await
.expect("responses stream to start");
while let Some(event) = stream.next().await {
if let Ok(ResponseEvent::Completed { .. }) = event {
break;
}
}
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn includes_base_instructions_override_in_request() {
skip_if_no_network!();
@@ -1796,6 +2104,7 @@ async fn azure_responses_request_includes_store_and_reasoning_ids() {
env_key: None,
env_key_instructions: None,
experimental_bearer_token: None,
auth: None,
wire_api: WireApi::Responses,
query_params: None,
http_headers: None,
@@ -2396,6 +2705,7 @@ async fn azure_overrides_assign_properties_used_for_responses_url() {
// Reuse the existing environment variable to avoid using unsafe code
env_key: Some(existing_env_var_with_random_value.to_string()),
experimental_bearer_token: None,
auth: None,
query_params: Some(std::collections::HashMap::from([(
"api-version".to_string(),
"2025-04-01-preview".to_string(),
@@ -2486,6 +2796,7 @@ async fn env_var_overrides_loaded_auth() {
)])),
env_key_instructions: None,
experimental_bearer_token: None,
auth: None,
wire_api: WireApi::Responses,
http_headers: Some(std::collections::HashMap::from([(
"Custom-Header".to_string(),
@@ -1674,6 +1674,7 @@ fn websocket_provider_with_connect_timeout(
env_key: None,
env_key_instructions: None,
experimental_bearer_token: None,
auth: None,
wire_api: WireApi::Responses,
query_params: None,
http_headers: None,
@@ -69,6 +69,7 @@ async fn continue_after_stream_error() {
env_key: Some("PATH".into()),
env_key_instructions: None,
experimental_bearer_token: None,
auth: None,
wire_api: WireApi::Responses,
query_params: None,
http_headers: None,
@@ -53,6 +53,7 @@ async fn retries_on_early_close() {
env_key: Some("PATH".into()),
env_key_instructions: None,
experimental_bearer_token: None,
auth: None,
wire_api: WireApi::Responses,
query_params: None,
http_headers: None,