From fe52d795c95321429c24f64be684856557491b47 Mon Sep 17 00:00:00 2001 From: alexsong-oai Date: Wed, 9 Sep 2026 02:56:10 +0000 Subject: [PATCH] Add AWS credential export commands for Amazon Bedrock (#44028) ## What changed Support `model_providers.amazon-bedrock.aws.credential_export` with `command`, `args`, and `timeout_ms` to supply SigV4 signing credentials from a command's JSON output. Accept both flat credential-process output and nested STS `Credentials` objects. Cache credentials in memory, refresh before expiration, and share exports across sessions with matching AWS configuration. On recoverable authentication failures, run the optional `aws.auth_refresh` command before exporting fresh credentials, coalescing concurrent recovery attempts. Bound command execution time and output size, and keep credential values out of errors. Reject combining `aws.credential_export` with `aws.profile`. Bedrock setup and login reject changes while an exporter is configured, preserving configuration and saved credentials. ## Testing Add coverage for output formats, expiration, credential precedence, caching, concurrent recovery, command failures, and secret redaction. Integration tests verify request signing with rotated credentials, bounded recovery attempts, and setup/login rejection across configuration layers. GitOrigin-RevId: 4ef799f88ab5e3b2578f8f41a5e56e4f30e2cd81 --- codex-rs/Cargo.lock | 3 + codex-rs/app-server/README.md | 7 + .../src/request_processors/bedrock_auth.rs | 13 + .../tests/suite/v2/bedrock_setup.rs | 82 +++- codex-rs/aws-auth/src/lib.rs | 132 ++++++ codex-rs/config/src/bedrock_runtime_tests.rs | 1 + codex-rs/core/config.schema.json | 37 ++ codex-rs/core/src/config/config_tests.rs | 2 +- codex-rs/core/tests/suite/external_auth.rs | 409 +++++++++++++++++- codex-rs/model-provider-info/src/lib.rs | 67 ++- .../src/model_provider_info_tests.rs | 108 ++++- codex-rs/model-provider/Cargo.toml | 7 +- .../model-provider/src/amazon_bedrock/auth.rs | 72 +++ .../src/amazon_bedrock/credential_export.rs | 250 +++++++++++ .../amazon_bedrock/credential_export_tests.rs | 348 +++++++++++++++ .../src/amazon_bedrock/error.rs | 8 + .../src/amazon_bedrock/error_tests.rs | 67 +++ .../src/amazon_bedrock/mantle.rs | 2 + .../model-provider/src/amazon_bedrock/mod.rs | 77 +++- .../src/amazon_bedrock/runtime_tests.rs | 1 + codex-rs/model-provider/src/provider.rs | 93 ++++ codex-rs/model-provider/src/shared_state.rs | 26 ++ codex-rs/tui/src/app/agents_overview_tests.rs | 17 +- codex-rs/tui/src/app/tests.rs | 43 +- .../rollout_history_tests.rs | 6 +- codex-rs/tui/src/status/tests.rs | 1 + 26 files changed, 1790 insertions(+), 89 deletions(-) create mode 100644 codex-rs/model-provider/src/amazon_bedrock/credential_export.rs create mode 100644 codex-rs/model-provider/src/amazon_bedrock/credential_export_tests.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 8eb7df52fa..9d53c1317b 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -3973,6 +3973,7 @@ name = "codex-model-provider" version = "0.0.0" dependencies = [ "base64 0.22.1", + "chrono", "codex-agent-identity", "codex-api", "codex-aws-auth", @@ -3987,8 +3988,10 @@ dependencies = [ "codex-utils-redacted-string", "http 1.4.0", "pretty_assertions", + "serde", "serde_json", "sha2 0.10.9", + "tempfile", "tokio", "tracing", "wiremock", diff --git a/codex-rs/app-server/README.md b/codex-rs/app-server/README.md index 80aede118c..e88e670154 100644 --- a/codex-rs/app-server/README.md +++ b/codex-rs/app-server/README.md @@ -31,3 +31,10 @@ after a client tries to archive or delete it. After the owner releases the worker, its saved conversation can be archived or deleted normally. Ordinary client-controlled threads keep their existing behavior. + +# Amazon Bedrock authentication + +If `model_providers.amazon-bedrock.aws.credential_export` is configured, Bedrock setup and +Bedrock login return an error without changing configuration or saved credentials. Remove the +exporter configuration before selecting another credential source. `aws.credential_export` and +`aws.profile` cannot be configured together. diff --git a/codex-rs/app-server/src/request_processors/bedrock_auth.rs b/codex-rs/app-server/src/request_processors/bedrock_auth.rs index 016b90bcf8..9aba8f9c52 100644 --- a/codex-rs/app-server/src/request_processors/bedrock_auth.rs +++ b/codex-rs/app-server/src/request_processors/bedrock_auth.rs @@ -99,6 +99,19 @@ pub(super) async fn ensure_user_model_provider_can_be_bedrock( ))); } + if layers + .effective_config() + .get("model_providers") + .and_then(|providers| providers.get(AMAZON_BEDROCK_PROVIDER_ID)) + .and_then(|provider| provider.get("aws")) + .and_then(|aws| aws.get("credential_export")) + .is_some() + { + return Err(invalid_request( + "Amazon Bedrock is configured to use `aws.credential_export`. Please clear this setting to use another sign-in method.", + )); + } + Ok(()) } diff --git a/codex-rs/app-server/tests/suite/v2/bedrock_setup.rs b/codex-rs/app-server/tests/suite/v2/bedrock_setup.rs index 87cfaf3ae4..21ce43474b 100644 --- a/codex-rs/app-server/tests/suite/v2/bedrock_setup.rs +++ b/codex-rs/app-server/tests/suite/v2/bedrock_setup.rs @@ -9,6 +9,7 @@ use codex_app_server_protocol::RequestId; use codex_config::types::AuthCredentialsStoreMode; use codex_login::AuthKeyringBackendKind; use codex_login::login_with_bedrock_api_key; +use codex_utils_absolute_path::AbsolutePathBuf; use pretty_assertions::assert_eq; use serde_json::json; use std::path::Path; @@ -122,11 +123,10 @@ async fn setup_bedrock_profile_and_environment() -> Result<()> { let codex_home = TempDir::new()?; let config_path = codex_home.path().join("config.toml"); let dotenv_path = codex_home.path().join(".env"); - std::fs::write( - &config_path, - "[model_providers.amazon-bedrock]\nhttp_headers = { X-Existing = \"preserved\" }\n\ - [model_providers.amazon-bedrock.aws]\nprofile = \"old\"\nregion = \"us-east-1\"\n", - )?; + let existing_config = "[model_providers.amazon-bedrock]\nhttp_headers = { X-Existing = \"preserved\" }\n\ + [model_providers.amazon-bedrock.aws]\nprofile = \"old\"\nregion = \"us-east-1\"\n\ + auth_refresh = { command = \"aws\", args = [\"login\"] }\n"; + std::fs::write(&config_path, existing_config)?; let existing_dotenv = "# existing configuration\nUNRELATED=value\nAWS_ACCESS_KEY_ID=old-id\n\ export AWS_SECRET_ACCESS_KEY=old-secret\nAWS_SESSION_TOKEN=stale-token\n"; std::fs::write(&dotenv_path, existing_dotenv)?; @@ -153,10 +153,15 @@ async fn setup_bedrock_profile_and_environment() -> Result<()> { let config: toml::Value = toml::from_str(&std::fs::read_to_string(&config_path)?)?; assert_eq!( serde_json::to_value(&config["model_providers"]["amazon-bedrock"]["aws"])?, - json!({"profile": "engineering", "region": "us-west-2"}) + json!({ + "profile": "engineering", + "region": "us-west-2", + "auth_refresh": {"command": "aws", "args": ["login"]} + }) ); assert_eq!(std::fs::read_to_string(&dotenv_path)?, existing_dotenv); + std::fs::write(&config_path, existing_config)?; let response: BedrockSetupResponse = app_server .request(|request_id| ClientRequest::BedrockSetup { request_id, @@ -170,7 +175,10 @@ async fn setup_bedrock_profile_and_environment() -> Result<()> { let config: toml::Value = toml::from_str(&std::fs::read_to_string(&config_path)?)?; assert_eq!( serde_json::to_value(&config["model_providers"]["amazon-bedrock"]["aws"])?, - json!({"region": "us-east-1"}) + json!({ + "region": "us-east-1", + "auth_refresh": {"command": "aws", "args": ["login"]} + }) ); let config: toml::Value = toml::from_str(&std::fs::read_to_string(&config_path)?)?; @@ -184,6 +192,7 @@ async fn setup_bedrock_profile_and_environment() -> Result<()> { [model_providers.amazon-bedrock.aws] region = "us-east-1" + auth_refresh = { command = "aws", args = ["login"] } } .into() ); @@ -319,5 +328,64 @@ async fn setup_bedrock_rejects_invalid_or_conflicting_credentials() -> Result<() ); assert!(!home.path().join(".env").exists()); + for source_file in ["selected.config.toml", "config.toml", "managed_config.toml"] { + let layered_home = TempDir::new()?; + let layered_home_path = + AbsolutePathBuf::from_absolute_path(layered_home.path().canonicalize()?)? + .into_path_buf(); + let config_path = layered_home_path.join("config.toml"); + let source_path = layered_home_path.join(source_file); + let active_path = layered_home_path.join("selected.config.toml"); + let exporter_config = "[model_providers.amazon-bedrock.aws.credential_export]\n\ + command = \"secret-command\"\nargs = [\"secret-argument\"]\n"; + std::fs::write(&config_path, "")?; + std::fs::write(&active_path, "")?; + std::fs::write(&source_path, exporter_config)?; + let active_config_path = active_path.to_string_lossy(); + let mut app_server = bedrock_app_server( + &layered_home_path, + &[( + "CODEX_APP_SERVER_TEST_USER_CONFIG_FILE", + Some(active_config_path.as_ref()), + )], + ) + .await?; + let expected_config = std::fs::read_to_string(&config_path)?; + let expected_active = std::fs::read_to_string(&active_path)?; + let expected_error = "Amazon Bedrock is configured to use `aws.credential_export`. Please clear this setting to use another sign-in method."; + + for (method, params) in [ + ( + "account/bedrock/setup", + json!({"type": "profile", "profile": "engineering", "region": "us-west-2"}), + ), + ( + "account/bedrock/setup", + json!({"type": "environment", "region": "us-west-2"}), + ), + ( + "account/login/start", + json!({"type": "amazonBedrock", "apiKey": "new-key", "region": "us-west-2"}), + ), + ( + "account/login/start", + json!({"type": "amazonBedrockAccessKeys", "accessKeyId": "new-id", "secretAccessKey": "new-secret", "region": "us-west-2"}), + ), + ] { + let request_id = app_server.send_raw_request(method, Some(params)).await?; + let error = timeout( + READ_TIMEOUT, + app_server.read_stream_until_error_message(RequestId::Integer(request_id)), + ) + .await??; + assert_eq!(error.error.message, expected_error); + } + assert_eq!(std::fs::read_to_string(&config_path)?, expected_config); + assert_eq!(std::fs::read_to_string(&source_path)?, exporter_config); + assert_eq!(std::fs::read_to_string(&active_path)?, expected_active); + assert!(!layered_home.path().join("auth.json").exists()); + assert!(!layered_home.path().join(".env").exists()); + } + Ok(()) } diff --git a/codex-rs/aws-auth/src/lib.rs b/codex-rs/aws-auth/src/lib.rs index 8f1bc2b5d7..f79ef91b60 100644 --- a/codex-rs/aws-auth/src/lib.rs +++ b/codex-rs/aws-auth/src/lib.rs @@ -2,6 +2,7 @@ mod config; mod discovery; mod signing; +use std::sync::Arc; use std::time::SystemTime; use aws_credential_types::provider::ProvideCredentials; @@ -44,6 +45,55 @@ impl std::fmt::Debug for AwsAccessKeys { } } +/// Supplies AWS access keys on demand without exposing AWS SDK credential types to callers. +/// +/// Implementations should return current credentials for every call so request signing can +/// observe credential refreshes. Errors must not contain credentials or command output. +pub trait AwsCredentialsProvider: std::fmt::Debug + Send + Sync { + fn credentials( + &self, + ) -> impl std::future::Future> + Send; +} + +#[derive(Debug)] +struct AwsCredentialsProviderAdapter

(Arc

); + +#[derive(Debug, Error)] +#[error("{0}")] +struct ProvidedCredentialsError(std::io::Error); + +impl ProvideCredentials for AwsCredentialsProviderAdapter

{ + fn provide_credentials<'a>( + &'a self, + ) -> aws_credential_types::provider::future::ProvideCredentials<'a> + where + Self: 'a, + { + aws_credential_types::provider::future::ProvideCredentials::new(async move { + let access_keys = self.0.credentials().await.map_err(|error| { + let error = ProvidedCredentialsError(error); + match error.0.kind() { + std::io::ErrorKind::InvalidData + | std::io::ErrorKind::InvalidInput + | std::io::ErrorKind::NotFound + | std::io::ErrorKind::PermissionDenied => { + aws_credential_types::provider::error::CredentialsError::invalid_configuration(error) + } + _ => aws_credential_types::provider::error::CredentialsError::provider_error(error), + } + })?; + + Ok(aws_credential_types::Credentials::new( + access_keys.access_key_id, + access_keys.secret_access_key, + access_keys.session_token, + /*expires_after*/ None, + "codex-bedrock-credential-export", + )) + }) + } +} + /// Generic HTTP request shape consumed by SigV4 signing. #[derive(Debug, Clone, PartialEq, Eq)] pub struct AwsRequestToSign { @@ -135,6 +185,16 @@ impl AwsAuthContext { Ok(context) } + pub async fn load_with_credentials_provider( + config: AwsAuthConfig, + provider: Arc, + ) -> Result { + let mut context = Self::load(config).await?; + context.credentials_provider = + SharedCredentialsProvider::new(AwsCredentialsProviderAdapter(provider)); + Ok(context) + } + pub async fn load_profile(config: AwsAuthConfig) -> Result { let profile = config .profile @@ -171,6 +231,16 @@ impl AwsAuthContext { } impl AwsAuthError { + /// Returns the caller-supplied credential error without exposing SDK error sources. + pub fn credentials_provider_error(&self) -> Option<&std::io::Error> { + let Self::Credentials(error) = self else { + return None; + }; + std::error::Error::source(error)? + .downcast_ref::() + .map(|error| &error.0) + } + /// Returns whether retrying the outbound request can reasonably recover from this auth error. pub fn is_retryable(&self) -> bool { match self { @@ -263,6 +333,68 @@ mod tests { assert!(signing::header_value(&signed.headers, "x-amz-date").is_some()); } + #[tokio::test] + async fn credentials_provider_adapter_converts_keys_and_provider_failures() { + #[derive(Debug)] + struct TestCredentialsProvider(Result); + + impl AwsCredentialsProvider for TestCredentialsProvider { + async fn credentials(&self) -> std::io::Result { + self.0 + .clone() + .map_err(|kind| std::io::Error::new(kind, "credential export failed")) + } + } + + let credentials = + AwsCredentialsProviderAdapter(Arc::new(TestCredentialsProvider(Ok(AwsAccessKeys { + access_key_id: "access-key-id".to_string(), + secret_access_key: "secret-access-key".to_string(), + session_token: Some("session-token".to_string()), + })))) + .provide_credentials() + .await + .expect("exported credentials should be available"); + + assert_eq!( + credentials, + Credentials::new( + "access-key-id", + "secret-access-key", + Some("session-token".to_string()), + /*expires_after*/ None, + "codex-bedrock-credential-export", + ) + ); + + for (kind, retryable) in [ + (std::io::ErrorKind::Other, true), + (std::io::ErrorKind::TimedOut, true), + (std::io::ErrorKind::InvalidData, false), + (std::io::ErrorKind::InvalidInput, false), + (std::io::ErrorKind::NotFound, false), + (std::io::ErrorKind::PermissionDenied, false), + ] { + let error = AwsCredentialsProviderAdapter(Arc::new(TestCredentialsProvider(Err(kind)))) + .provide_credentials() + .await + .expect_err("credential export failure should be propagated"); + let error = AwsAuthError::Credentials(error); + assert_eq!( + ( + error.is_retryable(), + error + .credentials_provider_error() + .map(|error| (error.kind(), error.to_string())), + ), + ( + retryable, + Some((kind, "credential export failed".to_string())) + ), + ); + } + } + #[test] fn credentials_provider_failures_are_retryable() { assert!( diff --git a/codex-rs/config/src/bedrock_runtime_tests.rs b/codex-rs/config/src/bedrock_runtime_tests.rs index 00b86a8028..ea764e771b 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()), + credential_export: None, auth_refresh: None, }) ); diff --git a/codex-rs/core/config.schema.json b/codex-rs/core/config.schema.json index 024ebbc285..c3b46a6b7d 100644 --- a/codex-rs/core/config.schema.json +++ b/codex-rs/core/config.schema.json @@ -419,6 +419,35 @@ ], "type": "object" }, + "AwsCredentialExportConfig": { + "additionalProperties": false, + "description": "Command used to export AWS signing credentials for a model provider.", + "properties": { + "args": { + "default": [], + "description": "Arguments passed to the credential export command.", + "items": { + "type": "string" + }, + "type": "array" + }, + "command": { + "description": "Executable to invoke directly, without a shell.", + "type": "string" + }, + "timeout_ms": { + "default": 30000, + "description": "Maximum time to wait for the credential export command to complete.", + "format": "uint64", + "minimum": 1.0, + "type": "integer" + } + }, + "required": [ + "command" + ], + "type": "object" + }, "BrowserUseConfigToml": { "additionalProperties": false, "properties": { @@ -2215,6 +2244,14 @@ ], "description": "Optional command used to reauthenticate after a refreshable AWS auth failure." }, + "credential_export": { + "allOf": [ + { + "$ref": "#/definitions/AwsCredentialExportConfig" + } + ], + "description": "Optional command whose exported credentials replace the AWS SDK credential chain." + }, "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 35e496f3cf..8ca391de86 100644 --- a/codex-rs/core/src/config/config_tests.rs +++ b/codex-rs/core/src/config/config_tests.rs @@ -1241,7 +1241,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`, `aws.region`, and `aws.auth_refresh`; other non-default provider fields are not supported" + "model_providers.amazon-bedrock only supports changing `base_url`, `auth`, `http_headers`, `aws.profile`, `aws.region`, `aws.credential_export`, 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 f5afbe77c0..5969992b54 100644 --- a/codex-rs/core/tests/suite/external_auth.rs +++ b/codex-rs/core/tests/suite/external_auth.rs @@ -6,6 +6,7 @@ use codex_login::ExternalAuthFuture; use codex_login::ExternalAuthRefreshContext; use codex_login::auth::BedrockAccessKeysAuth; use codex_model_provider_info::AwsAuthRefreshConfig; +use codex_model_provider_info::AwsCredentialExportConfig; use codex_model_provider_info::ModelProviderAwsAuthInfo; use codex_model_provider_info::ModelProviderInfo; use codex_model_provider_info::WireApi; @@ -16,16 +17,21 @@ use codex_protocol::user_input::UserInput; use codex_utils_redacted_string::RedactedString; use core_test_support::responses::ev_completed; use core_test_support::responses::ev_response_created; +use core_test_support::responses::mount_response_sequence; use core_test_support::responses::mount_sse_once; +use core_test_support::responses::mount_sse_sequence; use core_test_support::responses::sse; use core_test_support::responses::start_mock_server; use core_test_support::skip_if_no_network; use core_test_support::test_codex::test_codex; +use core_test_support::wait_for_event; use http::HeaderMap; use http::HeaderValue; use http::header::AUTHORIZATION; use pretty_assertions::assert_eq; use std::num::NonZeroU64; +use std::path::Path; +use std::path::PathBuf; use std::sync::Arc; use std::sync::Mutex; use tokio::process::Command; @@ -200,6 +206,7 @@ async fn amazon_bedrock_managed_access_keys_sign_requests() -> anyhow::Result<() ModelProviderInfo::create_amazon_bedrock_provider(Some(ModelProviderAwsAuthInfo { profile: None, region: Some("us-east-1".to_string()), + credential_export: None, auth_refresh: None, })); provider.base_url = Some(format!("{}/v1", server.uri())); @@ -231,6 +238,191 @@ async fn amazon_bedrock_managed_access_keys_sign_requests() -> anyhow::Result<() Ok(()) } +fn bedrock_credential_export_config( + credentials_path: &Path, + invocations_path: &Path, +) -> anyhow::Result { + #[cfg(unix)] + let (command, args) = { + let script_path = credentials_path.with_extension("sh"); + std::fs::write( + &script_path, + "#!/bin/sh\nprintf 'invoked\\n' >> \"$1\"\ncat \"$2\"\n", + )?; + 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)?; + ( + script_path.to_string_lossy().into_owned(), + vec![ + invocations_path.to_string_lossy().into_owned(), + credentials_path.to_string_lossy().into_owned(), + ], + ) + }; + + #[cfg(windows)] + let (command, args) = { + let script_path = credentials_path.with_extension("cmd"); + std::fs::write( + &script_path, + "@echo off\r\n>> \"%~1\" echo invoked\r\ntype \"%~2\"\r\n", + )?; + ( + "cmd.exe".to_string(), + vec![ + "/D".to_string(), + "/Q".to_string(), + "/C".to_string(), + script_path.to_string_lossy().into_owned(), + invocations_path.to_string_lossy().into_owned(), + credentials_path.to_string_lossy().into_owned(), + ], + ) + }; + + Ok(AwsCredentialExportConfig { + command, + args: args.into_iter().map(RedactedString::from).collect(), + timeout_ms: NonZeroU64::new(5_000).expect("timeout should be non-zero"), + }) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn amazon_bedrock_credential_export_precedence_and_caching() -> anyhow::Result<()> { + skip_if_no_network!(Ok(())); + + let fixture = tempfile::tempdir()?; + let credentials_path = fixture.path().join("credentials.json"); + let invocations_path = fixture.path().join("invocations.txt"); + std::fs::write( + &credentials_path, + serde_json::to_vec(&serde_json::json!({ + "Credentials": { + "AccessKeyId": "exported-access-key-id", + "SecretAccessKey": "exported-secret-access-key", + "SessionToken": "exported-session-token", + "Expiration": "2099-01-01T00:00:00Z", + }, + }))?, + )?; + + let server = start_mock_server().await; + let response_mock = mount_sse_sequence( + &server, + vec![ + sse(vec![ev_response_created("resp-1"), ev_completed("resp-1")]), + sse(vec![ev_response_created("resp-2"), ev_completed("resp-2")]), + ], + ) + .await; + let mut provider = + ModelProviderInfo::create_amazon_bedrock_provider(Some(ModelProviderAwsAuthInfo { + profile: None, + region: Some("us-east-1".to_string()), + credential_export: Some(bedrock_credential_export_config( + &credentials_path, + &invocations_path, + )?), + auth_refresh: None, + })); + provider.base_url = Some(format!("{}/v1", server.uri())); + provider.request_max_retries = Some(0); + provider.stream_max_retries = Some(2); + + let mut builder = test_codex() + .with_model("openai.gpt-5.5") + .with_auth(CodexAuth::BedrockAccessKeys(BedrockAccessKeysAuth { + access_key_id: "managed-access-key-id".to_string(), + secret_access_key: "managed-secret-access-key".to_string(), + session_token: Some("managed-session-token".to_string()), + })) + .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("first request").await?; + test.submit_turn("second request").await?; + + let requests = response_mock.requests(); + assert_eq!(requests.len(), 2); + for request in requests { + let authorization = request + .header("authorization") + .expect("exported AWS credentials should produce SigV4 authorization"); + assert!(authorization.starts_with("AWS4-HMAC-SHA256 ")); + assert!(authorization.contains("Credential=exported-access-key-id/")); + assert_eq!( + request.header("x-amz-security-token").as_deref(), + Some("exported-session-token") + ); + } + let invocations = std::fs::read_to_string(&invocations_path)?; + assert_eq!(invocations.lines().collect::>(), vec!["invoked"]); + + std::fs::write( + &credentials_path, + serde_json::to_vec(&serde_json::json!({ + "Credentials": { + "AccessKeyId": "rotated-access-key-id", + "SecretAccessKey": "rotated-secret-access-key", + "SessionToken": "rotated-session-token", + "Expiration": "2099-01-01T00:00:00Z", + }, + }))?, + )?; + let rotated_response_mock = mount_response_sequence( + &server, + vec![ + ResponseTemplate::new(403).set_body_string("ExpiredTokenException"), + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(sse(vec![ + ev_response_created("resp-3"), + ev_completed("resp-3"), + ])), + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(sse(vec![ + ev_response_created("resp-4"), + ev_completed("resp-4"), + ])), + ], + ) + .await; + + test.submit_turn("recover expired credentials").await?; + test.submit_turn("reuse rotated credentials").await?; + + let rotated_requests = rotated_response_mock.requests(); + assert_eq!(rotated_requests.len(), 3); + for (request, (access_key_id, session_token)) in rotated_requests.iter().zip([ + ("exported-access-key-id", "exported-session-token"), + ("rotated-access-key-id", "rotated-session-token"), + ("rotated-access-key-id", "rotated-session-token"), + ]) { + let authorization = request + .header("authorization") + .expect("exported AWS credentials should produce SigV4 authorization"); + assert!(authorization.contains(&format!("Credential={access_key_id}/"))); + assert_eq!( + request.header("x-amz-security-token").as_deref(), + Some(session_token) + ); + } + let invocations = std::fs::read_to_string(invocations_path)?; + assert_eq!( + invocations.lines().collect::>(), + vec!["invoked", "invoked"] + ); + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn amazon_bedrock_aws_auth_refresh_resigns() -> anyhow::Result<()> { skip_if_no_network!(Ok(())); @@ -252,32 +444,77 @@ async fn amazon_bedrock_aws_auth_refresh_resigns() -> anyhow::Result<()> { 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:?}"); + let path = std::env::join_paths(paths)?; + for mode in ["profile", "export", "export-failure", "export-expired"] { + let credentials = fixture.path().join(format!("{mode}-credentials")); + std::fs::write(&credentials, OLD)?; + let mut command = Command::new(&test_executable); + command + .arg("--exact") + .arg(TEST_NAME) + .env(SUBPROCESS_ENV, mode) + .env("PATH", &path) + .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(), "{mode}: {output:?}"); + } return Ok(()); } + let mode = std::env::var(SUBPROCESS_ENV)?; + let persistent_export_failure = mode == "export-failure"; + let credentials_path = PathBuf::from(std::env::var("AWS_SHARED_CREDENTIALS_FILE")?); + let refresh_invocations = credentials_path.with_extension("refreshes"); if std::env::args().any(|argument| argument == HELPER_ARG) { - std::fs::write(std::env::var("AWS_SHARED_CREDENTIALS_FILE")?, NEW)?; + std::fs::write(&credentials_path, NEW)?; + if !persistent_export_failure { + std::fs::write( + credentials_path.with_extension("json"), + serde_json::to_vec(&serde_json::json!({ + "Version": 1, + "AccessKeyId": "NEW", + "SecretAccessKey": "s", + "SessionToken": "exported-session-token", + }))?, + )?; + } + let mut invocations = std::fs::OpenOptions::new() + .create(true) + .append(true) + .open(&refresh_invocations)?; + std::io::Write::write_all(&mut invocations, b"refresh\n")?; return Ok(()); } + let export_mode = mode.starts_with("export"); + if mode == "export-expired" { + std::fs::write( + credentials_path.with_extension("json"), + serde_json::to_vec(&serde_json::json!({ + "Version": 1, + "AccessKeyId": "OLD", + "SecretAccessKey": "s", + "Expiration": "2000-01-01T00:00:00Z", + }))?, + )?; + } + let export_invocations = credentials_path.with_extension("exports"); + let credential_export = if export_mode { + // Login creates or replaces this file after the initial export fails. + Some(bedrock_credential_export_config( + &credentials_path.with_extension("json"), + &export_invocations, + )?) + } else { + None + }; let server = start_mock_server().await; Mock::given(method("POST")) .and(path("/v1/responses")) @@ -294,14 +531,19 @@ async fn amazon_bedrock_aws_auth_refresh_resigns() -> anyhow::Result<()> { .set_body_string(sse(vec![ev_response_created("r"), ev_completed("r")])) } }) - .expect(2) + .expect(match mode.as_str() { + "profile" => 2, + "export-failure" => 0, + _ => 1, + }) .mount(&server) .await; let mut provider = ModelProviderInfo::create_amazon_bedrock_provider(Some(ModelProviderAwsAuthInfo { - profile: Some("default".to_string()), + profile: (!export_mode).then(|| "default".to_string()), region: Some("us-east-1".to_string()), + credential_export, auth_refresh: Some(AwsAuthRefreshConfig { command: "aws".to_string(), args: Vec::from( @@ -312,7 +554,9 @@ async fn amazon_bedrock_aws_auth_refresh_resigns() -> anyhow::Result<()> { })); provider.base_url = Some(format!("{}/v1", server.uri())); provider.request_max_retries = Some(0); - provider.stream_max_retries = Some(0); + let stream_max_retries = if export_mode { 2 } else { 0 }; + provider.stream_max_retries = Some(stream_max_retries); + let recovery_attempts = stream_max_retries as usize + 1; let mut builder = test_codex() .with_model("openai.gpt-5.5") @@ -321,6 +565,47 @@ async fn amazon_bedrock_aws_auth_refresh_resigns() -> anyhow::Result<()> { config.model_provider = provider; }); let test = builder.build_with_auto_env(&server).await?; + if persistent_export_failure { + test.codex + .start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text { + text: "fail export even after login".into(), + text_elements: Vec::new(), + }])) + .await?; + let EventMsg::TurnComplete(completed) = wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await + else { + unreachable!("predicate guarantees a turn complete event"); + }; + let error = completed.error.expect("failed export should fail the turn"); + assert!( + error + .message + .contains("AWS credential export command exited with"), + "{error:?}" + ); + assert_eq!( + std::fs::read_to_string(refresh_invocations)?, + "refresh\n".repeat(recovery_attempts) + ); + assert_eq!( + std::fs::read_to_string(export_invocations)? + .lines() + .collect::>(), + vec!["invoked"; 2 * recovery_attempts] + ); + assert!( + server + .received_requests() + .await + .expect("mock server should capture requests") + .is_empty() + ); + server.verify().await; + return Ok(()); + } test.codex .start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text { text: "hello".to_string(), @@ -358,6 +643,86 @@ async fn amazon_bedrock_aws_auth_refresh_resigns() -> anyhow::Result<()> { ] ); server.verify().await; + assert_eq!(std::fs::read_to_string(&refresh_invocations)?, "refresh\n"); + if export_mode { + assert_eq!( + std::fs::read_to_string(&export_invocations)? + .lines() + .collect::>(), + vec!["invoked", "invoked"] + ); + let requests = server + .received_requests() + .await + .expect("mock server should capture the exported request"); + assert_eq!( + requests[0].headers.get("x-amz-security-token"), + Some(&HeaderValue::from_static("exported-session-token")) + ); + } + + server.reset().await; + let rejected = mount_response_sequence( + &server, + vec![ + ResponseTemplate::new(403).set_body_string("ExpiredTokenException"); + 2 * recovery_attempts + ], + ) + .await; + test.codex + .start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text { + text: "reject the refreshed credentials too".into(), + text_elements: Vec::new(), + }])) + .await?; + let EventMsg::Error(error) = + wait_for_event(&test.codex, |event| matches!(event, EventMsg::Error(_))).await + else { + unreachable!("predicate guarantees an error event"); + }; + assert!(error.message.contains("ExpiredTokenException"), "{error:?}"); + let EventMsg::TurnComplete(completed) = wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await + else { + unreachable!("predicate guarantees a turn complete event"); + }; + assert_eq!(completed.error, Some(error)); + assert_eq!(rejected.requests().len(), 2 * recovery_attempts); + assert_eq!( + std::fs::read_to_string(refresh_invocations)?, + "refresh\n".repeat(1 + recovery_attempts) + ); + if export_mode { + assert_eq!( + std::fs::read_to_string(export_invocations)? + .lines() + .collect::>(), + vec!["invoked"; 2 + recovery_attempts] + ); + + // A transient auth rejection can succeed on an outer stream retry. + server.reset().await; + let retried = mount_response_sequence( + &server, + vec![ + ResponseTemplate::new(403).set_body_string("ExpiredTokenException"), + ResponseTemplate::new(403).set_body_string("ExpiredTokenException"), + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(sse(vec![ + ev_response_created("retried"), + ev_completed("retried"), + ])), + ], + ) + .await; + test.submit_turn("retry a temporary auth failure").await?; + assert_eq!(retried.requests().len(), 3); + } + server.verify().await; Ok(()) } diff --git a/codex-rs/model-provider-info/src/lib.rs b/codex-rs/model-provider-info/src/lib.rs index 3fb3e87a2a..e486e86124 100644 --- a/codex-rs/model-provider-info/src/lib.rs +++ b/codex-rs/model-provider-info/src/lib.rs @@ -22,11 +22,14 @@ use serde::Serialize; use std::collections::HashMap; use std::fmt; use std::num::NonZeroU64; +use std::path::Component; +use std::path::Path; 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_CREDENTIAL_EXPORT_TIMEOUT_MS: u64 = 30_000; 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`. @@ -159,10 +162,39 @@ pub struct ModelProviderAwsAuthInfo { pub profile: Option, /// AWS region to use for provider-specific endpoints. pub region: Option, + /// Optional command whose exported credentials replace the AWS SDK credential chain. + pub credential_export: Option, /// Optional command used to reauthenticate after a refreshable AWS auth failure. pub auth_refresh: Option, } +/// Command used to export AWS signing credentials for a model provider. +#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq, JsonSchema)] +#[schemars(deny_unknown_fields)] +pub struct AwsCredentialExportConfig { + /// Executable to invoke directly, without a shell. + pub command: String, + /// Arguments passed to the credential export command. + #[serde(default)] + pub args: Vec, + /// Maximum time to wait for the credential export command to complete. + #[serde(default = "default_aws_credential_export_timeout_ms")] + pub timeout_ms: NonZeroU64, +} + +impl AwsCredentialExportConfig { + pub fn timeout(&self) -> Duration { + Duration::from_millis(self.timeout_ms.get()) + } +} + +fn default_aws_credential_export_timeout_ms() -> NonZeroU64 { + match NonZeroU64::new(DEFAULT_AWS_CREDENTIAL_EXPORT_TIMEOUT_MS) { + Some(timeout_ms) => timeout_ms, + None => panic!("AWS credential export timeout must be non-zero"), + } +} + /// Command used to refresh AWS credentials for a model provider. #[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq, JsonSchema)] #[schemars(deny_unknown_fields)] @@ -192,7 +224,7 @@ fn default_aws_auth_refresh_timeout_ms() -> NonZeroU64 { impl ModelProviderInfo { pub fn validate(&self) -> std::result::Result<(), String> { - if self.aws.is_some() { + if let Some(aws) = self.aws.as_ref() { if self.supports_websockets { // TODO(celia-oai): Support AWS SigV4 signing for WebSocket // upgrade requests before allowing AWS-authenticated providers @@ -221,8 +253,33 @@ impl ModelProviderInfo { )); } - if let Some(auth_refresh) = self.aws.as_ref().and_then(|aws| aws.auth_refresh.as_ref()) - { + if let Some(credential_export) = aws.credential_export.as_ref() { + if aws.profile.is_some() { + return Err( + "provider aws.credential_export cannot be combined with aws.profile" + .to_string(), + ); + } + if credential_export.command.trim().is_empty() { + return Err( + "provider aws.credential_export.command must not be empty".to_string() + ); + } + let command = Path::new(&credential_export.command); + let mut components = command.components(); + let is_bare_command = matches!( + (components.next(), components.next()), + (Some(Component::Normal(name)), None) if name == command.as_os_str() + ); + if !command.is_absolute() && !is_bare_command { + return Err( + "provider aws.credential_export.command must be an absolute path or a bare executable name" + .to_string(), + ); + } + } + + if let Some(auth_refresh) = 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()); } @@ -437,6 +494,7 @@ impl ModelProviderInfo { aws: Some(aws.unwrap_or(ModelProviderAwsAuthInfo { profile: None, region: None, + credential_export: None, auth_refresh: None, })), wire_api: WireApi::Responses, @@ -564,7 +622,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`, `aws.region`, and `aws.auth_refresh`; \ +`base_url`, `auth`, `http_headers`, `aws.profile`, `aws.region`, `aws.credential_export`, \ +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 d08af0ffb5..8e13869cb2 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 @@ -227,9 +227,12 @@ name = "Amazon Bedrock" base_url = "https://bedrock.example.com/v1" [aws] -profile = "codex-bedrock" region = "us-west-2" +[aws.credential_export] +command = "aws-vault" +args = ["--profile", "codex-bedrock"] + [aws.auth_refresh] command = "aws" args = ["login", "--profile", "codex-bedrock"] @@ -240,8 +243,13 @@ args = ["login", "--profile", "codex-bedrock"] assert_eq!( provider.aws, Some(ModelProviderAwsAuthInfo { - profile: Some("codex-bedrock".to_string()), + profile: None, region: Some("us-west-2".to_string()), + credential_export: Some(AwsCredentialExportConfig { + command: "aws-vault".to_string(), + args: vec!["--profile".into(), "codex-bedrock".into()], + timeout_ms: NonZeroU64::new(30_000).expect("timeout should be non-zero"), + }), auth_refresh: Some(AwsAuthRefreshConfig { command: "aws".to_string(), args: vec!["login".into(), "--profile".into(), "codex-bedrock".into()], @@ -265,6 +273,7 @@ fn test_create_amazon_bedrock_provider() { aws: Some(ModelProviderAwsAuthInfo { profile: None, region: None, + credential_export: None, auth_refresh: None, }), wire_api: WireApi::Responses, @@ -303,6 +312,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()), + credential_export: None, auth_refresh: None, })); @@ -318,6 +328,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()), + credential_export: None, auth_refresh: None, }), None, @@ -415,6 +426,11 @@ fn test_merge_configured_model_providers_adds_custom_provider() { #[test] fn test_merge_configured_model_providers_applies_amazon_bedrock_aws_override() { + let credential_export = AwsCredentialExportConfig { + command: "aws-vault".to_string(), + args: vec!["export".into(), "codex-bedrock".into()], + timeout_ms: NonZeroU64::new(30_000).expect("timeout should be non-zero"), + }; let auth_refresh = AwsAuthRefreshConfig { command: "aws".to_string(), args: vec!["login".into(), "--profile".into(), "codex-bedrock".into()], @@ -424,8 +440,9 @@ fn test_merge_configured_model_providers_applies_amazon_bedrock_aws_override() { AMAZON_BEDROCK_PROVIDER_ID.to_string(), ModelProviderInfo { aws: Some(ModelProviderAwsAuthInfo { - profile: Some("codex-bedrock".to_string()), + profile: None, region: Some("us-west-2".to_string()), + credential_export: Some(credential_export.clone()), auth_refresh: Some(auth_refresh.clone()), }), ..ModelProviderInfo::default() @@ -437,8 +454,9 @@ fn test_merge_configured_model_providers_applies_amazon_bedrock_aws_override() { .get_mut(AMAZON_BEDROCK_PROVIDER_ID) .expect("Amazon Bedrock provider should be built in") .aws = Some(ModelProviderAwsAuthInfo { - profile: Some("codex-bedrock".to_string()), + profile: None, region: Some("us-west-2".to_string()), + credential_export: Some(credential_export), auth_refresh: Some(auth_refresh), }); @@ -456,6 +474,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()), + credential_export: None, auth_refresh: None, }; let configured_model_providers = std::collections::HashMap::from([( @@ -493,6 +512,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()), + credential_export: None, auth_refresh: None, }), http_headers: Some(maplit::hashmap! { @@ -511,6 +531,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()), + credential_export: None, auth_refresh: None, }); expected_provider @@ -536,6 +557,7 @@ fn test_merge_configured_model_providers_rejects_amazon_bedrock_non_default_fiel aws: Some(ModelProviderAwsAuthInfo { profile: Some("codex-bedrock".to_string()), region: None, + credential_export: None, auth_refresh: None, }), ..ModelProviderInfo::default() @@ -548,7 +570,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`, `aws.region`, and `aws.auth_refresh`; other non-default provider fields are not supported" + "model_providers.amazon-bedrock only supports changing `base_url`, `auth`, `http_headers`, `aws.profile`, `aws.region`, `aws.credential_export`, and `aws.auth_refresh`; other non-default provider fields are not supported" .to_string() ) ); @@ -562,6 +584,7 @@ fn test_merge_configured_model_providers_allows_amazon_bedrock_default_fields() aws: Some(ModelProviderAwsAuthInfo { profile: None, region: None, + credential_export: None, auth_refresh: None, }), wire_api: WireApi::Responses, @@ -584,6 +607,7 @@ fn test_validate_provider_aws_rejects_conflicting_auth() { aws: Some(ModelProviderAwsAuthInfo { profile: None, region: None, + credential_export: None, auth_refresh: None, }), env_key: Some("AWS_BEARER_TOKEN_BEDROCK".to_string()), @@ -603,6 +627,7 @@ fn test_validate_provider_aws_rejects_websockets() { aws: Some(ModelProviderAwsAuthInfo { profile: None, region: None, + credential_export: None, auth_refresh: None, }), requires_openai_auth: false, @@ -633,6 +658,7 @@ fn test_validate_provider_aws_auth_refresh_command() { ModelProviderInfo::create_amazon_bedrock_provider(Some(ModelProviderAwsAuthInfo { profile: None, region: None, + credential_export: None, auth_refresh: Some(AwsAuthRefreshConfig { command: command.to_string(), args: Vec::new(), @@ -644,6 +670,78 @@ fn test_validate_provider_aws_auth_refresh_command() { } } +#[test] +fn test_validate_provider_aws_credential_export_command() { + let absolute_command = std::env::current_exe() + .expect("current executable should have a path") + .to_string_lossy() + .into_owned(); + for (command, expected) in [ + ( + String::new(), + Err("provider aws.credential_export.command must not be empty".to_string()), + ), + ( + " ".to_string(), + Err("provider aws.credential_export.command must not be empty".to_string()), + ), + ("aws-vault".to_string(), Ok(())), + (absolute_command, Ok(())), + ( + "./scripts/export-credentials".to_string(), + Err( + "provider aws.credential_export.command must be an absolute path or a bare executable name" + .to_string(), + ), + ), + ( + "scripts/export-credentials".to_string(), + Err( + "provider aws.credential_export.command must be an absolute path or a bare executable name" + .to_string(), + ), + ), + ( + "scripts/.".to_string(), + Err( + "provider aws.credential_export.command must be an absolute path or a bare executable name" + .to_string(), + ), + ), + ] { + let mut provider = + ModelProviderInfo::create_amazon_bedrock_provider(Some(ModelProviderAwsAuthInfo { + profile: None, + region: Some("us-west-2".to_string()), + credential_export: Some(AwsCredentialExportConfig { + command, + args: vec!["--secret-argument".into()], + timeout_ms: NonZeroU64::new(30_000).expect("timeout should be non-zero"), + }), + auth_refresh: Some(AwsAuthRefreshConfig { + command: "aws".to_string(), + args: Vec::new(), + timeout_ms: NonZeroU64::new(300_000).expect("timeout should be non-zero"), + }), + })); + + assert_eq!(provider.validate(), expected); + if expected.is_ok() { + provider.aws.as_mut().expect("AWS config").profile = Some("codex-bedrock".to_string()); + assert_eq!( + provider.validate(), + Err("provider aws.credential_export cannot be combined with aws.profile".to_string()), + ); + } + let export = provider + .aws + .as_ref() + .and_then(|aws| aws.credential_export.as_ref()) + .expect("credential export should be configured"); + assert!(!format!("{export:?}").contains("--secret-argument")); + } +} + #[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 c10bb3f0bb..9ba32d881a 100644 --- a/codex-rs/model-provider/Cargo.toml +++ b/codex-rs/model-provider/Cargo.toml @@ -13,6 +13,7 @@ path = "src/lib.rs" workspace = true [dependencies] +chrono = { workspace = true } codex-api = { workspace = true } codex-agent-identity = { workspace = true } codex-aws-auth = { workspace = true } @@ -25,14 +26,16 @@ codex-otel = { workspace = true } codex-protocol = { workspace = true } codex-response-debug-context = { workspace = true } http = { workspace = true } +serde = { workspace = true, features = ["derive"] } +serde_json = { workspace = true } sha2 = { workspace = true } -tokio = { workspace = true, features = ["process", "sync", "time"] } +tokio = { workspace = true, features = ["io-util", "process", "sync", "time"] } tracing = { workspace = true, features = ["log"] } [dev-dependencies] base64 = { workspace = true } codex-utils-redacted-string = { workspace = true } pretty_assertions = { workspace = true } -serde_json = { workspace = true } +tempfile = { workspace = true } tokio = { workspace = true, features = ["macros", "rt"] } wiremock = { workspace = true } diff --git a/codex-rs/model-provider/src/amazon_bedrock/auth.rs b/codex-rs/model-provider/src/amazon_bedrock/auth.rs index 706fb6e283..1ea6cc690d 100644 --- a/codex-rs/model-provider/src/amazon_bedrock/auth.rs +++ b/codex-rs/model-provider/src/amazon_bedrock/auth.rs @@ -19,6 +19,7 @@ use codex_protocol::error::Result; use http::HeaderMap; use crate::BearerAuthProvider; +use crate::shared_state::process_shared_state; use super::BedrockEndpoint; use super::mantle::aws_auth_config; @@ -34,6 +35,7 @@ const AWS_DEFAULT_REGION_ENV_VAR: &str = "AWS_DEFAULT_REGION"; #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub(super) enum BedrockAuthSource { CommandBearerToken, + CredentialExport, ConfiguredAwsProfile, ManagedBearerToken, ManagedAccessKeys, @@ -55,6 +57,12 @@ pub(super) fn auth_source( ) -> BedrockAuthSource { if provider_info.has_command_auth() { BedrockAuthSource::CommandBearerToken + } else if provider_info + .aws + .as_ref() + .is_some_and(|aws| aws.credential_export.is_some()) + { + BedrockAuthSource::CredentialExport } else if provider_info .aws .as_ref() @@ -93,6 +101,24 @@ pub(super) async fn resolve_auth_method( "Amazon Bedrock command authentication must be resolved by the model provider" .to_string(), )), + BedrockAuthSource::CredentialExport => { + let config = match endpoint { + BedrockEndpoint::Mantle => aws_auth_config(aws), + BedrockEndpoint::Runtime => runtime::aws_auth_config(aws), + }; + let credential_export = process_shared_state() + .aws_credential_export(aws) + .ok_or_else(|| { + CodexErr::Fatal( + "selected Amazon Bedrock credential exporter is no longer configured" + .to_string(), + ) + })?; + let context = AwsAuthContext::load_with_credentials_provider(config, credential_export) + .await + .map_err(aws_auth_error_to_codex_error)?; + Ok(BedrockAuthMethod::AwsSdkAuth { context }) + } BedrockAuthSource::ManagedBearerToken => { let Some(CodexAuth::BedrockApiKey(auth)) = managed_auth else { return Err(CodexErr::Fatal( @@ -234,6 +260,17 @@ fn aws_auth_error_to_codex_error(error: AwsAuthError) -> CodexErr { } fn aws_auth_error_to_auth_error(error: AwsAuthError) -> AuthError { + if let Some(source) = error.credentials_provider_error() { + // Command failures need provider recovery, not repeated HTTP transport attempts. + return AuthError::Build(if error.is_retryable() { + format!("failed to load AWS credentials: {source}") + } else { + format!( + "{} {source}", + super::error::CREDENTIAL_EXPORT_CONFIG_ERROR_PREFIX + ) + }); + } if error.is_retryable() { AuthError::Transient(error.to_string()) } else { @@ -303,9 +340,12 @@ impl AuthProvider for BedrockSigV4AuthProvider { #[cfg(test)] mod tests { + use std::num::NonZeroU64; + use codex_api::AuthProvider; use codex_login::auth::BedrockAccessKeysAuth; use codex_login::auth::BedrockApiKeyAuth; + use codex_model_provider_info::AwsCredentialExportConfig; use http::HeaderValue; use pretty_assertions::assert_eq; @@ -322,6 +362,18 @@ mod tests { ModelProviderInfo::create_amazon_bedrock_provider(Some(ModelProviderAwsAuthInfo { profile: Some("configured-profile".to_string()), region: Some("us-west-2".to_string()), + credential_export: None, + auth_refresh: None, + })); + let credential_export_provider = + ModelProviderInfo::create_amazon_bedrock_provider(Some(ModelProviderAwsAuthInfo { + profile: None, + region: Some("us-west-2".to_string()), + credential_export: Some(AwsCredentialExportConfig { + command: "export-credentials".to_string(), + args: Vec::new(), + timeout_ms: NonZeroU64::new(30_000).expect("timeout should be non-zero"), + }), auth_refresh: None, })); let managed_auth = @@ -342,6 +394,22 @@ mod tests { &[&str], BedrockAuthSource, )] = &[ + ( + &credential_export_provider, + Some(managed_auth.as_ref()), + &[ + AWS_BEARER_TOKEN_BEDROCK_ENV_VAR, + AWS_ACCESS_KEY_ID_ENV_VAR, + AWS_SECRET_ACCESS_KEY_ENV_VAR, + ], + BedrockAuthSource::CredentialExport, + ), + ( + &credential_export_provider, + Some(managed_access_keys.as_ref()), + &[AWS_BEARER_TOKEN_BEDROCK_ENV_VAR], + BedrockAuthSource::CredentialExport, + ), ( &configured_profile_provider, Some(managed_auth.as_ref()), @@ -418,6 +486,7 @@ mod tests { &ModelProviderAwsAuthInfo { profile: None, region: Some(" us-west-2 ".to_string()), + credential_export: None, auth_refresh: None, }, |name| match name { @@ -450,6 +519,7 @@ mod tests { &ModelProviderAwsAuthInfo { profile: None, region: None, + credential_export: None, auth_refresh: None, }, |name| match name { @@ -468,6 +538,7 @@ mod tests { &ModelProviderAwsAuthInfo { profile: None, region: None, + credential_export: None, auth_refresh: None, }, |name| match name { @@ -486,6 +557,7 @@ mod tests { &ModelProviderAwsAuthInfo { profile: None, region: None, + credential_export: None, auth_refresh: None, }, missing_env_var, diff --git a/codex-rs/model-provider/src/amazon_bedrock/credential_export.rs b/codex-rs/model-provider/src/amazon_bedrock/credential_export.rs new file mode 100644 index 0000000000..7cabae3376 --- /dev/null +++ b/codex-rs/model-provider/src/amazon_bedrock/credential_export.rs @@ -0,0 +1,250 @@ +//! Caches Bedrock signing credentials and provides an exclusive refresh guard. +//! The caller holds the guard across login and export to coalesce auth recovery. +//! Exported secrets remain in memory; command output is bounded and never logged. + +use std::io; +use std::ops::Deref; +use std::path::Component; +use std::path::Path; +use std::process::Stdio; +use std::sync::atomic::AtomicU64; +use std::sync::atomic::Ordering; +use std::time::Duration; +use std::time::SystemTime; + +use chrono::DateTime; +use chrono::Utc; +use codex_aws_auth::AwsAccessKeys; +use codex_aws_auth::AwsCredentialsProvider; +use codex_model_provider_info::AwsCredentialExportConfig; +use serde::Deserialize; +use tokio::io::AsyncReadExt; +use tokio::process::Command; +use tokio::sync::Mutex; +use tokio::sync::MutexGuard; + +const DEFAULT_CREDENTIAL_LIFETIME: Duration = Duration::from_secs(60 * 60); +const CREDENTIAL_REFRESH_WINDOW: Duration = Duration::from_secs(5 * 60); +const MAX_CREDENTIAL_OUTPUT_BYTES: u64 = 64 * 1024; + +#[derive(Debug)] +pub(crate) struct AwsCredentialExport { + config: AwsCredentialExportConfig, + cached: Mutex>, + recovery_generation: AtomicU64, +} + +#[derive(Debug)] +struct CachedCredentials { + access_keys: AwsAccessKeys, + refresh_at: SystemTime, +} + +/// Keeps cached credentials invalidated while the caller reauthenticates and exports. +pub(super) struct CredentialExportRefresh<'a> { + exporter: &'a AwsCredentialExport, + cached: MutexGuard<'a, Option>, +} + +impl CredentialExportRefresh<'_> { + pub(super) async fn refresh(mut self) -> io::Result<()> { + *self.cached = Some(self.exporter.export_credentials().await?); + self.exporter + .recovery_generation + .fetch_add(1, Ordering::Release); + Ok(()) + } +} + +#[derive(Deserialize)] +#[serde(untagged)] +enum CredentialOutput { + Sts { + #[serde(rename = "Credentials")] + credentials: ExportedCredentials, + }, + Process(ExportedCredentials), +} + +#[derive(Deserialize)] +#[serde(rename_all = "PascalCase")] +struct ExportedCredentials { + access_key_id: String, + secret_access_key: String, + session_token: Option, + expiration: Option, +} + +impl CachedCredentials { + fn from_output(output: &[u8], now: SystemTime) -> io::Result { + let output: CredentialOutput = serde_json::from_slice(output).map_err(|_| { + // Serde's detailed errors may quote credential values. + io::Error::new( + io::ErrorKind::InvalidData, + "AWS credential export returned invalid credentials JSON", + ) + })?; + let credentials = match output { + CredentialOutput::Sts { credentials } | CredentialOutput::Process(credentials) => { + credentials + } + }; + if credentials.access_key_id.trim().is_empty() + || credentials.secret_access_key.trim().is_empty() + { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "AWS credential export returned empty access keys", + )); + } + let refresh_at = if let Some(expiration) = credentials.expiration { + let expiration = DateTime::parse_from_rfc3339(&expiration).map_err(|_| { + io::Error::new( + io::ErrorKind::InvalidData, + "AWS credential export Expiration must be an RFC 3339 timestamp", + ) + })?; + let lifetime = expiration + .signed_duration_since(DateTime::::from(now)) + .to_std() + .ok() + .filter(|lifetime| !lifetime.is_zero()) + .ok_or_else(|| { + io::Error::other("AWS credential export returned expired credentials") + })?; + now.checked_add(lifetime.saturating_sub(CREDENTIAL_REFRESH_WINDOW)) + .ok_or_else(|| { + io::Error::new( + io::ErrorKind::InvalidData, + "AWS credential export Expiration is out of range", + ) + })? + } else { + now + DEFAULT_CREDENTIAL_LIFETIME + }; + Ok(Self { + access_keys: AwsAccessKeys { + access_key_id: credentials.access_key_id, + secret_access_key: credentials.secret_access_key, + session_token: credentials.session_token, + }, + refresh_at, + }) + } +} + +impl AwsCredentialExport { + pub(crate) fn new(config: AwsCredentialExportConfig) -> Self { + Self { + config, + cached: Mutex::new(None), + recovery_generation: AtomicU64::new(0), + } + } + + /// Returns None when another overlapping recovery has already completed. + /// Dropping the guard before refresh succeeds leaves the cache empty. + pub(super) async fn begin_refresh(&self) -> Option> { + let generation = self.recovery_generation.load(Ordering::Acquire); + let mut cached = self.cached.lock().await; + if self.recovery_generation.load(Ordering::Acquire) != generation { + return None; + } + *cached = None; + Some(CredentialExportRefresh { + exporter: self, + cached, + }) + } + + async fn export_credentials(&self) -> io::Result { + let program = Path::new(&self.config.command); + let mut components = program.components(); + let is_bare_command = matches!( + (components.next(), components.next()), + (Some(Component::Normal(name)), None) if name == program.as_os_str() + ); + if self.config.command.trim().is_empty() || (!program.is_absolute() && !is_bare_command) { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + "AWS credential export command must be an absolute path or a bare executable name", + )); + } + + let mut command = Command::new(program); + command + .args(self.config.args.iter().map(Deref::deref)) + .stdin(Stdio::null()) + .stdout(Stdio::piped()) + .stderr(Stdio::null()) + .kill_on_drop(true); + + let mut child = command.spawn().map_err(|error| { + io::Error::new( + error.kind(), + format!( + "failed to run AWS credential export command `{}`: {error}", + self.config.command + ), + ) + })?; + let stdout = child + .stdout + .take() + .ok_or_else(|| io::Error::other("AWS credential export stdout is unavailable"))?; + let output = tokio::time::timeout(self.config.timeout(), async { + let mut output = Vec::new(); + stdout + .take(MAX_CREDENTIAL_OUTPUT_BYTES + 1) + .read_to_end(&mut output) + .await?; + if output.len() as u64 > MAX_CREDENTIAL_OUTPUT_BYTES { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "AWS credential export output exceeds 64 KiB", + )); + } + let status = child.wait().await?; + if !status.success() { + return Err(io::Error::other(format!( + "AWS credential export command exited with {status}" + ))); + } + Ok(output) + }) + .await + .map_err(|_| { + io::Error::new( + io::ErrorKind::TimedOut, + format!( + "AWS credential export command timed out after {} ms", + self.config.timeout_ms + ), + ) + })??; + CachedCredentials::from_output(&output, SystemTime::now()) + } +} + +impl AwsCredentialsProvider for AwsCredentialExport { + #[expect( + clippy::await_holding_invalid_type, + reason = "cache misses serialize command execution to avoid duplicate credential exports" + )] + async fn credentials(&self) -> std::io::Result { + let mut cached = self.cached.lock().await; + if let Some(credentials) = cached.as_ref() + && SystemTime::now() < credentials.refresh_at + { + return Ok(credentials.access_keys.clone()); + } + let credentials = self.export_credentials().await?; + let access_keys = credentials.access_keys.clone(); + *cached = Some(credentials); + Ok(access_keys) + } +} + +#[cfg(test)] +#[path = "credential_export_tests.rs"] +mod tests; diff --git a/codex-rs/model-provider/src/amazon_bedrock/credential_export_tests.rs b/codex-rs/model-provider/src/amazon_bedrock/credential_export_tests.rs new file mode 100644 index 0000000000..b027757a61 --- /dev/null +++ b/codex-rs/model-provider/src/amazon_bedrock/credential_export_tests.rs @@ -0,0 +1,348 @@ +use std::future::Future; +use std::num::NonZeroU64; +use std::sync::Arc; +use std::task::Context; +use std::task::Waker; + +use codex_model_provider_info::AwsAuthRefreshConfig; +use codex_model_provider_info::ModelProviderAwsAuthInfo; +use codex_model_provider_info::ModelProviderInfo; +use pretty_assertions::assert_eq; +use serde_json::json; + +use crate::provider::ModelProvider; +use crate::provider::ProviderUnauthorizedRecovery; + +use super::super::AmazonBedrockModelProvider; +use super::super::AwsAuthRecovery; +use super::*; + +const SECRET: &str = "credential-export-secret"; +const VALID_OUTPUT: &[u8] = br#"{"AccessKeyId":"exported","SecretAccessKey":"credential-export-secret","SessionToken":"credential-export-secret"}"#; + +fn exporter_fixture( + output: &[u8], + mode: &str, +) -> io::Result<(tempfile::TempDir, AwsCredentialExport)> { + let directory = tempfile::tempdir()?; + let output_path = directory.path().join("output.json"); + std::fs::write(&output_path, output)?; + + #[cfg(unix)] + let (command, mut args) = { + let script = directory.path().join("export.sh"); + std::fs::write( + &script, + r#"printf 'invoked\n' >> "$1" +printf '%s\n' "$4" >&2 +cat "$2" +case "$3" in + fail) exit 7 ;; + wait) while :; do :; done ;; +esac +"#, + )?; + ( + "sh".to_string(), + vec![script.to_string_lossy().into_owned()], + ) + }; + #[cfg(windows)] + let (command, mut args) = { + let script = directory.path().join("export.cmd"); + std::fs::write( + &script, + "@echo off\r\n>> \"%~1\" echo invoked\r\n>&2 echo %~4\r\ntype \"%~2\"\r\nif \"%~3\"==\"fail\" exit /b 7\r\nif \"%~3\"==\"wait\" goto wait\r\nexit /b 0\r\n:wait\r\ngoto wait\r\n", + )?; + ( + "cmd.exe".to_string(), + vec![ + "/D".to_string(), + "/Q".to_string(), + "/C".to_string(), + script.to_string_lossy().into_owned(), + ], + ) + }; + args.extend([ + directory + .path() + .join("invocations") + .to_string_lossy() + .into_owned(), + output_path.to_string_lossy().into_owned(), + mode.to_string(), + SECRET.to_string(), + ]); + Ok(( + directory, + AwsCredentialExport::new(AwsCredentialExportConfig { + command, + args: args.into_iter().map(Into::into).collect(), + timeout_ms: NonZeroU64::new(5_000).expect("non-zero timeout"), + }), + )) +} + +#[test] +fn credential_output_formats_expiration_and_sanitized_errors() { + let now = SystemTime::UNIX_EPOCH + Duration::from_secs(1_700_000_000); + for (optional_fields, refresh_after) in [ + (json!({}), 3_600), + (json!({"SessionToken": null, "Expiration": null}), 3_600), + ( + json!({"SessionToken": SECRET, "Expiration": "2023-11-14T23:13:20+00:00"}), + 3_300, + ), + ( + json!({"SessionToken": SECRET, "Expiration": "2023-11-14T22:18:20Z"}), + 0, + ), + (json!({"Expiration": "2023-11-14T22:15:20Z"}), 0), + ] { + let mut credentials = json!({"AccessKeyId": "exported", "SecretAccessKey": SECRET}); + credentials + .as_object_mut() + .unwrap() + .extend(optional_fields.as_object().unwrap().clone()); + let mut process_output = credentials.clone(); + process_output["Version"] = json!(1); + for output in [ + process_output, + json!({"Credentials": credentials, "AssumedRoleUser": {"Arn": "test-role"}}), + ] { + let cached = CachedCredentials::from_output(output.to_string().as_bytes(), now) + .expect("valid credentials should parse"); + assert_eq!( + (cached.access_keys, cached.refresh_at), + ( + AwsAccessKeys { + access_key_id: "exported".to_string(), + secret_access_key: SECRET.to_string(), + session_token: optional_fields + .get("SessionToken") + .and_then(serde_json::Value::as_str) + .map(str::to_string), + }, + now + Duration::from_secs(refresh_after) + ) + ); + } + } + + let mut invalid_outputs = vec![( + format!("{{\"SecretAccessKey\":\"{SECRET}\""), + io::ErrorKind::InvalidData, + )]; + for (field, value) in [ + ("AccessKeyId", None), + ("SecretAccessKey", None), + ("AccessKeyId", Some(json!(" "))), + ("SecretAccessKey", Some(json!(" "))), + ("AccessKeyId", Some(json!([SECRET]))), + ("SecretAccessKey", Some(json!([SECRET]))), + ("SessionToken", Some(json!([SECRET]))), + ("Expiration", Some(json!([SECRET]))), + ("Expiration", Some(json!(SECRET))), + ] { + let mut credentials: serde_json::Value = serde_json::from_slice(VALID_OUTPUT).unwrap(); + match value { + Some(value) => { + credentials[field] = value; + } + None => { + credentials.as_object_mut().unwrap().remove(field); + } + } + invalid_outputs.extend( + [ + credentials.to_string(), + json!({"Credentials": credentials}).to_string(), + ] + .map(|output| (output, io::ErrorKind::InvalidData)), + ); + } + for expiration in ["2023-11-14T22:13:20Z", "2023-11-14T22:00:00Z"] { + let credentials = + json!({"AccessKeyId": "exported", "SecretAccessKey": SECRET, "Expiration": expiration}); + invalid_outputs.extend( + [ + credentials.to_string(), + json!({"Credentials": credentials}).to_string(), + ] + .map(|output| (output, io::ErrorKind::Other)), + ); + } + for (output, expected_kind) in invalid_outputs { + let error = CachedCredentials::from_output(output.as_bytes(), now) + .expect_err("invalid credentials should be rejected"); + assert_eq!(error.kind(), expected_kind); + assert!(!error.to_string().contains(SECRET)); + } +} + +#[tokio::test] +async fn credential_export_process_errors_are_bounded_and_sanitized() -> io::Result<()> { + let oversized = format!( + "{SECRET}{}", + "x".repeat(MAX_CREDENTIAL_OUTPUT_BYTES as usize) + ); + for (mode, output, timeout_ms, expected_kind) in [ + ("missing", VALID_OUTPUT, 5_000, io::ErrorKind::NotFound), + ("relative", VALID_OUTPUT, 5_000, io::ErrorKind::InvalidInput), + ( + "normalized-relative", + VALID_OUTPUT, + 5_000, + io::ErrorKind::InvalidInput, + ), + ("fail", VALID_OUTPUT, 5_000, io::ErrorKind::Other), + ("wait", VALID_OUTPUT, 100, io::ErrorKind::TimedOut), + ( + "wait", + oversized.as_bytes(), + 5_000, + io::ErrorKind::InvalidData, + ), + ] { + let (directory, mut exporter) = exporter_fixture(output, mode)?; + exporter.config.timeout_ms = NonZeroU64::new(timeout_ms).expect("non-zero timeout"); + if mode == "missing" { + exporter.config.command = directory + .path() + .join("missing") + .to_string_lossy() + .into_owned(); + } else if mode == "relative" { + exporter.config.command = std::path::Path::new(".") + .join("export-credentials") + .to_string_lossy() + .into_owned(); + } else if mode == "normalized-relative" { + exporter.config.command = std::path::Path::new("scripts") + .join(".") + .to_string_lossy() + .into_owned(); + } + let error = exporter + .credentials() + .await + .expect_err("export should fail"); + assert_eq!(error.kind(), expected_kind); + assert!(!error.to_string().contains(SECRET)); + assert!(!format!("{exporter:?}").contains(SECRET)); + assert!(exporter.cached.lock().await.is_none()); + } + Ok(()) +} + +#[tokio::test] +async fn credential_export_reuses_cache_and_coalesces_refreshes() -> io::Result<()> { + let mut output = VALID_OUTPUT.to_vec(); + output.resize(MAX_CREDENTIAL_OUTPUT_BYTES as usize, b' '); + let (directory, exporter) = exporter_fixture(&output, "success")?; + let original = AwsAccessKeys { + access_key_id: "exported".to_string(), + secret_access_key: SECRET.to_string(), + session_token: Some(SECRET.to_string()), + }; + let (first, second) = tokio::join!(exporter.credentials(), exporter.credentials()); + assert_eq!([first?, second?], [original.clone(), original.clone()]); + let rotated = json!({"AccessKeyId": "rotated", "SecretAccessKey": SECRET}); + std::fs::write(directory.path().join("output.json"), rotated.to_string())?; + assert_eq!(exporter.credentials().await?, original); + + exporter + .cached + .lock() + .await + .as_mut() + .expect("populated cache") + .refresh_at = SystemTime::UNIX_EPOCH; + let rotated = AwsAccessKeys { + access_key_id: "rotated".to_string(), + secret_access_key: SECRET.to_string(), + session_token: None, + }; + assert_eq!(exporter.credentials().await?, rotated); + + let mut provider = AmazonBedrockModelProvider::new( + ModelProviderInfo::create_amazon_bedrock_provider(Some(ModelProviderAwsAuthInfo { + profile: None, + region: Some("us-west-2".to_string()), + credential_export: Some(exporter.config.clone()), + auth_refresh: None, + })), + /*auth_manager*/ None, + ); + let exporter = Arc::new(exporter); + provider.credential_export = Some(exporter.clone()); + { + let first = provider.recover_from_unauthorized(); + let second = provider.recover_from_unauthorized(); + tokio::pin!(first, second); + { + // Both callers must observe the same generation before either can export. + let _cached = exporter.cached.lock().await; + let mut context = Context::from_waker(Waker::noop()); + assert!(first.as_mut().poll(&mut context).is_pending()); + assert!(second.as_mut().poll(&mut context).is_pending()); + } + let (first, second) = tokio::join!(first, second); + assert_eq!( + [ + first.expect("first recovery"), + second.expect("second recovery") + ], + [ProviderUnauthorizedRecovery::Recovered; 2], + ); + } + assert_eq!(exporter.credentials().await?, rotated); + assert_eq!( + std::fs::read_to_string(directory.path().join("invocations"))? + .lines() + .collect::>(), + vec!["invoked", "invoked", "invoked"] + ); + assert!(!format!("{exporter:?}").contains(SECRET)); + + provider.auth_recovery = Some(Arc::new(AwsAuthRecovery::new(AwsAuthRefreshConfig { + command: "not-aws".to_string(), + args: Vec::new(), + timeout_ms: NonZeroU64::new(5_000).expect("non-zero timeout"), + }))); + let credentials = exporter.credentials(); + let refresh = provider.recover_from_unauthorized(); + tokio::pin!(credentials, refresh); + { + // A normal export completing must not make a queued login recovery look completed. + let mut cached = exporter.cached.lock().await; + cached.as_mut().expect("populated cache").refresh_at = SystemTime::UNIX_EPOCH; + let mut context = Context::from_waker(Waker::noop()); + assert!(credentials.as_mut().poll(&mut context).is_pending()); + assert!(refresh.as_mut().poll(&mut context).is_pending()); + } + let (credentials, refresh) = tokio::join!(credentials, refresh); + assert_eq!(credentials?, rotated); + assert!( + !refresh + .expect_err("queued login recovery must still run") + .is_retryable() + ); + assert!(exporter.cached.lock().await.is_none()); + + // Cancellation also leaves the invalidated cache empty and releases the guard. + exporter.credentials().await?; + let refresh = exporter.begin_refresh().await.expect("new recovery"); + let credentials = exporter.credentials(); + tokio::pin!(credentials); + assert!( + credentials + .as_mut() + .poll(&mut Context::from_waker(Waker::noop())) + .is_pending() + ); + drop(refresh); + assert_eq!(credentials.await?, rotated); + Ok(()) +} diff --git a/codex-rs/model-provider/src/amazon_bedrock/error.rs b/codex-rs/model-provider/src/amazon_bedrock/error.rs index ad67f1ff89..dc71b47979 100644 --- a/codex-rs/model-provider/src/amazon_bedrock/error.rs +++ b/codex-rs/model-provider/src/amazon_bedrock/error.rs @@ -4,6 +4,9 @@ use codex_protocol::error::CodexErr; use codex_protocol::error::CodexErrorDetails; use http::StatusCode; +pub(super) const CREDENTIAL_EXPORT_CONFIG_ERROR_PREFIX: &str = + "AWS credential export configuration error:"; + pub(super) const BEDROCK_EXPIRED_SIGNATURE_MESSAGE: &str = concat!( "Amazon Bedrock rejected the request because its AWS signature has expired. ", "Refresh your AWS credentials and retry. If `AWS_BEARER_TOKEN_BEDROCK` is set, ", @@ -11,6 +14,11 @@ pub(super) const BEDROCK_EXPIRED_SIGNATURE_MESSAGE: &str = concat!( ); pub(super) fn map_api_error(error: ApiError) -> CodexErr { + if let ApiError::Transport(TransportError::Build(message)) = &error + && message.starts_with(CREDENTIAL_EXPORT_CONFIG_ERROR_PREFIX) + { + return CodexErr::InvalidRequest(message.clone()); + } let error = codex_api::map_api_error(error); if let CodexErrorDetails::UnexpectedStatus(response) = error.details() && response.status == StatusCode::UNAUTHORIZED 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 692be488e4..7ecfeda5a5 100644 --- a/codex-rs/model-provider/src/amazon_bedrock/error_tests.rs +++ b/codex-rs/model-provider/src/amazon_bedrock/error_tests.rs @@ -1,14 +1,22 @@ +use std::num::NonZeroU64; + use codex_api::ApiError; use codex_api::TransportError; +use codex_model_provider_info::AwsCredentialExportConfig; +use codex_model_provider_info::ModelProviderAwsAuthInfo; +use codex_model_provider_info::ModelProviderInfo; use codex_protocol::error::CodexErrorDetails; use http::HeaderMap; use http::HeaderValue; use http::StatusCode; use pretty_assertions::assert_eq; +use super::AmazonBedrockModelProvider; use super::error::BEDROCK_EXPIRED_SIGNATURE_MESSAGE; +use super::error::CREDENTIAL_EXPORT_CONFIG_ERROR_PREFIX; use super::error::is_refreshable_auth_error; use super::error::map_api_error; +use crate::ModelProvider; const BEDROCK_RESPONSES_URL: &str = "https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses"; @@ -144,3 +152,62 @@ fn classifies_only_refreshable_bedrock_auth_failures() { assert_eq!(is_refreshable_auth_error(&error), expected, "{error}"); } } + +#[test] +fn credential_export_errors_distinguish_configuration_failures_from_retryable_auth() { + let provider = AmazonBedrockModelProvider::new( + ModelProviderInfo::create_amazon_bedrock_provider(Some(ModelProviderAwsAuthInfo { + profile: None, + region: Some("us-west-2".to_string()), + credential_export: Some(AwsCredentialExportConfig { + command: "export-credentials".to_string(), + args: Vec::new(), + timeout_ms: NonZeroU64::new(5_000).expect("non-zero timeout"), + }), + auth_refresh: None, + })), + /*auth_manager*/ None, + ); + for (error, recoverable, retryable) in [ + ( + ApiError::Transport(TransportError::Build(format!( + "{CREDENTIAL_EXPORT_CONFIG_ERROR_PREFIX} invalid credentials JSON" + ))), + false, + false, + ), + ( + ApiError::Transport(TransportError::Build( + "failed to load AWS credentials: export command exited with status 1".to_string(), + )), + true, + true, + ), + ( + http_error(StatusCode::UNAUTHORIZED, "ExpiredTokenException"), + true, + true, + ), + ( + http_error(StatusCode::FORBIDDEN, "ExpiredTokenException"), + true, + true, + ), + ( + http_error(StatusCode::FORBIDDEN, "AccessDeniedException"), + false, + true, + ), + ( + ApiError::Transport(TransportError::Network("offline".to_string())), + false, + true, + ), + ] { + let ApiError::Transport(transport) = &error else { + unreachable!("all test cases are transport errors"); + }; + assert_eq!(provider.is_recoverable_auth_error(transport), recoverable); + assert_eq!(provider.map_api_error(error).is_retryable(), retryable); + } +} diff --git a/codex-rs/model-provider/src/amazon_bedrock/mantle.rs b/codex-rs/model-provider/src/amazon_bedrock/mantle.rs index 60cc69d4c9..1bd2ddee65 100644 --- a/codex-rs/model-provider/src/amazon_bedrock/mantle.rs +++ b/codex-rs/model-provider/src/amazon_bedrock/mantle.rs @@ -94,6 +94,7 @@ mod tests { aws_auth_config(&ModelProviderAwsAuthInfo { profile: Some("codex-bedrock".to_string()), region: None, + credential_export: None, auth_refresh: None, }), AwsAuthConfig { @@ -110,6 +111,7 @@ mod tests { aws_auth_config(&ModelProviderAwsAuthInfo { profile: None, region: Some(" us-west-2 ".to_string()), + credential_export: None, auth_refresh: None, }), AwsAuthConfig { diff --git a/codex-rs/model-provider/src/amazon_bedrock/mod.rs b/codex-rs/model-provider/src/amazon_bedrock/mod.rs index 1543348e95..3a38eb2639 100644 --- a/codex-rs/model-provider/src/amazon_bedrock/mod.rs +++ b/codex-rs/model-provider/src/amazon_bedrock/mod.rs @@ -1,6 +1,7 @@ mod auth; mod auth_refresh; mod catalog; +mod credential_export; mod error; mod mantle; mod runtime; @@ -43,6 +44,7 @@ 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; +pub(crate) use credential_export::AwsCredentialExport; use mantle::bedrock_mantle_runtime_base_url; pub use mantle::is_supported_amazon_bedrock_region; use runtime::bedrock_runtime_base_url; @@ -61,6 +63,7 @@ pub(crate) struct AmazonBedrockModelProvider { aws: ModelProviderAwsAuthInfo, endpoint: BedrockEndpoint, auth_manager: Option>, + credential_export: Option>, auth_recovery: Option>, } @@ -80,11 +83,20 @@ impl AmazonBedrockModelProvider { .unwrap_or(ModelProviderAwsAuthInfo { profile: None, region: None, + credential_export: None, auth_refresh: None, }); + let auth_source = auth::auth_source(&provider_info, auth_manager.as_deref(), std::env::var); + let credential_export = if auth_source == auth::BedrockAuthSource::CredentialExport { + process_shared_state().aws_credential_export(&aws) + } else { + None + }; let uses_aws_sdk_auth = matches!( - auth::auth_source(&provider_info, auth_manager.as_deref(), std::env::var), - auth::BedrockAuthSource::ConfiguredAwsProfile | auth::BedrockAuthSource::AwsSdk + auth_source, + auth::BedrockAuthSource::CredentialExport + | auth::BedrockAuthSource::ConfiguredAwsProfile + | auth::BedrockAuthSource::AwsSdk ); let auth_recovery = if uses_aws_sdk_auth && aws.auth_refresh.is_some() { process_shared_state().aws_auth_recovery(&aws) @@ -97,6 +109,7 @@ impl AmazonBedrockModelProvider { aws, endpoint, auth_manager, + credential_export, auth_recovery, } } @@ -125,11 +138,13 @@ impl AmazonBedrockModelProvider { } fn uses_aws_auth_recovery(&self) -> bool { - self.auth_recovery.is_some() - && matches!( - self.auth_source(), - auth::BedrockAuthSource::ConfiguredAwsProfile | auth::BedrockAuthSource::AwsSdk - ) + let source = self.auth_source(); + source == auth::BedrockAuthSource::CredentialExport + || (self.auth_recovery.is_some() + && matches!( + source, + auth::BedrockAuthSource::ConfiguredAwsProfile | auth::BedrockAuthSource::AwsSdk + )) } async fn auth(&self) -> Option { @@ -140,7 +155,8 @@ impl AmazonBedrockModelProvider { }, auth::BedrockAuthSource::ManagedBearerToken | auth::BedrockAuthSource::ManagedAccessKeys => self.managed_auth(), - auth::BedrockAuthSource::ConfiguredAwsProfile + auth::BedrockAuthSource::CredentialExport + | auth::BedrockAuthSource::ConfiguredAwsProfile | auth::BedrockAuthSource::EnvBearerToken | auth::BedrockAuthSource::EnvAwsCredentials | auth::BedrockAuthSource::AwsSdk => None, @@ -231,7 +247,8 @@ impl ModelProvider for AmazonBedrockModelProvider { auth::BedrockAuthSource::CommandBearerToken | auth::BedrockAuthSource::ManagedBearerToken | auth::BedrockAuthSource::ManagedAccessKeys => self.auth_manager.clone(), - auth::BedrockAuthSource::ConfiguredAwsProfile + auth::BedrockAuthSource::CredentialExport + | auth::BedrockAuthSource::ConfiguredAwsProfile | auth::BedrockAuthSource::EnvBearerToken | auth::BedrockAuthSource::EnvAwsCredentials | auth::BedrockAuthSource::AwsSdk => None, @@ -257,16 +274,39 @@ impl ModelProvider for AmazonBedrockModelProvider { &self, ) -> ModelProviderFuture<'_, Result> { Box::pin(async move { - let Some(recovery) = self - .auth_recovery - .as_ref() - .filter(|_| self.uses_aws_auth_recovery()) - else { + if !self.uses_aws_auth_recovery() { return Ok(ProviderUnauthorizedRecovery::NotConfigured); - }; + } - recovery.refresh().await.map_err(|error| { - if error.kind() == std::io::ErrorKind::InvalidInput { + // Hold the cache guard across both steps so concurrent callers share recovery. + let export_refresh = if let Some(exporter) = &self.credential_export { + let refresh = exporter.begin_refresh().await; + if refresh.is_none() { + // Another caller completed recovery while we were waiting. + return Ok(ProviderUnauthorizedRecovery::Recovered); + } + refresh + } else { + None + }; + let result: std::io::Result<()> = async { + if let Some(recovery) = &self.auth_recovery { + recovery.refresh().await?; + } + if let Some(exporter) = export_refresh { + exporter.refresh().await?; + } + Ok(()) + } + .await; + result.map_err(|error| { + if matches!( + error.kind(), + std::io::ErrorKind::InvalidData + | std::io::ErrorKind::InvalidInput + | std::io::ErrorKind::NotFound + | std::io::ErrorKind::PermissionDenied + ) { CodexErr::InvalidRequest(error.to_string()) } else { CodexErr::Io(error) @@ -388,6 +428,7 @@ 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()), + credential_export: None, auth_refresh: Some(AwsAuthRefreshConfig { command: "aws".to_string(), args: vec!["login".into()], @@ -423,6 +464,7 @@ mod tests { regional_provider_info.aws = Some(ModelProviderAwsAuthInfo { profile: None, region: Some("us-west-2".to_string()), + credential_export: None, auth_refresh: None, }); let regional_provider = @@ -448,6 +490,7 @@ mod tests { let aws = ModelProviderAwsAuthInfo { profile: None, region: Some("us-west-2".to_string()), + credential_export: None, auth_refresh: Some(AwsAuthRefreshConfig { command: "aws".to_string(), args: vec!["login".into()], 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 84f278bed8..73b4752844 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()), + credential_export: None, auth_refresh: None, }), AwsAuthConfig { diff --git a/codex-rs/model-provider/src/provider.rs b/codex-rs/model-provider/src/provider.rs index 027850b788..6239bf0cb5 100644 --- a/codex-rs/model-provider/src/provider.rs +++ b/codex-rs/model-provider/src/provider.rs @@ -514,13 +514,17 @@ impl ModelProvider for ConfiguredModelProvider { #[cfg(test)] mod tests { + use std::future::Future; use std::num::NonZeroU64; + use std::task::Context; + use std::task::Waker; use codex_http_client::HttpClientFactory; 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::AwsCredentialExportConfig; use codex_model_provider_info::ModelProviderAwsAuthInfo; use codex_model_provider_info::WireApi; use codex_model_provider_info::create_oss_provider_with_base_url; @@ -766,6 +770,7 @@ mod tests { ModelProviderInfo::create_amazon_bedrock_provider(Some(ModelProviderAwsAuthInfo { profile: Some("codex-bedrock".to_string()), region: None, + credential_export: None, auth_refresh: None, })), Some(AuthManager::from_auth_for_testing(CodexAuth::from_api_key( @@ -823,6 +828,7 @@ mod tests { let aws = ModelProviderAwsAuthInfo { profile: Some("codex-bedrock".to_string()), region: Some("us-west-2".to_string()), + credential_export: None, auth_refresh: Some(AwsAuthRefreshConfig { command: "aws".to_string(), args: Vec::from( @@ -903,6 +909,93 @@ mod tests { ProviderUnauthorizedRecovery::Recovered ); assert_eq!(read_counter(), "11"); + + let fixture = tempfile::tempdir().expect("export fixture should be created"); + let export_counter = fixture.path().join("exports"); + let export_gate = fixture.path().join("release"); + #[cfg(unix)] + let (command, mut args) = ( + "sh".to_string(), + vec![ + "-c".to_string(), + r#"printf '1\n' >> "$1" +while [ ! -e "$2" ]; do :; done +printf '%s\n' '{"AccessKeyId":"exported","SecretAccessKey":"secret"}' +"# + .to_string(), + "export-credentials".to_string(), + ], + ); + #[cfg(windows)] + let (command, mut args) = { + let script = fixture.path().join("export.cmd"); + std::fs::write( + &script, + concat!( + "@echo off\r\n>> \"%~1\" echo 1\r\n", + ":wait\r\nif not exist \"%~2\" goto wait\r\n", + "echo {\"AccessKeyId\":\"exported\",\"SecretAccessKey\":\"secret\"}\r\n", + ), + ) + .expect("export script should be written"); + ( + "cmd.exe".to_string(), + vec![ + "/D".to_string(), + "/Q".to_string(), + "/C".to_string(), + script.to_string_lossy().into_owned(), + ], + ) + }; + args.extend( + [&export_counter, &export_gate].map(|path| path.to_string_lossy().into_owned()), + ); + let provider_info = + ModelProviderInfo::create_amazon_bedrock_provider(Some(ModelProviderAwsAuthInfo { + profile: None, + credential_export: Some(AwsCredentialExportConfig { + command, + args: args.into_iter().map(RedactedString::from).collect(), + timeout_ms: NonZeroU64::new(5_000).expect("timeout should be non-zero"), + }), + ..aws.clone() + })); + let first_export = create_model_provider(provider_info.clone(), /*auth_manager*/ None); + let second_export = create_model_provider(provider_info, /*auth_manager*/ None); + let first_refresh = first_export.recover_from_unauthorized(); + let second_refresh = second_export.recover_from_unauthorized(); + tokio::pin!(first_refresh, second_refresh); + tokio::time::timeout(std::time::Duration::from_secs(10), async { + tokio::select! { + result = &mut first_refresh => panic!("export should wait for its gate: {result:?}"), + () = async { + while !export_counter.exists() { + tokio::task::yield_now().await; + } + } => {} + } + }).await.expect("export should start after login"); + assert_eq!(read_counter(), "111"); + { + // Queue another recovery after login finishes, while export still holds the lock. + let mut context = Context::from_waker(Waker::noop()); + assert!(second_refresh.as_mut().poll(&mut context).is_pending()); + } + std::fs::write(&export_gate, []).expect("export gate should open"); + let (first_result, second_result) = tokio::join!(first_refresh, second_refresh); + assert_eq!( + [first_result, second_result].map(|result| result.expect("provider should recover")), + [ProviderUnauthorizedRecovery::Recovered; 2] + ); + assert_eq!(read_counter(), "111"); + assert_eq!( + std::fs::read_to_string(export_counter) + .expect("read export counter") + .lines() + .collect::>(), + vec!["1"] + ); std::fs::remove_file(&counter).expect("refresh invocation counter should be removed"); let different_profile = ModelProviderAwsAuthInfo { diff --git a/codex-rs/model-provider/src/shared_state.rs b/codex-rs/model-provider/src/shared_state.rs index 13d348819c..52f8842793 100644 --- a/codex-rs/model-provider/src/shared_state.rs +++ b/codex-rs/model-provider/src/shared_state.rs @@ -6,10 +6,12 @@ use std::sync::Weak; use codex_model_provider_info::ModelProviderAwsAuthInfo; use crate::amazon_bedrock::AwsAuthRecovery; +use crate::amazon_bedrock::AwsCredentialExport; /// Provider-owned runtime state shared across independently configured sessions. #[derive(Debug, Default)] pub(crate) struct ModelProviderSharedState { + aws_credential_exports: Mutex)>>, aws_auth_recoveries: Mutex)>>, } @@ -19,6 +21,30 @@ pub(crate) fn process_shared_state() -> &'static ModelProviderSharedState { } impl ModelProviderSharedState { + pub(crate) fn aws_credential_export( + &self, + aws: &ModelProviderAwsAuthInfo, + ) -> Option> { + let config = aws.credential_export.as_ref()?; + let mut exports = self + .aws_credential_exports + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + exports.retain(|(_, export)| export.strong_count() != 0); + + if let Some(export) = exports + .iter() + .find(|(cached_aws, _)| cached_aws == aws) + .and_then(|(_, export)| export.upgrade()) + { + return Some(export); + } + + let export = Arc::new(AwsCredentialExport::new(config.clone())); + exports.push((aws.clone(), Arc::downgrade(&export))); + Some(export) + } + pub(crate) fn aws_auth_recovery( &self, aws: &ModelProviderAwsAuthInfo, diff --git a/codex-rs/tui/src/app/agents_overview_tests.rs b/codex-rs/tui/src/app/agents_overview_tests.rs index 6dfe993cb7..452bee5912 100644 --- a/codex-rs/tui/src/app/agents_overview_tests.rs +++ b/codex-rs/tui/src/app/agents_overview_tests.rs @@ -1378,10 +1378,13 @@ async fn failed_root_switch_keeps_background_requests_on_the_active_session() -> #[tokio::test] async fn root_switch_loads_local_preferences_from_disk() -> Result<()> { - let mut app = make_test_app().await; - let mut app_server = - crate::start_embedded_app_server_for_picker(app.chat_widget.config_ref()).await?; - let previous = app_server.start_thread(&app.config).await?; + // Keep the large setup and root-switch futures off the test thread's stack. + let mut app = Box::pin(make_test_app()).await; + let mut app_server = Box::pin(crate::start_embedded_app_server_for_picker( + app.chat_widget.config_ref(), + )) + .await?; + let previous = Box::pin(app_server.start_thread(&app.config)).await?; app.enqueue_primary_thread_session(previous.session, previous.turns) .await?; let target_thread_id = ThreadId::from_string( @@ -1399,13 +1402,11 @@ async fn root_switch_loads_local_preferences_from_disk() -> Result<()> { app.local_settings.user_config_path.as_path(), "[tui]\ntheme = \"dracula\"\nresume_cwd = \"session\"\n[history]\npersistence = \"none\"\n", )?; - let config = app - .rebuild_config_for_cwd(app.config.cwd.to_path_buf()) - .await?; + let config = Box::pin(app.rebuild_config_for_cwd(app.config.cwd.to_path_buf())).await?; let expected = crate::local_settings::LocalSettings::from(&config); let mut tui = crate::tui::test_support::make_test_tui()?; - app.select_agents_overview_thread(&mut tui, &mut app_server, target_thread_id) + Box::pin(app.select_agents_overview_thread(&mut tui, &mut app_server, target_thread_id)) .await?; assert_eq!(app.current_displayed_thread_id(), Some(target_thread_id)); diff --git a/codex-rs/tui/src/app/tests.rs b/codex-rs/tui/src/app/tests.rs index 5c29e17017..2f8bfd4408 100644 --- a/codex-rs/tui/src/app/tests.rs +++ b/codex-rs/tui/src/app/tests.rs @@ -2377,8 +2377,9 @@ fn selected_and_resumed_threads_use_server_capability_for_v1_and_v2_children() - .enable_all() .build()?; - runtime.block_on(async { - let (mut app, mut app_event_rx, _op_rx) = make_test_app_with_channels().await; + // Keep the scenario and its large setup/resume futures off the test thread's stack. + runtime.block_on(Box::pin(async { + let (mut app, mut app_event_rx, _op_rx) = Box::pin(make_test_app_with_channels()).await; let root_thread_id = ThreadId::new(); let rollout_dir = app .config @@ -2450,16 +2451,17 @@ fn selected_and_resumed_threads_use_server_capability_for_v1_and_v2_children() - } // Cold-resume the persisted V2 root so its children are registered before selection. - let mut app_server = - crate::start_embedded_app_server_for_picker(app.chat_widget.config_ref()).await?; - let root = app_server - .resume_thread( - &app.local_settings, - app.config.clone(), - root_thread_id, - crate::app_server_session::ResumeModelSettings::PreserveExistingThread, - ) - .await?; + let mut app_server = Box::pin(crate::start_embedded_app_server_for_picker( + app.chat_widget.config_ref(), + )) + .await?; + let root = Box::pin(app_server.resume_thread( + &app.local_settings, + app.config.clone(), + root_thread_id, + crate::app_server_session::ResumeModelSettings::PreserveExistingThread, + )) + .await?; app.enqueue_primary_thread_session(root.session, root.turns) .await?; for (child_thread_id, multi_agent_version) in child_thread_ids @@ -2531,14 +2533,13 @@ fn selected_and_resumed_threads_use_server_capability_for_v1_and_v2_children() - .any(|event| matches!(event, AppEvent::CodexOp(Op::UserTurn { .. }))) ); - let resumed = app_server - .resume_thread( - &app.local_settings, - app.config.clone(), - child_thread_ids[1], - app.resume_model_settings(), - ) - .await?; + let resumed = Box::pin(app_server.resume_thread( + &app.local_settings, + app.config.clone(), + child_thread_ids[1], + app.resume_model_settings(), + )) + .await?; assert!(resumed.blocks_direct_input); app.replace_chat_widget_with_app_server_thread( &mut tui, @@ -2560,7 +2561,7 @@ fn selected_and_resumed_threads_use_server_capability_for_v1_and_v2_children() - .any(|event| matches!(event, AppEvent::CodexOp(Op::UserTurn { .. }))) ); Ok(()) - }) + })) } #[test] diff --git a/codex-rs/tui/src/app_server_session/rollout_history_tests.rs b/codex-rs/tui/src/app_server_session/rollout_history_tests.rs index 0250b747c7..64db875be3 100644 --- a/codex-rs/tui/src/app_server_session/rollout_history_tests.rs +++ b/codex-rs/tui/src/app_server_session/rollout_history_tests.rs @@ -286,7 +286,8 @@ async fn cached_legacy_resume_revalidates_history_across_migration_settings() -> [(false, false), (false, true), (true, false), (true, true)] { let codex_home = tempfile::tempdir().expect("tempdir"); - let config = build_config(&codex_home).await; + // Keep the large setup futures off the test thread's stack. + let config = Box::pin(build_config(&codex_home)).await; let legacy_thread_id = ThreadId::from_string( &create_fake_rollout( codex_home.path(), @@ -314,7 +315,8 @@ async fn cached_legacy_resume_revalidates_history_across_migration_settings() -> let maintenance_guard = codex_rollout::try_acquire_rollout_maintenance_lock(codex_home.path())? .expect("acquire rollout maintenance lock"); - let mut app_server = crate::start_embedded_app_server_for_picker(&startup_config).await?; + let mut app_server = + Box::pin(crate::start_embedded_app_server_for_picker(&startup_config)).await?; app_server.remember_thread_history_mode(legacy_thread_id, ThreadHistoryMode::Legacy); let local_settings = crate::local_settings::LocalSettings::from(&resume_config); let next_request_id = app_server.next_request_id; diff --git a/codex-rs/tui/src/status/tests.rs b/codex-rs/tui/src/status/tests.rs index b6482d462b..41069d770a 100644 --- a/codex-rs/tui/src/status/tests.rs +++ b/codex-rs/tui/src/status/tests.rs @@ -744,6 +744,7 @@ async fn status_uses_server_provider_id_and_auth_requirement() { ModelProviderInfo::create_amazon_bedrock_provider(Some(ModelProviderAwsAuthInfo { profile: None, region: Some("eu-west-1".to_string()), + credential_export: None, auth_refresh: None, })); config.model_provider.base_url =