mirror of
https://github.com/openai/codex.git
synced 2026-11-13 06:45:07 +00:00
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:
1 parent
f1bb4c168d
commit
41d3dc56a0
33 files changed
+611
-11
No files matched your search
@@ -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",
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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(_)
|
||||
|
||||
@@ -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(())
|
||||
}
|
||||
|
||||
Reference in new issue
Block a user