openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-18 17:18:52 +00:00
Code Issues Packages Projects Releases Wiki Activity

Refresh expired AWS credentials for Bedrock (#39410)

Browse Source 
## Why

Bedrock sessions that use the AWS SDK credential chain need a way to recover when credentials expire during a request.

## What changed

- Add `aws.auth_refresh` provider configuration with an `aws` command, arguments, and a configurable timeout.
- Run the command for refreshable Bedrock authentication failures, reload the SDK credentials, re-sign the request, and retry it.
- Share refresh state across matching provider configurations so concurrent failures invoke the command only once. Bearer tokens, command auth, and static environment credentials do not use this recovery path.

## Testing

- Add coverage for configuration validation, refreshable error classification, concurrent refresh sharing, and an end-to-end retry signed with refreshed credentials.

GitOrigin-RevId: 0302fe3aabdbc1097e7bd62a74d407ba38a3cc57
This commit is contained in:
alexsong-oai
2026-08-19 04:46:58 +00:00
committed by copyberry
parent 6cc2ba8a95
commit 3929c99a97
18 changed files with 752 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/config/src/bedrock_runtime_tests.rs
+1
View File
@@ -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,
})
);
}
codex-rs/core/config.schema.json
+37
View File
@@ -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"
codex-rs/core/src/config/config_tests.rs
+1 -1
View File
@@ -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"
));
}
codex-rs/core/tests/suite/external_auth.rs
+98
View File
@@ -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::<Vec<_>>();
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(()));
codex-rs/model-provider-info/src/lib.rs
+44 -2
View File
@@ -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<String>,
/// AWS region to use for provider-specific endpoints.
pub region: Option<String>,
/// Optional command used to reauthenticate after a refreshable AWS auth failure.
pub auth_refresh: Option<AwsAuthRefreshConfig>,
}
/// 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<String>,
/// 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"
));
}
codex-rs/model-provider-info/src/model_provider_info_tests.rs
+64 -2
View File
@@ -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();
codex-rs/model-provider/Cargo.toml
+1 -1
View File
@@ -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]
codex-rs/model-provider/src/amazon_bedrock/auth.rs
+68 -1
View File
@@ -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<String, std::env::VarError> + 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,
)
codex-rs/model-provider/src/amazon_bedrock/auth_refresh.rs
+82
View File
@@ -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}"
)))
}
}
}
codex-rs/model-provider/src/amazon_bedrock/error.rs
+25
View File
@@ -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,
}
}
codex-rs/model-provider/src/amazon_bedrock/error_tests.rs
+69
View File
@@ -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}");
}
}
codex-rs/model-provider/src/amazon_bedrock/mantle.rs
+2
View File
@@ -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,
codex-rs/model-provider/src/amazon_bedrock/mod.rs
+65 -2
View File
@@ -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<Arc<AuthManager>>,
auth_recovery: Option<Arc<AwsAuthRecovery>>,
}
impl AmazonBedrockModelProvider {
@@ -63,7 +69,6 @@ impl AmazonBedrockModelProvider {
provider_info: ModelProviderInfo,
auth_manager: Option<Arc<AuthManager>>,
) -> 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<CodexAuth> {
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<ProviderUnauthorizedRecovery>> {
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<CodexAuth>> {
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()),
);
codex-rs/model-provider/src/amazon_bedrock/runtime_tests.rs
+1
View File
@@ -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()),
codex-rs/model-provider/src/lib.rs
+1
View File
@@ -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;
codex-rs/model-provider/src/provider.rs
+147
View File
@@ -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::<Vec<_>>();
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();
codex-rs/model-provider/src/shared_state.rs
+45
View File
@@ -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<Vec<(ModelProviderAwsAuthInfo, Weak<AwsAuthRecovery>)>>,
}
pub(crate) fn process_shared_state() -> &'static ModelProviderSharedState {
static STATE: OnceLock<ModelProviderSharedState> = OnceLock::new();
STATE.get_or_init(ModelProviderSharedState::default)
}
impl ModelProviderSharedState {
pub(crate) fn aws_auth_recovery(
&self,
aws: &ModelProviderAwsAuthInfo,
) -> Option<Arc<AwsAuthRecovery>> {
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)
}
}
codex-rs/tui/src/status/tests.rs
+1
View File
@@ -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());