Surface model provider authentication recovery progress (#41239)

## What changed

- Emit turn-scoped authentication recovery started and completed events when a model provider refreshes expired credentials.
- Add stable `modelProvider/authRecoveryStarted` and `modelProvider/authRecoveryCompleted` app-server notifications with the thread, turn, provider, and user-facing message.
- Show recovery progress in the TUI and `codex exec`, including Amazon Bedrock session reauthentication.

## Testing

- Cover provider recovery success and failure events, app-server routing, client rendering, and Amazon Bedrock credential refresh.

GitOrigin-RevId: 3010c38d0676f18bced27761cf86dd38344d09f9
This commit is contained in:
alexsong-oai authored and copyberry committed 2026-08-28 01:14:02 +00:00
1 parent f1bb4c168d
commit 41d3dc56a0
33 files changed
+611 -11

No files matched your search

+52 -1
View File
@@ -30,6 +30,7 @@ use std::sync::OnceLock;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
use async_channel::Sender;
use codex_api::AgentIdentityTelemetry;
use codex_api::ApiError;
use codex_api::AuthProvider;
@@ -83,6 +84,9 @@ use codex_protocol::config_types::Verbosity as VerbosityConfig;
use codex_protocol::models::ResponseItem;
use codex_protocol::openai_models::ModelInfo;
use codex_protocol::openai_models::ReasoningEffort as ReasoningEffortConfig;
use codex_protocol::protocol::AuthRecoveryEvent;
use codex_protocol::protocol::Event as ProtocolEvent;
use codex_protocol::protocol::EventMsg;
use codex_protocol::protocol::InternalSessionSource;
use codex_protocol::protocol::SessionSource;
use codex_protocol::protocol::W3cTraceContext;
@@ -291,6 +295,7 @@ pub struct ModelClient {
agent_identity_policy: AgentIdentityAuthPolicy,
prompt_cache_key_override: Option<String>,
free_guardian_enabled: bool,
event_sender: Option<Sender<ProtocolEvent>>,
http_client_factory: HttpClientFactory,
}
@@ -512,6 +517,7 @@ impl ModelClient {
agent_identity_policy,
prompt_cache_key_override: None,
free_guardian_enabled: false,
event_sender: None,
http_client_factory,
}
}
@@ -521,11 +527,13 @@ impl ModelClient {
self
}
pub(crate) fn with_prompt_cache_key_override(
pub(crate) fn with_session_context(
mut self,
prompt_cache_key_override: Option<String>,
event_sender: Sender<ProtocolEvent>,
) -> Self {
self.prompt_cache_key_override = prompt_cache_key_override;
self.event_sender = Some(event_sender);
self
}
@@ -1672,6 +1680,8 @@ impl ModelClientSession {
&mut provider_auth_recovery_attempted,
session_telemetry,
&self.client.state.provider,
self.client.event_sender.as_ref(),
responses_metadata.turn_id.as_deref(),
)
.await?,
);
@@ -1805,6 +1815,8 @@ impl ModelClientSession {
&mut provider_auth_recovery_attempted,
session_telemetry,
&provider,
self.client.event_sender.as_ref(),
responses_metadata.turn_id.as_deref(),
)
.await?,
);
@@ -2380,18 +2392,57 @@ struct WebsocketConnectParams<'a> {
endpoint: ResponsesEndpoint,
}
fn emit_auth_recovery_event(
event_sender: Option<&Sender<ProtocolEvent>>,
turn_id: Option<&str>,
provider: &SharedModelProvider,
message: &str,
event: fn(AuthRecoveryEvent) -> EventMsg,
) {
if let (Some(sender), Some(turn_id)) = (event_sender, turn_id) {
let _ = sender.try_send(ProtocolEvent {
id: turn_id.to_string(),
msg: event(AuthRecoveryEvent {
provider: provider.info().name.clone(),
message: message.to_string(),
}),
});
}
}
async fn handle_unauthorized(
transport: TransportError,
auth_recovery: &mut Option<UnauthorizedRecovery>,
provider_auth_recovery_attempted: &mut bool,
session_telemetry: &SessionTelemetry,
provider: &SharedModelProvider,
event_sender: Option<&Sender<ProtocolEvent>>,
turn_id: Option<&str>,
) -> Result<UnauthorizedRecoveryExecution> {
let debug = extract_response_debug_context(&transport);
if !*provider_auth_recovery_attempted {
*provider_auth_recovery_attempted = true;
let messages = provider.auth_recovery_messages();
if let Some(messages) = messages {
emit_auth_recovery_event(
event_sender,
turn_id,
provider,
messages.started,
EventMsg::AuthRecoveryStarted,
);
}
match provider.recover_from_unauthorized().await {
Ok(ProviderUnauthorizedRecovery::Recovered) => {
if let Some(messages) = messages {
emit_auth_recovery_event(
event_sender,
turn_id,
provider,
messages.succeeded,
EventMsg::AuthRecoveryCompleted,
);
}
return Ok(UnauthorizedRecoveryExecution {
mode: "provider",
phase: "provider_refresh",
+38
View File
@@ -31,6 +31,7 @@ use codex_model_provider::BearerAuthProvider;
use codex_model_provider::ModelProvider;
use codex_model_provider::ModelProviderFuture;
use codex_model_provider::ProviderAccountResult;
use codex_model_provider::ProviderAuthRecoveryMessages;
use codex_model_provider::ProviderUnauthorizedRecovery;
use codex_model_provider::SharedModelProvider;
use codex_model_provider::create_model_provider;
@@ -895,6 +896,8 @@ async fn bedrock_unauthorized_error_uses_provider_mapping() {
&mut provider_auth_recovery_attempted,
&test_session_telemetry(),
&provider,
/*event_sender*/ None,
/*turn_id*/ None,
)
.await
.expect_err("expired Bedrock signature should fail");
@@ -931,6 +934,13 @@ impl ModelProvider for TestRecoveryProvider {
self.inner.account_state()
}
fn auth_recovery_messages(&self) -> Option<ProviderAuthRecoveryMessages> {
Some(ProviderAuthRecoveryMessages {
started: "Refreshing provider authentication.",
succeeded: "Provider authentication recovered.",
})
}
fn recover_from_unauthorized(
&self,
) -> ModelProviderFuture<'_, codex_protocol::error::Result<ProviderUnauthorizedRecovery>> {
@@ -975,12 +985,15 @@ async fn provider_owned_auth_recovery_is_bounded_and_preserves_unauthorized_fail
let mut auth_recovery = None;
let mut provider_auth_recovery_attempted = false;
let telemetry = test_session_telemetry();
let (event_sender, event_receiver) = async_channel::unbounded();
let result = super::handle_unauthorized(
unauthorized(),
&mut auth_recovery,
&mut provider_auth_recovery_attempted,
&telemetry,
&provider,
Some(&event_sender),
Some("turn-1"),
)
.await;
@@ -998,6 +1011,8 @@ async fn provider_owned_auth_recovery_is_bounded_and_preserves_unauthorized_fail
&mut provider_auth_recovery_attempted,
&telemetry,
&provider,
Some(&event_sender),
Some("turn-1"),
)
.await
.expect_err("provider recovery should not run more than once")
@@ -1011,6 +1026,29 @@ async fn provider_owned_auth_recovery_is_bounded_and_preserves_unauthorized_fail
other => panic!("unexpected error after provider recovery: {other}"),
}
assert_eq!(attempts.load(Ordering::Relaxed), 1);
let events = std::iter::from_fn(|| event_receiver.try_recv().ok())
.map(|event| serde_json::to_value(event).expect("recovery event should serialize"))
.collect::<Vec<_>>();
let mut expected = vec![json!({
"id": "turn-1",
"msg": {
"type": "auth_recovery_started",
"provider": provider.info().name,
"message": "Refreshing provider authentication.",
}
})];
if !should_fail {
expected.push(json!({
"id": "turn-1",
"msg": {
"type": "auth_recovery_completed",
"provider": provider.info().name,
"message": "Provider authentication recovered.",
}
}));
}
assert_eq!(events, expected);
}
}
+2 -1
View File
@@ -1422,11 +1422,12 @@ impl Session {
config.http_client_factory(),
)
.with_free_guardian_enabled(config.free_guardian_enabled())
.with_prompt_cache_key_override(
.with_session_context(
crate::guardian::prompt_cache_key_override_for_review_session(
&session_configuration.session_source,
session_configuration.parent_thread_id,
),
tx_event.clone(),
),
executed_tool_calls: executed_tool_calls.clone(),
code_mode_service: crate::tools::code_mode::CodeModeService::new(
+2
View File
@@ -1797,6 +1797,8 @@ pub(super) fn realtime_text_for_event(msg: &EventMsg) -> Option<(String, Option<
}
EventMsg::Error(_)
| EventMsg::Warning(_)
| EventMsg::AuthRecoveryStarted(_)
| EventMsg::AuthRecoveryCompleted(_)
| EventMsg::GuardianWarning(_)
| EventMsg::RealtimeConversationStarted(_)
| EventMsg::RealtimeConversationSdp(_)
+40 -1
View File
@@ -1,3 +1,4 @@
use codex_core::TurnInputRequest;
use codex_login::AuthHeaders;
use codex_login::CodexAuth;
use codex_login::ExternalAuth;
@@ -9,6 +10,9 @@ use codex_model_provider_info::ModelProviderAwsAuthInfo;
use codex_model_provider_info::ModelProviderInfo;
use codex_model_provider_info::WireApi;
use codex_model_provider_info::create_oss_provider_with_base_url;
use codex_protocol::protocol::AuthRecoveryEvent;
use codex_protocol::protocol::EventMsg;
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;
@@ -317,7 +321,42 @@ async fn amazon_bedrock_aws_auth_refresh_resigns() -> anyhow::Result<()> {
config.model_provider = provider;
});
let test = builder.build_with_auto_env(&server).await?;
test.submit_turn("hello").await?;
test.codex
.start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text {
text: "hello".to_string(),
text_elements: Vec::new(),
}]))
.await?;
let mut recovery_events = Vec::new();
loop {
match core_test_support::wait_for_event(&test.codex, |_| true).await {
EventMsg::AuthRecoveryStarted(event) => recovery_events.push(("started", event)),
EventMsg::AuthRecoveryCompleted(event) => recovery_events.push(("completed", event)),
EventMsg::TurnComplete(_) => break,
_ => {}
}
}
assert_eq!(
recovery_events,
vec![
(
"started",
AuthRecoveryEvent {
provider: "Amazon Bedrock".to_string(),
message: "AWS session has expired. Reauthenticating...".to_string(),
},
),
(
"completed",
AuthRecoveryEvent {
provider: "Amazon Bedrock".to_string(),
message: "Signed in with AWS.".to_string(),
},
),
]
);
server.verify().await;
Ok(())
}