mirror of
https://github.com/openai/codex.git
synced 2026-09-20 12:47:38 +00:00
3929c99a97
## Why Bedrock sessions that use the AWS SDK credential chain need a way to recover when credentials expire during a request. ## What changed - Add `aws.auth_refresh` provider configuration with an `aws` command, arguments, and a configurable timeout. - Run the command for refreshable Bedrock authentication failures, reload the SDK credentials, re-sign the request, and retry it. - Share refresh state across matching provider configurations so concurrent failures invoke the command only once. Bearer tokens, command auth, and static environment credentials do not use this recovery path. ## Testing - Add coverage for configuration validation, refreshable error classification, concurrent refresh sharing, and an end-to-end retry signed with refreshed credentials. GitOrigin-RevId: 0302fe3aabdbc1097e7bd62a74d407ba38a3cc57
46 lines
1.4 KiB
Rust
46 lines
1.4 KiB
Rust
use std::sync::Arc;
|
|
use std::sync::Mutex;
|
|
use std::sync::OnceLock;
|
|
use std::sync::Weak;
|
|
|
|
use codex_model_provider_info::ModelProviderAwsAuthInfo;
|
|
|
|
use crate::amazon_bedrock::AwsAuthRecovery;
|
|
|
|
/// Provider-owned runtime state shared across independently configured sessions.
|
|
#[derive(Debug, Default)]
|
|
pub(crate) struct ModelProviderSharedState {
|
|
aws_auth_recoveries: Mutex<Vec<(ModelProviderAwsAuthInfo, Weak<AwsAuthRecovery>)>>,
|
|
}
|
|
|
|
pub(crate) fn process_shared_state() -> &'static ModelProviderSharedState {
|
|
static STATE: OnceLock<ModelProviderSharedState> = OnceLock::new();
|
|
STATE.get_or_init(ModelProviderSharedState::default)
|
|
}
|
|
|
|
impl ModelProviderSharedState {
|
|
pub(crate) fn aws_auth_recovery(
|
|
&self,
|
|
aws: &ModelProviderAwsAuthInfo,
|
|
) -> Option<Arc<AwsAuthRecovery>> {
|
|
let config = aws.auth_refresh.as_ref()?;
|
|
let mut recoveries = self
|
|
.aws_auth_recoveries
|
|
.lock()
|
|
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
|
recoveries.retain(|(_, recovery)| recovery.strong_count() != 0);
|
|
|
|
if let Some(recovery) = recoveries
|
|
.iter()
|
|
.find(|(cached_aws, _)| cached_aws == aws)
|
|
.and_then(|(_, recovery)| recovery.upgrade())
|
|
{
|
|
return Some(recovery);
|
|
}
|
|
|
|
let recovery = Arc::new(AwsAuthRecovery::new(config.clone()));
|
|
recoveries.push((aws.clone(), Arc::downgrade(&recovery)));
|
|
Some(recovery)
|
|
}
|
|
}
|