diff --git a/codex-rs/config/src/bedrock_runtime_tests.rs b/codex-rs/config/src/bedrock_runtime_tests.rs index 8bba3aca39..00b86a8028 100644 --- a/codex-rs/config/src/bedrock_runtime_tests.rs +++ b/codex-rs/config/src/bedrock_runtime_tests.rs @@ -23,6 +23,7 @@ region = "us-west-2" Some(ModelProviderAwsAuthInfo { profile: Some("runtime-profile".to_string()), region: Some("us-west-2".to_string()), + auth_refresh: None, }) ); } diff --git a/codex-rs/core/config.schema.json b/codex-rs/core/config.schema.json index 7c3c0edd75..74905788f3 100644 --- a/codex-rs/core/config.schema.json +++ b/codex-rs/core/config.schema.json @@ -352,6 +352,35 @@ }, "type": "object" }, + "AwsAuthRefreshConfig": { + "additionalProperties": false, + "description": "Command used to refresh AWS credentials for a model provider.", + "properties": { + "args": { + "default": [], + "description": "Arguments passed to the refresh command.", + "items": { + "type": "string" + }, + "type": "array" + }, + "command": { + "description": "Executable to invoke directly, without a shell.", + "type": "string" + }, + "timeout_ms": { + "default": 300000, + "description": "Maximum time to wait for the refresh command to complete.", + "format": "uint64", + "minimum": 1.0, + "type": "integer" + } + }, + "required": [ + "command" + ], + "type": "object" + }, "BundledSkillsConfig": { "additionalProperties": false, "properties": { @@ -1871,6 +1900,14 @@ "additionalProperties": false, "description": "AWS SigV4 auth configuration for a model provider.", "properties": { + "auth_refresh": { + "allOf": [ + { + "$ref": "#/definitions/AwsAuthRefreshConfig" + } + ], + "description": "Optional command used to reauthenticate after a refreshable AWS auth failure." + }, "profile": { "description": "AWS profile name to use. When unset, the AWS SDK default chain decides.", "type": "string" diff --git a/codex-rs/core/src/config/config_tests.rs b/codex-rs/core/src/config/config_tests.rs index d3e147cd60..4dea8dafac 100644 --- a/codex-rs/core/src/config/config_tests.rs +++ b/codex-rs/core/src/config/config_tests.rs @@ -1085,7 +1085,7 @@ supports_websockets = true assert_eq!(err.kind(), std::io::ErrorKind::InvalidData); assert!(err.to_string().contains( - "model_providers.amazon-bedrock only supports changing `base_url`, `auth`, `http_headers`, `aws.profile`, and `aws.region`; other non-default provider fields are not supported" + "model_providers.amazon-bedrock only supports changing `base_url`, `auth`, `http_headers`, `aws.profile`, `aws.region`, and `aws.auth_refresh`; other non-default provider fields are not supported" )); } diff --git a/codex-rs/core/tests/suite/external_auth.rs b/codex-rs/core/tests/suite/external_auth.rs index 33adde4e7c..9fad8edd71 100644 --- a/codex-rs/core/tests/suite/external_auth.rs +++ b/codex-rs/core/tests/suite/external_auth.rs @@ -3,6 +3,9 @@ use codex_login::CodexAuth; use codex_login::ExternalAuth; use codex_login::ExternalAuthFuture; use codex_login::ExternalAuthRefreshContext; +use codex_model_provider_info::AwsAuthRefreshConfig; +use codex_model_provider_info::ModelProviderAwsAuthInfo; +use codex_model_provider_info::ModelProviderInfo; use codex_model_provider_info::WireApi; use codex_model_provider_info::create_oss_provider_with_base_url; use core_test_support::responses::ev_completed; @@ -16,8 +19,10 @@ use http::HeaderMap; use http::HeaderValue; use http::header::AUTHORIZATION; use pretty_assertions::assert_eq; +use std::num::NonZeroU64; use std::sync::Arc; use std::sync::Mutex; +use tokio::process::Command; use wiremock::Mock; use wiremock::ResponseTemplate; use wiremock::matchers::header; @@ -175,6 +180,99 @@ async fn custom_provider_uses_explicit_bearer_without_ambient_account() -> anyho Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn amazon_bedrock_aws_auth_refresh_resigns() -> anyhow::Result<()> { + skip_if_no_network!(Ok(())); + + const TEST_NAME: &str = "suite::external_auth::amazon_bedrock_aws_auth_refresh_resigns"; + const SUBPROCESS_ENV: &str = "CODEX_BEDROCK_AWS_REFRESH_TEST"; + const HELPER_ARG: &str = "CODEX_BEDROCK_AWS_REFRESH_COMMAND"; + const OLD: &str = "[default]\naws_access_key_id=OLD\naws_secret_access_key=s\n"; + const NEW: &str = "[default]\naws_access_key_id=NEW\naws_secret_access_key=s\n"; + + if std::env::var_os(SUBPROCESS_ENV).is_none() { + let fixture = tempfile::tempdir()?; + let test_executable = std::env::current_exe()?; + let aws_executable = fixture + .path() + .join(format!("aws{}", std::env::consts::EXE_SUFFIX)); + std::fs::hard_link(&test_executable, &aws_executable) + .or_else(|_| std::fs::copy(&test_executable, &aws_executable).map(|_| ()))?; + let inherited_path = std::env::var_os("PATH").unwrap_or_default(); + let mut paths = std::env::split_paths(&inherited_path).collect::>(); + paths.insert(/*index*/ 0, fixture.path().to_path_buf()); + let credentials = fixture.path().join("credentials"); + std::fs::write(&credentials, OLD)?; + let mut command = Command::new(test_executable); + command + .arg("--exact") + .arg(TEST_NAME) + .env(SUBPROCESS_ENV, "1") + .env("PATH", std::env::join_paths(paths)?) + .env("AWS_SHARED_CREDENTIALS_FILE", &credentials) + .env("AWS_CONFIG_FILE", &credentials) + .env("AWS_EC2_METADATA_DISABLED", "true") + .env_remove("AWS_ACCESS_KEY_ID") + .env_remove("AWS_SECRET_ACCESS_KEY") + .env_remove("AWS_BEARER_TOKEN_BEDROCK") + .env_remove("AWS_WEB_IDENTITY_TOKEN_FILE") + .env_remove("AWS_ROLE_ARN"); + let output = command.output().await?; + assert!(output.status.success(), "{output:?}"); + return Ok(()); + } + + if std::env::args().any(|argument| argument == HELPER_ARG) { + std::fs::write(std::env::var("AWS_SHARED_CREDENTIALS_FILE")?, NEW)?; + return Ok(()); + } + + let server = start_mock_server().await; + Mock::given(method("POST")) + .and(path("/v1/responses")) + .respond_with(|request: &wiremock::Request| { + let authorization = request.headers["authorization"] + .to_str() + .expect("SigV4 authorization should be valid"); + if authorization.contains("Credential=OLD/") { + ResponseTemplate::new(401).set_body_string("ExpiredTokenException") + } else { + assert!(authorization.contains("Credential=NEW/")); + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(sse(vec![ev_response_created("r"), ev_completed("r")])) + } + }) + .expect(2) + .mount(&server) + .await; + + let mut provider = + ModelProviderInfo::create_amazon_bedrock_provider(Some(ModelProviderAwsAuthInfo { + profile: Some("default".to_string()), + region: Some("us-east-1".to_string()), + auth_refresh: Some(AwsAuthRefreshConfig { + command: "aws".to_string(), + args: Vec::from(["--exact", TEST_NAME, "--skip", HELPER_ARG].map(str::to_string)), + timeout_ms: NonZeroU64::new(30_000).expect("timeout should be non-zero"), + }), + })); + provider.base_url = Some(format!("{}/v1", server.uri())); + provider.request_max_retries = Some(0); + provider.stream_max_retries = Some(0); + + let mut builder = test_codex() + .with_model("openai.gpt-5.5") + .with_config(move |config| { + config.model_provider_id = provider.name.clone(); + config.model_provider = provider; + }); + let test = builder.build_with_auto_env(&server).await?; + test.submit_turn("hello").await?; + server.verify().await; + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn external_auth_401_retry_uses_refreshed_chatgpt_headers() -> anyhow::Result<()> { skip_if_no_network!(Ok(())); diff --git a/codex-rs/model-provider-info/src/lib.rs b/codex-rs/model-provider-info/src/lib.rs index 14f55c2042..3d58aadd7a 100644 --- a/codex-rs/model-provider-info/src/lib.rs +++ b/codex-rs/model-provider-info/src/lib.rs @@ -20,11 +20,13 @@ use serde::Deserialize; use serde::Serialize; use std::collections::HashMap; use std::fmt; +use std::num::NonZeroU64; use std::time::Duration; const DEFAULT_STREAM_IDLE_TIMEOUT_MS: u64 = 300_000; const DEFAULT_STREAM_MAX_RETRIES: u64 = 5; const DEFAULT_REQUEST_MAX_RETRIES: u64 = 4; +const DEFAULT_AWS_AUTH_REFRESH_TIMEOUT_MS: u64 = 300_000; pub const DEFAULT_WEBSOCKET_CONNECT_TIMEOUT_MS: u64 = 15_000; /// Hard cap for user-configured `stream_max_retries`. const MAX_STREAM_MAX_RETRIES: u64 = 100; @@ -155,6 +157,35 @@ pub struct ModelProviderAwsAuthInfo { pub profile: Option, /// AWS region to use for provider-specific endpoints. pub region: Option, + /// Optional command used to reauthenticate after a refreshable AWS auth failure. + pub auth_refresh: Option, +} + +/// Command used to refresh AWS credentials for a model provider. +#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq, JsonSchema)] +#[schemars(deny_unknown_fields)] +pub struct AwsAuthRefreshConfig { + /// Executable to invoke directly, without a shell. + pub command: String, + /// Arguments passed to the refresh command. + #[serde(default)] + pub args: Vec, + /// Maximum time to wait for the refresh command to complete. + #[serde(default = "default_aws_auth_refresh_timeout_ms")] + pub timeout_ms: NonZeroU64, +} + +impl AwsAuthRefreshConfig { + pub fn timeout(&self) -> Duration { + Duration::from_millis(self.timeout_ms.get()) + } +} + +fn default_aws_auth_refresh_timeout_ms() -> NonZeroU64 { + match NonZeroU64::new(DEFAULT_AWS_AUTH_REFRESH_TIMEOUT_MS) { + Some(timeout_ms) => timeout_ms, + None => panic!("AWS auth refresh timeout must be non-zero"), + } } impl ModelProviderInfo { @@ -187,6 +218,16 @@ impl ModelProviderInfo { conflicts.join(", ") )); } + + if let Some(auth_refresh) = self.aws.as_ref().and_then(|aws| aws.auth_refresh.as_ref()) + { + if auth_refresh.command.trim().is_empty() { + return Err("provider aws.auth_refresh.command must not be empty".to_string()); + } + if auth_refresh.command != "aws" { + return Err("provider aws.auth_refresh.command must be `aws`".to_string()); + } + } } let Some(auth) = self.auth.as_ref() else { @@ -387,6 +428,7 @@ impl ModelProviderInfo { aws: Some(aws.unwrap_or(ModelProviderAwsAuthInfo { profile: None, region: None, + auth_refresh: None, })), wire_api: WireApi::Responses, query_params: None, @@ -504,8 +546,8 @@ pub fn merge_configured_model_providers( if provider != ModelProviderInfo::default() { return Err(format!( "model_providers.{key} only supports changing \ -`base_url`, `auth`, `http_headers`, `aws.profile`, and `aws.region`; other non-default \ -provider fields are not supported" +`base_url`, `auth`, `http_headers`, `aws.profile`, `aws.region`, and `aws.auth_refresh`; \ +other non-default provider fields are not supported" )); } diff --git a/codex-rs/model-provider-info/src/model_provider_info_tests.rs b/codex-rs/model-provider-info/src/model_provider_info_tests.rs index 26950e41ac..cb8dd5cfc9 100644 --- a/codex-rs/model-provider-info/src/model_provider_info_tests.rs +++ b/codex-rs/model-provider-info/src/model_provider_info_tests.rs @@ -216,6 +216,10 @@ base_url = "https://bedrock.example.com/v1" [aws] profile = "codex-bedrock" region = "us-west-2" + +[aws.auth_refresh] +command = "aws" +args = ["login", "--profile", "codex-bedrock"] "#; let provider: ModelProviderInfo = toml::from_str(provider_toml).unwrap(); @@ -225,6 +229,15 @@ region = "us-west-2" Some(ModelProviderAwsAuthInfo { profile: Some("codex-bedrock".to_string()), region: Some("us-west-2".to_string()), + auth_refresh: Some(AwsAuthRefreshConfig { + command: "aws".to_string(), + args: vec![ + "login".to_string(), + "--profile".to_string(), + "codex-bedrock".to_string(), + ], + timeout_ms: NonZeroU64::new(300_000).expect("timeout should be non-zero"), + }), }) ); } @@ -243,6 +256,7 @@ fn test_create_amazon_bedrock_provider() { aws: Some(ModelProviderAwsAuthInfo { profile: None, region: None, + auth_refresh: None, }), wire_api: WireApi::Responses, query_params: None, @@ -280,6 +294,7 @@ fn test_create_amazon_bedrock_runtime_provider_with_aws_configuration() { ModelProviderInfo::create_amazon_bedrock_runtime_provider(Some(ModelProviderAwsAuthInfo { profile: Some("runtime-profile".to_string()), region: Some("us-west-2".to_string()), + auth_refresh: None, })); assert_eq!( @@ -294,6 +309,7 @@ fn test_create_amazon_bedrock_runtime_provider_with_aws_configuration() { Some(ModelProviderAwsAuthInfo { profile: Some("runtime-profile".to_string()), region: Some("us-west-2".to_string()), + auth_refresh: None, }), None, false, @@ -389,13 +405,23 @@ fn test_merge_configured_model_providers_adds_custom_provider() { } #[test] -fn test_merge_configured_model_providers_applies_amazon_bedrock_profile_override() { +fn test_merge_configured_model_providers_applies_amazon_bedrock_aws_override() { + let auth_refresh = AwsAuthRefreshConfig { + command: "aws".to_string(), + args: vec![ + "login".to_string(), + "--profile".to_string(), + "codex-bedrock".to_string(), + ], + timeout_ms: NonZeroU64::new(10_000).expect("timeout should be non-zero"), + }; let configured_model_providers = std::collections::HashMap::from([( AMAZON_BEDROCK_PROVIDER_ID.to_string(), ModelProviderInfo { aws: Some(ModelProviderAwsAuthInfo { profile: Some("codex-bedrock".to_string()), region: Some("us-west-2".to_string()), + auth_refresh: Some(auth_refresh.clone()), }), ..ModelProviderInfo::default() }, @@ -408,6 +434,7 @@ fn test_merge_configured_model_providers_applies_amazon_bedrock_profile_override .aws = Some(ModelProviderAwsAuthInfo { profile: Some("codex-bedrock".to_string()), region: Some("us-west-2".to_string()), + auth_refresh: Some(auth_refresh), }); assert_eq!( @@ -424,6 +451,7 @@ fn test_merge_configured_model_providers_applies_runtime_overrides_independently let runtime_aws = ModelProviderAwsAuthInfo { profile: Some("runtime-profile".to_string()), region: Some("eu-west-1".to_string()), + auth_refresh: None, }; let configured_model_providers = std::collections::HashMap::from([( AMAZON_BEDROCK_RUNTIME_PROVIDER_ID.to_string(), @@ -460,6 +488,7 @@ fn test_merge_configured_model_providers_applies_amazon_bedrock_transport_overri aws: Some(ModelProviderAwsAuthInfo { profile: Some("codex-bedrock".to_string()), region: Some("us-west-2".to_string()), + auth_refresh: None, }), http_headers: Some(maplit::hashmap! { "x-example-header".to_string() => "value".to_string(), @@ -477,6 +506,7 @@ fn test_merge_configured_model_providers_applies_amazon_bedrock_transport_overri expected_provider.aws = Some(ModelProviderAwsAuthInfo { profile: Some("codex-bedrock".to_string()), region: Some("us-west-2".to_string()), + auth_refresh: None, }); expected_provider .http_headers @@ -501,6 +531,7 @@ fn test_merge_configured_model_providers_rejects_amazon_bedrock_non_default_fiel aws: Some(ModelProviderAwsAuthInfo { profile: Some("codex-bedrock".to_string()), region: None, + auth_refresh: None, }), ..ModelProviderInfo::default() }, @@ -512,7 +543,7 @@ fn test_merge_configured_model_providers_rejects_amazon_bedrock_non_default_fiel configured_model_providers, ), Err( - "model_providers.amazon-bedrock only supports changing `base_url`, `auth`, `http_headers`, `aws.profile`, and `aws.region`; other non-default provider fields are not supported" + "model_providers.amazon-bedrock only supports changing `base_url`, `auth`, `http_headers`, `aws.profile`, `aws.region`, and `aws.auth_refresh`; other non-default provider fields are not supported" .to_string() ) ); @@ -526,6 +557,7 @@ fn test_merge_configured_model_providers_allows_amazon_bedrock_default_fields() aws: Some(ModelProviderAwsAuthInfo { profile: None, region: None, + auth_refresh: None, }), wire_api: WireApi::Responses, ..ModelProviderInfo::default() @@ -547,6 +579,7 @@ fn test_validate_provider_aws_rejects_conflicting_auth() { aws: Some(ModelProviderAwsAuthInfo { profile: None, region: None, + auth_refresh: None, }), env_key: Some("AWS_BEARER_TOKEN_BEDROCK".to_string()), supports_websockets: false, @@ -565,6 +598,7 @@ fn test_validate_provider_aws_rejects_websockets() { aws: Some(ModelProviderAwsAuthInfo { profile: None, region: None, + auth_refresh: None, }), requires_openai_auth: false, supports_websockets: true, @@ -577,6 +611,34 @@ fn test_validate_provider_aws_rejects_websockets() { ); } +#[test] +fn test_validate_provider_aws_auth_refresh_command() { + for (command, expected) in [ + ( + " ", + Err("provider aws.auth_refresh.command must not be empty".to_string()), + ), + ( + "other-command", + Err("provider aws.auth_refresh.command must be `aws`".to_string()), + ), + ("aws", Ok(())), + ] { + let provider = + ModelProviderInfo::create_amazon_bedrock_provider(Some(ModelProviderAwsAuthInfo { + profile: None, + region: None, + auth_refresh: Some(AwsAuthRefreshConfig { + command: command.to_string(), + args: Vec::new(), + timeout_ms: NonZeroU64::new(300_000).expect("timeout should be non-zero"), + }), + })); + + assert_eq!(provider.validate(), expected); + } +} + #[test] fn test_deserialize_provider_auth_config_allows_zero_refresh_interval() { let base_dir = tempdir().unwrap(); diff --git a/codex-rs/model-provider/Cargo.toml b/codex-rs/model-provider/Cargo.toml index b6ee6acbfc..bbb665a19e 100644 --- a/codex-rs/model-provider/Cargo.toml +++ b/codex-rs/model-provider/Cargo.toml @@ -25,7 +25,7 @@ codex-otel = { workspace = true } codex-protocol = { workspace = true } codex-response-debug-context = { workspace = true } http = { workspace = true } -tokio = { workspace = true, features = ["sync", "time"] } +tokio = { workspace = true, features = ["process", "sync", "time"] } tracing = { workspace = true, features = ["log"] } [dev-dependencies] diff --git a/codex-rs/model-provider/src/amazon_bedrock/auth.rs b/codex-rs/model-provider/src/amazon_bedrock/auth.rs index 1420f14654..c082434002 100644 --- a/codex-rs/model-provider/src/amazon_bedrock/auth.rs +++ b/codex-rs/model-provider/src/amazon_bedrock/auth.rs @@ -9,8 +9,11 @@ use codex_aws_auth::AwsRequestToSign; use codex_http_client::Request; use codex_http_client::RequestBody; use codex_http_client::RequestCompression; +use codex_login::AuthManager; +use codex_login::CodexAuth; use codex_login::auth::BedrockApiKeyAuth; use codex_model_provider_info::ModelProviderAwsAuthInfo; +use codex_model_provider_info::ModelProviderInfo; use codex_protocol::error::CodexErr; use codex_protocol::error::Result; use http::HeaderMap; @@ -22,16 +25,50 @@ use super::mantle::aws_auth_config; use super::mantle::region_from_config; use super::runtime; -const AWS_BEARER_TOKEN_BEDROCK_ENV_VAR: &str = "AWS_BEARER_TOKEN_BEDROCK"; +pub(super) const AWS_BEARER_TOKEN_BEDROCK_ENV_VAR: &str = "AWS_BEARER_TOKEN_BEDROCK"; +const AWS_ACCESS_KEY_ID_ENV_VAR: &str = "AWS_ACCESS_KEY_ID"; +const AWS_SECRET_ACCESS_KEY_ENV_VAR: &str = "AWS_SECRET_ACCESS_KEY"; const AWS_REGION_ENV_VAR: &str = "AWS_REGION"; const AWS_DEFAULT_REGION_ENV_VAR: &str = "AWS_DEFAULT_REGION"; +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(super) enum BedrockAuthSource { + CommandBearerToken, + ManagedBearerToken, + EnvBearerToken, + EnvAwsCredentials, + AwsSdk, +} + pub(super) enum BedrockAuthMethod { ManagedBearerToken { token: String, region: String }, EnvBearerToken { token: String, region: String }, AwsSdkAuth { context: AwsAuthContext }, } +pub(super) fn auth_source( + provider_info: &ModelProviderInfo, + auth_manager: Option<&AuthManager>, + env_var: impl Fn(&'static str) -> std::result::Result + Copy, +) -> BedrockAuthSource { + if provider_info.has_command_auth() { + BedrockAuthSource::CommandBearerToken + } else if matches!( + auth_manager.and_then(AuthManager::auth_cached), + Some(CodexAuth::BedrockApiKey(_)) + ) { + BedrockAuthSource::ManagedBearerToken + } else if non_empty_env_var_from(AWS_BEARER_TOKEN_BEDROCK_ENV_VAR, env_var).is_some() { + BedrockAuthSource::EnvBearerToken + } else if non_empty_env_var_from(AWS_ACCESS_KEY_ID_ENV_VAR, env_var).is_some() + && non_empty_env_var_from(AWS_SECRET_ACCESS_KEY_ENV_VAR, env_var).is_some() + { + BedrockAuthSource::EnvAwsCredentials + } else { + BedrockAuthSource::AwsSdk + } +} + pub(super) async fn resolve_auth_method( managed_auth: Option<&BedrockApiKeyAuth>, aws: &ModelProviderAwsAuthInfo, @@ -199,6 +236,32 @@ mod tests { Err(std::env::VarError::NotPresent) } + #[test] + fn bedrock_auth_source_distinguishes_static_environment_credentials() { + let provider = ModelProviderInfo::create_amazon_bedrock_provider(/*aws*/ None); + let cases: &[(&[&str], BedrockAuthSource)] = &[ + ( + &[AWS_ACCESS_KEY_ID_ENV_VAR, AWS_SECRET_ACCESS_KEY_ENV_VAR], + BedrockAuthSource::EnvAwsCredentials, + ), + (&[AWS_ACCESS_KEY_ID_ENV_VAR], BedrockAuthSource::AwsSdk), + ( + &["AWS_PROFILE", AWS_REGION_ENV_VAR], + BedrockAuthSource::AwsSdk, + ), + ]; + + for (variables, expected) in cases { + let actual = auth_source(&provider, /*auth_manager*/ None, |name| { + variables + .contains(&name) + .then(|| "configured".to_string()) + .ok_or(std::env::VarError::NotPresent) + }); + assert_eq!(actual, *expected, "{variables:?}"); + } + } + #[test] fn bedrock_bearer_auth_prefers_configured_region_and_uses_header() { let token = "bedrock-api-key-test".to_string(); @@ -206,6 +269,7 @@ mod tests { &ModelProviderAwsAuthInfo { profile: None, region: Some(" us-west-2 ".to_string()), + auth_refresh: None, }, |name| match name { AWS_REGION_ENV_VAR => Ok("eu-west-1".to_string()), @@ -237,6 +301,7 @@ mod tests { &ModelProviderAwsAuthInfo { profile: None, region: None, + auth_refresh: None, }, |name| match name { AWS_REGION_ENV_VAR => Ok(" eu-central-1 ".to_string()), @@ -254,6 +319,7 @@ mod tests { &ModelProviderAwsAuthInfo { profile: None, region: None, + auth_refresh: None, }, |name| match name { AWS_DEFAULT_REGION_ENV_VAR => Ok("ap-northeast-1".to_string()), @@ -271,6 +337,7 @@ mod tests { &ModelProviderAwsAuthInfo { profile: None, region: None, + auth_refresh: None, }, missing_env_var, ) diff --git a/codex-rs/model-provider/src/amazon_bedrock/auth_refresh.rs b/codex-rs/model-provider/src/amazon_bedrock/auth_refresh.rs new file mode 100644 index 0000000000..213be9f47d --- /dev/null +++ b/codex-rs/model-provider/src/amazon_bedrock/auth_refresh.rs @@ -0,0 +1,82 @@ +use std::io::IsTerminal; +use std::process::Stdio; +use std::sync::atomic::AtomicU64; +use std::sync::atomic::Ordering; + +use codex_model_provider_info::AwsAuthRefreshConfig; +use tokio::process::Command; +use tokio::sync::Semaphore; + +/// Shares the configured AWS reauthentication command across Bedrock providers. +#[derive(Debug)] +pub(crate) struct AwsAuthRecovery { + config: AwsAuthRefreshConfig, + refresh_lock: Semaphore, + generation: AtomicU64, +} + +impl AwsAuthRecovery { + pub(crate) fn new(config: AwsAuthRefreshConfig) -> Self { + Self { + config, + refresh_lock: Semaphore::new(/*permits*/ 1), + generation: AtomicU64::new(0), + } + } + + pub(crate) async fn refresh(&self) -> std::io::Result<()> { + if self.config.command != "aws" { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "AWS auth refresh command must be `aws`", + )); + } + + let generation = self.generation.load(Ordering::Acquire); + let _refresh_lock = self + .refresh_lock + .acquire() + .await + .map_err(std::io::Error::other)?; + if self.generation.load(Ordering::Acquire) != generation { + return Ok(()); + } + + let mut command = Command::new(&self.config.command); + command + .args(&self.config.args) + .stdin(if std::io::stdin().is_terminal() { + Stdio::inherit() + } else { + Stdio::null() + }) + // Keep interactive prompts visible without corrupting app-server's JSON-RPC stdout. + .stdout(std::io::stderr()) + .stderr(Stdio::inherit()) + .kill_on_drop(true); + + let status = tokio::time::timeout(self.config.timeout(), command.status()) + .await + .map_err(|_| { + std::io::Error::other(format!( + "AWS auth refresh command timed out after {} ms", + self.config.timeout_ms + )) + })? + .map_err(|error| { + std::io::Error::other(format!( + "failed to run AWS auth refresh command `{}`: {error}", + self.config.command + )) + })?; + + if status.success() { + self.generation.fetch_add(1, Ordering::Release); + Ok(()) + } else { + Err(std::io::Error::other(format!( + "AWS auth refresh command exited with {status}" + ))) + } + } +} diff --git a/codex-rs/model-provider/src/amazon_bedrock/error.rs b/codex-rs/model-provider/src/amazon_bedrock/error.rs index 51428be7f5..ad67f1ff89 100644 --- a/codex-rs/model-provider/src/amazon_bedrock/error.rs +++ b/codex-rs/model-provider/src/amazon_bedrock/error.rs @@ -1,4 +1,5 @@ use codex_api::ApiError; +use codex_api::TransportError; use codex_protocol::error::CodexErr; use codex_protocol::error::CodexErrorDetails; use http::StatusCode; @@ -25,3 +26,27 @@ pub(super) fn map_api_error(error: ApiError) -> CodexErr { } error } + +pub(super) fn is_refreshable_auth_error(error: &TransportError) -> bool { + match error { + TransportError::Build(message) | TransportError::Network(message) => { + message.starts_with("failed to load AWS credentials:") + } + TransportError::Http { status, .. } if *status == StatusCode::UNAUTHORIZED => true, + TransportError::Http { + status, + body: Some(body), + .. + } if *status == StatusCode::FORBIDDEN => { + let body = body.to_ascii_lowercase(); + [ + "expiredtoken", + "unrecognizedclientexception", + "invalidclienttokenid", + ] + .iter() + .any(|error_code| body.contains(error_code)) + } + _ => false, + } +} diff --git a/codex-rs/model-provider/src/amazon_bedrock/error_tests.rs b/codex-rs/model-provider/src/amazon_bedrock/error_tests.rs index 053697907a..692be488e4 100644 --- a/codex-rs/model-provider/src/amazon_bedrock/error_tests.rs +++ b/codex-rs/model-provider/src/amazon_bedrock/error_tests.rs @@ -7,6 +7,7 @@ use http::StatusCode; use pretty_assertions::assert_eq; use super::error::BEDROCK_EXPIRED_SIGNATURE_MESSAGE; +use super::error::is_refreshable_auth_error; use super::error::map_api_error; const BEDROCK_RESPONSES_URL: &str = "https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses"; @@ -75,3 +76,71 @@ fn signature_errors_with_other_statuses_remain_generic() { }; assert_eq!(response.user_message, None); } + +#[test] +fn classifies_only_refreshable_bedrock_auth_failures() { + let cases = [ + ( + StatusCode::UNAUTHORIZED, + Some("ExpiredTokenException"), + true, + ), + ( + StatusCode::UNAUTHORIZED, + Some("AccessDeniedException"), + true, + ), + (StatusCode::UNAUTHORIZED, Some(""), true), + (StatusCode::UNAUTHORIZED, None, true), + (StatusCode::FORBIDDEN, Some("ExpiredTokenException"), true), + ( + StatusCode::FORBIDDEN, + Some("UnrecognizedClientException"), + true, + ), + (StatusCode::FORBIDDEN, Some("InvalidClientTokenId"), true), + (StatusCode::FORBIDDEN, Some("AccessDeniedException"), false), + ( + StatusCode::FORBIDDEN, + Some("The security token included in the request is invalid"), + false, + ), + ( + StatusCode::TOO_MANY_REQUESTS, + Some("ExpiredTokenException"), + false, + ), + (StatusCode::BAD_REQUEST, Some("RequestExpired"), false), + ]; + + for (status, body, expected) in cases { + let error = TransportError::Http { + status, + url: Some(BEDROCK_RESPONSES_URL.to_string()), + headers: None, + body: body.map(str::to_string), + }; + assert_eq!( + is_refreshable_auth_error(&error), + expected, + "{status}: {body:?}" + ); + } + + for (error, expected) in [ + ( + TransportError::Build("failed to load AWS credentials: expired".to_string()), + true, + ), + ( + TransportError::Network("failed to load AWS credentials: unavailable".to_string()), + true, + ), + ( + TransportError::Build("request URL is not a valid URI".to_string()), + false, + ), + ] { + assert_eq!(is_refreshable_auth_error(&error), expected, "{error}"); + } +} diff --git a/codex-rs/model-provider/src/amazon_bedrock/mantle.rs b/codex-rs/model-provider/src/amazon_bedrock/mantle.rs index a75f124cb8..00f4d46476 100644 --- a/codex-rs/model-provider/src/amazon_bedrock/mantle.rs +++ b/codex-rs/model-provider/src/amazon_bedrock/mantle.rs @@ -92,6 +92,7 @@ mod tests { aws_auth_config(&ModelProviderAwsAuthInfo { profile: Some("codex-bedrock".to_string()), region: None, + auth_refresh: None, }), AwsAuthConfig { profile: Some("codex-bedrock".to_string()), @@ -107,6 +108,7 @@ mod tests { aws_auth_config(&ModelProviderAwsAuthInfo { profile: None, region: Some(" us-west-2 ".to_string()), + auth_refresh: None, }), AwsAuthConfig { profile: None, diff --git a/codex-rs/model-provider/src/amazon_bedrock/mod.rs b/codex-rs/model-provider/src/amazon_bedrock/mod.rs index f7707d0ab8..e2af6b3350 100644 --- a/codex-rs/model-provider/src/amazon_bedrock/mod.rs +++ b/codex-rs/model-provider/src/amazon_bedrock/mod.rs @@ -1,4 +1,5 @@ mod auth; +mod auth_refresh; mod catalog; mod error; mod mantle; @@ -11,6 +12,7 @@ use std::sync::Arc; use codex_api::ApiError; use codex_api::Provider; use codex_api::SharedAuthProvider; +use codex_api::TransportError; use codex_login::AuthManager; use codex_login::CodexAuth; use codex_login::auth::BedrockApiKeyAuth; @@ -34,8 +36,11 @@ use crate::provider::ModelProviderFuture; use crate::provider::ProviderAccountResult; use crate::provider::ProviderAccountState; use crate::provider::ProviderCapabilities; +use crate::provider::ProviderUnauthorizedRecovery; use crate::provider::RemoteCompactionSupport; +use crate::shared_state::process_shared_state; use auth::resolve_provider_auth as resolve_bedrock_provider_auth; +pub(crate) use auth_refresh::AwsAuthRecovery; use catalog::normalize_bedrock_catalog; pub(crate) use catalog::static_model_catalog; use mantle::bedrock_mantle_runtime_base_url; @@ -53,9 +58,10 @@ pub(super) enum BedrockEndpoint { #[derive(Clone, Debug)] pub(crate) struct AmazonBedrockModelProvider { pub(crate) info: ModelProviderInfo, - pub(crate) aws: ModelProviderAwsAuthInfo, + aws: ModelProviderAwsAuthInfo, endpoint: BedrockEndpoint, auth_manager: Option>, + auth_recovery: Option>, } impl AmazonBedrockModelProvider { @@ -63,7 +69,6 @@ impl AmazonBedrockModelProvider { provider_info: ModelProviderInfo, auth_manager: Option>, ) -> Self { - let auth_manager = auth_manager_for_provider(auth_manager, &provider_info); let endpoint = if provider_info.is_amazon_bedrock_runtime() { BedrockEndpoint::Runtime } else { @@ -75,12 +80,23 @@ impl AmazonBedrockModelProvider { .unwrap_or(ModelProviderAwsAuthInfo { profile: None, region: None, + auth_refresh: None, }); + let uses_aws_sdk_auth = + auth::auth_source(&provider_info, auth_manager.as_deref(), std::env::var) + == auth::BedrockAuthSource::AwsSdk; + let auth_recovery = if uses_aws_sdk_auth && aws.auth_refresh.is_some() { + process_shared_state().aws_auth_recovery(&aws) + } else { + None + }; + let auth_manager = auth_manager_for_provider(auth_manager, &provider_info); Self { info: provider_info, aws, endpoint, auth_manager, + auth_recovery, } } @@ -99,6 +115,12 @@ impl AmazonBedrockModelProvider { }) } + fn uses_aws_auth_recovery(&self) -> bool { + self.auth_recovery.is_some() + && auth::auth_source(&self.info, self.auth_manager.as_deref(), std::env::var) + == auth::BedrockAuthSource::AwsSdk + } + async fn auth(&self) -> Option { if self.info.has_command_auth() { match self.auth_manager.as_ref() { @@ -193,6 +215,36 @@ impl ModelProvider for AmazonBedrockModelProvider { } } + fn is_recoverable_auth_error(&self, error: &TransportError) -> bool { + matches!( + error, + TransportError::Http { status, .. } if *status == http::StatusCode::UNAUTHORIZED + ) || (self.uses_aws_auth_recovery() && error::is_refreshable_auth_error(error)) + } + + fn recover_from_unauthorized( + &self, + ) -> ModelProviderFuture<'_, Result> { + Box::pin(async move { + let Some(recovery) = self + .auth_recovery + .as_ref() + .filter(|_| self.uses_aws_auth_recovery()) + else { + return Ok(ProviderUnauthorizedRecovery::NotConfigured); + }; + + recovery.refresh().await.map_err(|error| { + if error.kind() == std::io::ErrorKind::InvalidInput { + CodexErr::InvalidRequest(error.to_string()) + } else { + CodexErr::Io(error) + } + })?; + Ok(ProviderUnauthorizedRecovery::Recovered) + }) + } + fn auth(&self) -> ModelProviderFuture<'_, Option> { Box::pin(AmazonBedrockModelProvider::auth(self)) } @@ -254,6 +306,7 @@ mod error_tests; mod tests { use std::num::NonZeroU64; + use codex_model_provider_info::AwsAuthRefreshConfig; use codex_protocol::config_types::ModelProviderAuthInfo; use http::HeaderValue; use pretty_assertions::assert_eq; @@ -298,6 +351,11 @@ mod tests { provider_info.aws = Some(ModelProviderAwsAuthInfo { profile: Some("aws-profile-that-should-not-be-loaded".to_string()), region: Some("us-west-2".to_string()), + auth_refresh: Some(AwsAuthRefreshConfig { + command: "aws".to_string(), + args: vec!["login".to_string()], + timeout_ms: NonZeroU64::new(1_000).expect("timeout should be non-zero"), + }), }); let provider = AmazonBedrockModelProvider::new(provider_info, /*auth_manager*/ None); @@ -337,6 +395,11 @@ mod tests { ModelProviderInfo::create_amazon_bedrock_provider(Some(ModelProviderAwsAuthInfo { profile: Some("aws-profile-that-should-not-be-loaded".to_string()), region: Some("us-west-2".to_string()), + auth_refresh: Some(AwsAuthRefreshConfig { + command: "aws".to_string(), + args: vec!["login".to_string()], + timeout_ms: NonZeroU64::new(1_000).expect("timeout should be non-zero"), + }), })), Some(auth_manager.clone()), ); diff --git a/codex-rs/model-provider/src/amazon_bedrock/runtime_tests.rs b/codex-rs/model-provider/src/amazon_bedrock/runtime_tests.rs index 6ba3439481..84f278bed8 100644 --- a/codex-rs/model-provider/src/amazon_bedrock/runtime_tests.rs +++ b/codex-rs/model-provider/src/amazon_bedrock/runtime_tests.rs @@ -19,6 +19,7 @@ fn aws_auth_config_uses_bedrock_service() { aws_auth_config(&ModelProviderAwsAuthInfo { profile: Some("runtime-profile".to_string()), region: Some(" eu-west-1 ".to_string()), + auth_refresh: None, }), AwsAuthConfig { profile: Some("runtime-profile".to_string()), diff --git a/codex-rs/model-provider/src/lib.rs b/codex-rs/model-provider/src/lib.rs index 22b5da74ed..b7d72a6f17 100644 --- a/codex-rs/model-provider/src/lib.rs +++ b/codex-rs/model-provider/src/lib.rs @@ -3,6 +3,7 @@ mod auth; mod bearer_auth_provider; mod models_endpoint; mod provider; +mod shared_state; pub use amazon_bedrock::is_supported_amazon_bedrock_region; pub use auth::AgentIdentitySessionFallback; diff --git a/codex-rs/model-provider/src/provider.rs b/codex-rs/model-provider/src/provider.rs index f4fc6b6688..65f4204d56 100644 --- a/codex-rs/model-provider/src/provider.rs +++ b/codex-rs/model-provider/src/provider.rs @@ -494,6 +494,7 @@ mod tests { use codex_http_client::OutboundProxyPolicy; use codex_login::auth::AgentIdentityAuthPolicy; use codex_login::auth::BedrockApiKeyAuth; + use codex_model_provider_info::AwsAuthRefreshConfig; use codex_model_provider_info::ModelProviderAwsAuthInfo; use codex_model_provider_info::WireApi; use codex_model_provider_info::create_oss_provider_with_base_url; @@ -515,6 +516,7 @@ mod tests { use super::*; use crate::auth::AgentIdentitySessionFallback; + use crate::shared_state::process_shared_state; fn provider_info_with_command_auth() -> ModelProviderInfo { ModelProviderInfo { @@ -731,6 +733,7 @@ mod tests { ModelProviderInfo::create_amazon_bedrock_provider(Some(ModelProviderAwsAuthInfo { profile: Some("codex-bedrock".to_string()), region: None, + auth_refresh: None, })), Some(AuthManager::from_auth_for_testing(CodexAuth::from_api_key( "openai-api-key", @@ -740,6 +743,150 @@ mod tests { assert!(provider.auth_manager().is_none()); } + #[tokio::test] + async fn shared_bedrock_auth_refresh_is_reused_only_for_matching_configuration() { + const TEST_NAME: &str = "provider::tests::shared_bedrock_auth_refresh_is_reused_only_for_matching_configuration"; + const HELPER_ARG: &str = "CODEX_BEDROCK_SHARED_AUTH_REFRESH_COMMAND"; + const SUBPROCESS_ARG: &str = "CODEX_BEDROCK_SHARED_AUTH_REFRESH_SUBPROCESS"; + let arguments = std::env::args().collect::>(); + if let Some(index) = arguments.iter().position(|argument| argument == HELPER_ARG) { + let counter = &arguments[index + 2]; + std::thread::sleep(std::time::Duration::from_millis(50)); + let mut counter = std::fs::OpenOptions::new() + .create(true) + .append(true) + .open(counter) + .expect("refresh invocation counter should open"); + std::io::Write::write_all(&mut counter, b"1").expect("write invocation counter"); + return; + } + + let counter = std::env::temp_dir().join(format!("bedrock-refresh-{}", std::process::id())); + if !arguments.iter().any(|argument| argument == SUBPROCESS_ARG) { + std::fs::create_dir(&counter).expect("AWS command directory should be created"); + let executable = std::env::current_exe().expect("test executable should be available"); + let aws = counter.join(format!("aws{}", std::env::consts::EXE_SUFFIX)); + std::fs::hard_link(&executable, &aws) + .or_else(|_| std::fs::copy(&executable, &aws).map(|_| ())) + .expect("test executable should be installed as aws"); + let existing_path = std::env::var_os("PATH").unwrap_or_default(); + let path = std::env::join_paths( + std::iter::once(counter.clone()).chain(std::env::split_paths(&existing_path)), + ) + .expect("test executable PATH should be valid"); + let output = tokio::process::Command::new(executable) + .args(["--exact", TEST_NAME, "--skip", SUBPROCESS_ARG]) + .env("PATH", path) + .env_remove("AWS_ACCESS_KEY_ID") + .env_remove("AWS_SECRET_ACCESS_KEY") + .output() + .await + .expect("isolated AWS refresh test should run"); + std::fs::remove_dir_all(&counter).expect("AWS command directory should be removed"); + assert!(output.status.success(), "{output:?}"); + return; + } + let _ = std::fs::remove_file(&counter); + let aws = ModelProviderAwsAuthInfo { + profile: Some("codex-bedrock".to_string()), + region: Some("us-west-2".to_string()), + auth_refresh: Some(AwsAuthRefreshConfig { + command: "aws".to_string(), + args: Vec::from( + [ + "--exact", + TEST_NAME, + "--skip", + HELPER_ARG, + "--skip", + counter.to_str().expect("counter path should be UTF-8"), + ] + .map(str::to_string), + ), + timeout_ms: NonZeroU64::new(10_000).expect("timeout should be non-zero"), + }), + }; + let provider_info = ModelProviderInfo::create_amazon_bedrock_provider(Some(aws.clone())); + + let first = create_model_provider( + provider_info.clone(), + Some(AuthManager::from_auth_for_testing(CodexAuth::from_api_key( + "openai-api-key", + ))), + ); + let second = create_model_provider(provider_info.clone(), /*auth_manager*/ None); + let shared_state = process_shared_state(); + let shared_recovery = shared_state + .aws_auth_recovery(&aws) + .expect("test provider should have refresh configured"); + assert_eq!(Arc::strong_count(&shared_recovery), 3); + assert!(first.auth_manager().is_none() && second.auth_manager().is_none()); + + for provider in [&first, &second] { + for (status, body, recoverable) in [ + (http::StatusCode::UNAUTHORIZED, "ExpiredToken", true), + (http::StatusCode::FORBIDDEN, "InvalidClientTokenId", true), + (http::StatusCode::FORBIDDEN, "AccessDeniedException", false), + ] { + let error = TransportError::Http { + status, + url: None, + headers: None, + body: Some(body.to_string()), + }; + assert_eq!(provider.is_recoverable_auth_error(&error), recoverable); + } + } + + let mut invalid_aws = aws.clone(); + invalid_aws.auth_refresh.as_mut().expect("refresh").command = "not-aws".into(); + let invalid_provider = create_model_provider( + ModelProviderInfo::create_amazon_bedrock_provider(Some(invalid_aws)), + /*auth_manager*/ None, + ); + let error = invalid_provider + .recover_from_unauthorized() + .await + .expect_err("non-aws command should be rejected"); + assert!(!error.is_retryable()); + assert_eq!(error.to_string(), "AWS auth refresh command must be `aws`"); + + let (first_result, second_result) = tokio::join!( + first.recover_from_unauthorized(), + second.recover_from_unauthorized() + ); + assert_eq!( + [first_result, second_result].map(|result| result.expect("provider should recover")), + [ProviderUnauthorizedRecovery::Recovered; 2] + ); + let read_counter = || std::fs::read_to_string(&counter).expect("read counter"); + assert_eq!(read_counter(), "1"); + + assert_eq!( + first + .recover_from_unauthorized() + .await + .expect("later generation should recover"), + ProviderUnauthorizedRecovery::Recovered + ); + assert_eq!(read_counter(), "11"); + std::fs::remove_file(&counter).expect("refresh invocation counter should be removed"); + + let different_profile = ModelProviderAwsAuthInfo { + profile: Some("another-bedrock-profile".to_string()), + ..aws.clone() + }; + let other_recovery = shared_state + .aws_auth_recovery(&different_profile) + .expect("test provider should have refresh configured"); + assert!(!Arc::ptr_eq(&shared_recovery, &other_recovery)); + + let released_recovery = Arc::downgrade(&shared_recovery); + drop((first, second, shared_recovery)); + assert!(released_recovery.upgrade().is_none()); + assert!(shared_state.aws_auth_recovery(&aws).is_some()); + } + #[tokio::test] async fn create_model_provider_uses_managed_auth_for_amazon_bedrock_provider() { let auth = bedrock_api_key_auth(); diff --git a/codex-rs/model-provider/src/shared_state.rs b/codex-rs/model-provider/src/shared_state.rs new file mode 100644 index 0000000000..13d348819c --- /dev/null +++ b/codex-rs/model-provider/src/shared_state.rs @@ -0,0 +1,45 @@ +use std::sync::Arc; +use std::sync::Mutex; +use std::sync::OnceLock; +use std::sync::Weak; + +use codex_model_provider_info::ModelProviderAwsAuthInfo; + +use crate::amazon_bedrock::AwsAuthRecovery; + +/// Provider-owned runtime state shared across independently configured sessions. +#[derive(Debug, Default)] +pub(crate) struct ModelProviderSharedState { + aws_auth_recoveries: Mutex)>>, +} + +pub(crate) fn process_shared_state() -> &'static ModelProviderSharedState { + static STATE: OnceLock = OnceLock::new(); + STATE.get_or_init(ModelProviderSharedState::default) +} + +impl ModelProviderSharedState { + pub(crate) fn aws_auth_recovery( + &self, + aws: &ModelProviderAwsAuthInfo, + ) -> Option> { + let config = aws.auth_refresh.as_ref()?; + let mut recoveries = self + .aws_auth_recoveries + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + recoveries.retain(|(_, recovery)| recovery.strong_count() != 0); + + if let Some(recovery) = recoveries + .iter() + .find(|(cached_aws, _)| cached_aws == aws) + .and_then(|(_, recovery)| recovery.upgrade()) + { + return Some(recovery); + } + + let recovery = Arc::new(AwsAuthRecovery::new(config.clone())); + recoveries.push((aws.clone(), Arc::downgrade(&recovery))); + Some(recovery) + } +} diff --git a/codex-rs/tui/src/status/tests.rs b/codex-rs/tui/src/status/tests.rs index ca28311421..161bbf4c74 100644 --- a/codex-rs/tui/src/status/tests.rs +++ b/codex-rs/tui/src/status/tests.rs @@ -738,6 +738,7 @@ async fn status_model_provider_uses_bedrock_runtime_base_url_and_gates_usage_lin ModelProviderInfo::create_amazon_bedrock_provider(Some(ModelProviderAwsAuthInfo { profile: None, region: Some("eu-west-1".to_string()), + auth_refresh: None, })); config.model_provider.base_url = Some("https://bedrock-mantle.us-east-1.api.aws/openai/v1".to_string());