[codex-core] close revoked auth review gaps [ci changed_files]

This commit is contained in:
Cooper Gamble
2026-05-21 23:26:51 +00:00
parent 053f6633b1
commit 4eedd05bba
9 changed files with 1283 additions and 82 deletions
+118 -30
View File
@@ -1283,6 +1283,7 @@ impl ModelClientSession {
stream,
session_telemetry.clone(),
inference_trace_attempt,
/*revoked_auth_recovery*/ None,
);
return Ok(stream);
}
@@ -1466,6 +1467,7 @@ impl ModelClientSession {
stream_result,
session_telemetry.clone(),
inference_trace_attempt,
auth_recovery,
);
self.websocket_session.last_response_rx = Some(last_request_rx);
return Ok(WebsocketStreamOutcome::Stream(stream));
@@ -1526,37 +1528,52 @@ impl ModelClientSession {
}
let disabled_trace = InferenceTraceContext::disabled();
match self
.stream_responses_websocket(
prompt,
model_info,
session_telemetry,
effort,
summary,
service_tier,
turn_metadata_header,
/*warmup*/ true,
current_span_w3c_trace_context(),
&disabled_trace,
)
.await
{
Ok(WebsocketStreamOutcome::Stream(mut stream)) => {
// Wait for the v2 warmup request to complete before sending the first turn request.
while let Some(event) = stream.next().await {
match event {
Ok(ResponseEvent::Completed { .. }) => break,
Err(err) => return Err(err),
_ => {}
let mut retried_after_auth_recovery = false;
loop {
match self
.stream_responses_websocket(
prompt,
model_info,
session_telemetry,
effort,
summary,
service_tier.clone(),
turn_metadata_header,
/*warmup*/ true,
current_span_w3c_trace_context(),
&disabled_trace,
)
.await
{
Ok(WebsocketStreamOutcome::Stream(mut stream)) => {
// Wait for the v2 warmup request to complete before sending the first turn request.
let mut retry_after_auth_recovery = false;
while let Some(event) = stream.next().await {
match event {
Ok(ResponseEvent::Completed { .. }) => return Ok(()),
Err(err)
if !retried_after_auth_recovery
&& is_websocket_auth_recovery_retry(&err) =>
{
retried_after_auth_recovery = true;
retry_after_auth_recovery = true;
break;
}
Err(err) => return Err(err),
_ => {}
}
}
if retry_after_auth_recovery {
continue;
}
return Ok(());
}
Ok(())
Ok(WebsocketStreamOutcome::FallbackToHttp) => {
self.try_switch_fallback_transport(session_telemetry, model_info);
return Ok(());
}
Err(err) => return Err(err),
}
Ok(WebsocketStreamOutcome::FallbackToHttp) => {
self.try_switch_fallback_transport(session_telemetry, model_info);
Ok(())
}
Err(err) => Err(err),
}
}
@@ -1736,11 +1753,22 @@ fn parent_thread_id_header_value(session_source: &SessionSource) -> Option<Strin
const RESPONSE_STREAM_CHANNEL_CAPACITY: usize = 1600;
const STREAM_DROPPED_REASON: &str = "response stream dropped before provider terminal event";
const WEBSOCKET_AUTH_RECOVERY_RETRY_REASON: &str =
"responses websocket auth recovered after unauthorized response";
pub(crate) fn is_websocket_auth_recovery_retry(err: &CodexErr) -> bool {
matches!(
err,
CodexErr::Stream(message, _)
if message == WEBSOCKET_AUTH_RECOVERY_RETRY_REASON
)
}
fn map_response_stream(
api_stream: codex_api::ResponseStream,
session_telemetry: SessionTelemetry,
inference_trace_attempt: InferenceTraceAttempt,
revoked_auth_recovery: Option<UnauthorizedRecovery>,
) -> (ResponseStream, oneshot::Receiver<LastResponse>) {
let codex_api::ResponseStream {
rx_event,
@@ -1755,6 +1783,7 @@ fn map_response_stream(
api_stream,
session_telemetry,
inference_trace_attempt,
revoked_auth_recovery,
)
}
@@ -1763,6 +1792,7 @@ fn map_response_events<S>(
api_stream: S,
session_telemetry: SessionTelemetry,
inference_trace_attempt: InferenceTraceAttempt,
mut revoked_auth_recovery: Option<UnauthorizedRecovery>,
) -> (ResponseStream, oneshot::Receiver<LastResponse>)
where
S: futures::Stream<Item = std::result::Result<ResponseEvent, ApiError>>
@@ -1873,7 +1903,12 @@ where
if let Some(upstream_request_id) = upstream_request_id {
feedback_tags!(last_model_request_id = upstream_request_id);
}
let mapped = map_api_error(err);
let mapped = map_response_stream_error(
err,
&mut revoked_auth_recovery,
&session_telemetry,
)
.await;
inference_trace_attempt.record_failed(
&mapped,
upstream_request_id,
@@ -1905,6 +1940,59 @@ where
)
}
async fn map_response_stream_error(
err: ApiError,
revoked_auth_recovery: &mut Option<UnauthorizedRecovery>,
session_telemetry: &SessionTelemetry,
) -> CodexErr {
match err {
ApiError::Transport(
transport @ TransportError::Http {
status: StatusCode::UNAUTHORIZED,
..
},
) if stream_error_has_revoked_access_token(&transport) => {
let recovery_transport = clone_http_transport_error(&transport);
match handle_unauthorized(recovery_transport, revoked_auth_recovery, session_telemetry)
.await
{
Ok(_) => CodexErr::Stream(
WEBSOCKET_AUTH_RECOVERY_RETRY_REASON.to_string(),
/*requested_delay*/ None,
),
Err(err) => err,
}
}
err => map_api_error(err),
}
}
fn stream_error_has_revoked_access_token(transport: &TransportError) -> bool {
let debug = extract_response_debug_context(transport);
access_token_revocation_recovery_reason(
debug.auth_error_code.as_deref(),
debug.auth_error_type.as_deref(),
)
.is_some()
}
fn clone_http_transport_error(transport: &TransportError) -> TransportError {
match transport {
TransportError::Http {
status,
url,
headers,
body,
} => TransportError::Http {
status: *status,
url: url.clone(),
headers: headers.clone(),
body: body.clone(),
},
_ => unreachable!("stream auth recovery only clones HTTP transport errors"),
}
}
/// Handles a 401 response by optionally refreshing ChatGPT tokens once.
///
/// When refresh succeeds, the caller should retry the API call; otherwise
@@ -1986,7 +2074,7 @@ async fn handle_unauthorized(
debug.auth_error_type.as_deref(),
);
if let Some(recovery_reason) = revocation_recovery_reason
&& let Some(recovery) = auth_recovery.as_ref()
&& let Some(recovery) = auth_recovery.as_mut()
&& recovery.handles_invalidated_access_token_auth()
{
let mode = recovery.mode_name();
+212 -4
View File
@@ -2,6 +2,7 @@ use super::AuthRequestTelemetryContext;
use super::ModelClient;
use super::PendingUnauthorizedRetry;
use super::UnauthorizedRecoveryExecution;
use super::WEBSOCKET_AUTH_RECOVERY_RETRY_REASON;
use super::X_CODEX_INSTALLATION_ID_HEADER;
use super::X_CODEX_PARENT_THREAD_ID_HEADER;
use super::X_CODEX_TURN_METADATA_HEADER;
@@ -11,12 +12,18 @@ use super::handle_unauthorized;
use crate::AttestationContext;
use crate::AttestationProvider;
use crate::GenerateAttestationFuture;
use async_trait::async_trait;
use base64::Engine;
use codex_api::ApiError;
use codex_api::ResponseEvent;
use codex_api::TransportError;
use codex_app_server_protocol::AuthMode;
use codex_login::AuthManager;
use codex_login::CodexAuth;
use codex_login::ExternalAuth;
use codex_login::ExternalAuthRefreshContext;
use codex_login::ExternalAuthRefreshReason;
use codex_login::ExternalAuthTokens;
use codex_model_provider::BearerAuthProvider;
use codex_model_provider_info::CHATGPT_CODEX_BASE_URL;
use codex_model_provider_info::ModelProviderInfo;
@@ -131,12 +138,12 @@ fn test_session_telemetry() -> SessionTelemetry {
)
}
fn write_managed_chatgpt_auth(codex_home: &Path, access_token: &str) {
fn fake_chatgpt_jwt(signature: &str) -> String {
let encode_json = |value: serde_json::Value| {
base64::engine::general_purpose::URL_SAFE_NO_PAD
.encode(serde_json::to_vec(&value).expect("managed auth JWT segment should serialize"))
};
let id_token = format!(
format!(
"{}.{}.{}",
encode_json(json!({"alg": "none", "typ": "JWT"})),
encode_json(json!({
@@ -147,8 +154,12 @@ fn write_managed_chatgpt_auth(codex_home: &Path, access_token: &str) {
"user_id": "user-12345"
}
})),
base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(b"sig"),
);
base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(signature.as_bytes()),
)
}
fn write_managed_chatgpt_auth(codex_home: &Path, access_token: &str) {
let id_token = fake_chatgpt_jwt("managed-auth");
let auth_json = json!({
"tokens": {
"id_token": id_token,
@@ -178,6 +189,35 @@ async fn managed_chatgpt_auth_manager(access_token: &str) -> (TempDir, Arc<AuthM
(codex_home, manager)
}
struct RefreshingExternalChatgptAuth {
refreshed_token: String,
refresh_count: AtomicUsize,
}
#[async_trait]
impl ExternalAuth for RefreshingExternalChatgptAuth {
fn auth_mode(&self) -> AuthMode {
AuthMode::Chatgpt
}
async fn refresh(
&self,
context: ExternalAuthRefreshContext,
) -> std::io::Result<ExternalAuthTokens> {
assert_eq!(
context.reason,
ExternalAuthRefreshReason::Unauthorized,
"websocket revoked-token recovery should use unauthorized external refresh"
);
self.refresh_count.fetch_add(1, Ordering::SeqCst);
Ok(ExternalAuthTokens::chatgpt(
self.refreshed_token.clone(),
"account_id",
Some("pro".to_string()),
))
}
}
#[derive(Default)]
struct TagCollectorVisitor {
tags: BTreeMap<String, String>,
@@ -401,6 +441,7 @@ async fn dropped_response_stream_traces_cancelled_partial_output() -> anyhow::Re
api_stream,
test_session_telemetry(),
attempt,
/*revoked_auth_recovery*/ None,
);
let observed = stream
@@ -450,6 +491,7 @@ async fn response_stream_records_last_model_feedback_ids() {
api_stream,
test_session_telemetry(),
InferenceTraceAttempt::disabled(),
/*revoked_auth_recovery*/ None,
);
while stream.next().await.is_some() {}
@@ -491,6 +533,7 @@ async fn dropped_backpressured_response_stream_traces_cancelled_partial_output()
api_stream,
test_session_telemetry(),
attempt,
/*revoked_auth_recovery*/ None,
);
// Fill the mapper channel with non-terminal events, then yield one output
@@ -745,6 +788,171 @@ async fn token_revoked_error_code_401_clears_auth_and_requires_relogin() {
assert!(manager.auth_cached().is_none());
}
#[tokio::test]
async fn token_revoked_error_type_401_clears_auth_and_requires_relogin() {
let (_codex_home, manager) = managed_chatgpt_auth_manager("revoked-access-token").await;
let mut recovery = Some(manager.unauthorized_recovery());
let err = handle_unauthorized(
TransportError::Http {
status: StatusCode::UNAUTHORIZED,
url: Some("https://chatgpt.com/backend-api/codex/responses".to_string()),
headers: None,
body: Some(r#"{"error":{"type":"token_revoked"}}"#.to_string()),
},
&mut recovery,
&test_session_telemetry(),
)
.await
.expect_err("revoked access tokens should force relogin");
let CodexErr::RefreshTokenFailed(failed) = err else {
panic!("expected revoked access token to force relogin, got {err:?}");
};
assert_eq!(failed.reason, RefreshTokenFailedReason::Revoked);
assert_eq!(
failed.message,
"Your ChatGPT session is no longer valid. Please sign in again."
);
assert!(manager.auth_cached().is_none());
}
#[tokio::test]
async fn websocket_stream_token_revoked_401_clears_auth_and_requires_relogin() {
let (_codex_home, manager) = managed_chatgpt_auth_manager("revoked-access-token").await;
let api_stream = futures::stream::iter([Err(ApiError::Transport(TransportError::Http {
status: StatusCode::UNAUTHORIZED,
url: None,
headers: None,
body: Some(r#"{"type":"error","status":401,"error":{"type":"token_revoked"}}"#.to_string()),
}))]);
let (mut stream, _) = super::map_response_events(
/*upstream_request_id*/ None,
api_stream,
test_session_telemetry(),
InferenceTraceAttempt::disabled(),
Some(manager.unauthorized_recovery()),
);
let err = stream
.next()
.await
.expect("revoked websocket stream should emit an error")
.expect_err("revoked websocket stream should require relogin");
let CodexErr::RefreshTokenFailed(failed) = err else {
panic!("expected revoked websocket stream to force relogin, got {err:?}");
};
assert_eq!(failed.reason, RefreshTokenFailedReason::Revoked);
assert_eq!(
failed.message,
"Your ChatGPT session is no longer valid. Please sign in again."
);
assert!(manager.auth_cached().is_none());
}
#[tokio::test]
async fn websocket_stream_token_revoked_401_retries_after_loading_replacement_auth() {
let (codex_home, manager) = managed_chatgpt_auth_manager("revoked-access-token").await;
write_managed_chatgpt_auth(codex_home.path(), "replacement-access-token");
let api_stream = futures::stream::iter([Err(ApiError::Transport(TransportError::Http {
status: StatusCode::UNAUTHORIZED,
url: None,
headers: None,
body: Some(r#"{"type":"error","status":401,"error":{"type":"token_revoked"}}"#.to_string()),
}))]);
let (mut stream, _) = super::map_response_events(
/*upstream_request_id*/ None,
api_stream,
test_session_telemetry(),
InferenceTraceAttempt::disabled(),
Some(manager.unauthorized_recovery()),
);
let err = stream
.next()
.await
.expect("recovered websocket stream should emit a retryable error")
.expect_err("recovered websocket stream should ask the turn loop to retry");
let CodexErr::Stream(message, None) = err else {
panic!(
"expected recovered websocket revocation to surface a retryable stream error, got {err:?}"
);
};
assert_eq!(message, WEBSOCKET_AUTH_RECOVERY_RETRY_REASON);
assert_eq!(
manager
.auth_cached()
.expect("replacement auth should remain cached")
.get_token()
.expect("replacement token should resolve"),
"replacement-access-token"
);
}
#[tokio::test]
async fn websocket_stream_token_revoked_401_refreshes_external_auth_before_retry() {
let codex_home = TempDir::new().expect("external auth tempdir");
let initial_access_token = fake_chatgpt_jwt("external-initial");
let refreshed_access_token = fake_chatgpt_jwt("external-refreshed");
codex_login::auth::login_with_chatgpt_auth_tokens(
codex_home.path(),
&initial_access_token,
"account_id",
Some("pro"),
)
.expect("external ChatGPT auth should seed");
let manager = AuthManager::shared(
codex_home.path().to_path_buf(),
/*enable_codex_api_key_env*/ false,
codex_login::AuthCredentialsStoreMode::File,
/*chatgpt_base_url*/ None,
)
.await;
let external_auth = Arc::new(RefreshingExternalChatgptAuth {
refreshed_token: refreshed_access_token.clone(),
refresh_count: AtomicUsize::new(0),
});
manager.set_external_auth(external_auth.clone());
let api_stream = futures::stream::iter([Err(ApiError::Transport(TransportError::Http {
status: StatusCode::UNAUTHORIZED,
url: None,
headers: None,
body: Some(r#"{"type":"error","status":401,"error":{"type":"token_revoked"}}"#.to_string()),
}))]);
let (mut stream, _) = super::map_response_events(
/*upstream_request_id*/ None,
api_stream,
test_session_telemetry(),
InferenceTraceAttempt::disabled(),
Some(manager.unauthorized_recovery()),
);
let err = stream
.next()
.await
.expect("refreshed websocket stream should emit a retryable error")
.expect_err("refreshed websocket stream should ask the turn loop to retry");
let CodexErr::Stream(message, None) = err else {
panic!(
"expected external websocket revocation to surface a retryable stream error, got {err:?}"
);
};
assert_eq!(message, WEBSOCKET_AUTH_RECOVERY_RETRY_REASON);
assert_eq!(external_auth.refresh_count.load(Ordering::SeqCst), 1);
assert_eq!(
manager
.auth_cached()
.expect("refreshed external auth should remain cached")
.get_token()
.expect("refreshed token should resolve"),
refreshed_access_token
);
}
#[tokio::test]
async fn token_invalidated_401_retries_when_persisted_auth_changed() {
let (codex_home, manager) = managed_chatgpt_auth_manager("revoked-access-token").await;
+8
View File
@@ -186,6 +186,7 @@ async fn run_compact_task_inner_impl(
let max_retries = turn_context.provider.info().stream_max_retries();
let mut retries = 0;
let mut websocket_auth_recovery_retries = 0;
let mut client_session = sess.services.model_client.new_session();
// Reuse one client session so turn-scoped state (sticky routing, websocket incremental
// request tracking)
@@ -235,6 +236,13 @@ async fn run_compact_task_inner_impl(
sess.send_event(&turn_context, event).await;
return Err(e);
}
Err(e)
if crate::client::is_websocket_auth_recovery_retry(&e)
&& websocket_auth_recovery_retries == 0 =>
{
websocket_auth_recovery_retries += 1;
continue;
}
Err(e) => {
if retries < max_retries {
retries += 1;
+10
View File
@@ -943,6 +943,7 @@ async fn run_sampling_request(
)
.await;
let mut retries = 0;
let mut websocket_auth_recovery_retries = 0;
let mut initial_input = Some(input);
loop {
let prompt_input = if let Some(input) = initial_input.take() {
@@ -992,6 +993,15 @@ async fn run_sampling_request(
return Err(err);
}
// Auth recovery already loaded replacement credentials. Give that one immediate retry
// without spending the transport reconnect budget.
if crate::client::is_websocket_auth_recovery_retry(&err)
&& websocket_auth_recovery_retries == 0
{
websocket_auth_recovery_retries += 1;
continue;
}
// Use the configured provider-specific stream retry budget.
let max_retries = turn_context.provider.info().stream_max_retries();
if retries >= max_retries
+4 -1
View File
@@ -467,7 +467,10 @@ impl TestCodexBuilder {
let installation_id = resolve_installation_id(&config.codex_home).await?;
let thread_manager = ThreadManager::new(
&config,
codex_core::test_support::auth_manager_from_auth(auth.clone()),
codex_core::test_support::auth_manager_from_auth_with_home(
auth.clone(),
config.codex_home.to_path_buf(),
),
SessionSource::Exec,
Arc::clone(&environment_manager),
empty_extension_registry(),
+2 -1
View File
@@ -1099,7 +1099,7 @@ async fn revoked_chatgpt_auth_user_turn_clears_auth_and_requests_relogin() -> an
.await?
.expect("managed ChatGPT auth should load");
let mut model_provider = built_in_model_providers(/* openai_base_url */ None)["openai"].clone();
let mut model_provider = built_in_model_providers(/*openai_base_url*/ None)["openai"].clone();
model_provider.base_url = Some(format!("{}/api/codex", server.uri()));
model_provider.supports_websockets = false;
let mut builder = test_codex()
@@ -1120,6 +1120,7 @@ async fn revoked_chatgpt_auth_user_turn_clears_auth_and_requests_relogin() -> an
}],
final_output_json_schema: None,
responsesapi_client_metadata: None,
thread_settings: Default::default(),
})
.await?;
@@ -1,6 +1,7 @@
#![allow(clippy::expect_used, clippy::unwrap_used)]
use codex_api::WS_REQUEST_HEADER_TRACEPARENT_CLIENT_METADATA_KEY;
use codex_api::WS_REQUEST_HEADER_TRACESTATE_CLIENT_METADATA_KEY;
use codex_config::types::AuthCredentialsStoreMode;
use codex_core::ModelClient;
use codex_core::ModelClientSession;
use codex_core::Prompt;
@@ -66,6 +67,42 @@ const TEST_INSTALLATION_ID: &str = "11111111-1111-4111-8111-111111111111";
const X_CODEX_WS_STREAM_REQUEST_START_MS_CLIENT_METADATA_KEY: &str =
"x-codex-ws-stream-request-start-ms";
fn write_managed_chatgpt_auth(codex_home: &TempDir, access_token: &str) {
use base64::Engine as _;
let encode_json = |value: serde_json::Value| {
base64::engine::general_purpose::URL_SAFE_NO_PAD
.encode(serde_json::to_vec(&value).expect("managed auth JWT segment should serialize"))
};
let id_token = format!(
"{}.{}.{}",
encode_json(json!({"alg": "none", "typ": "JWT"})),
encode_json(json!({
"email": "user@example.com",
"https://api.openai.com/auth": {
"chatgpt_account_id": "account_id",
"chatgpt_user_id": "user-12345",
"user_id": "user-12345"
}
})),
base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(b"sig"),
);
let auth_json = json!({
"tokens": {
"id_token": id_token,
"access_token": access_token,
"refresh_token": "test-refresh-token",
"account_id": "account_id"
},
"last_refresh": chrono::Utc::now(),
});
std::fs::write(
codex_home.path().join("auth.json"),
serde_json::to_vec_pretty(&auth_json).expect("managed auth should serialize"),
)
.expect("managed auth should persist");
}
fn assert_request_trace_matches(body: &serde_json::Value, expected_trace: &W3cTraceContext) {
let client_metadata = body["client_metadata"]
.as_object()
@@ -1429,6 +1466,360 @@ async fn responses_websocket_invalid_request_error_with_status_is_forwarded() {
server.shutdown().await;
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn responses_websocket_revoked_managed_auth_clears_auth_and_requests_relogin() {
skip_if_no_network!();
let revoked_token_error = json!({
"type": "error",
"status": 401,
"error": {
"type": "token_revoked",
"message": "revoked"
}
});
let server = start_websocket_server(vec![vec![
vec![
ev_response_created("resp-prewarm"),
ev_completed("resp-prewarm"),
],
vec![revoked_token_error],
]])
.await;
let codex_home = Arc::new(TempDir::new().expect("managed auth tempdir"));
write_managed_chatgpt_auth(codex_home.as_ref(), "revoked-access-token");
let auth = CodexAuth::from_auth_storage(
codex_home.path(),
AuthCredentialsStoreMode::File,
/*chatgpt_base_url*/ None,
)
.await
.expect("managed ChatGPT auth should load")
.expect("managed ChatGPT auth should exist");
let mut builder = test_codex()
.with_home(codex_home.clone())
.with_auth(auth)
.with_config(|config| {
config.model_provider.request_max_retries = Some(0);
config.model_provider.stream_max_retries = Some(0);
});
let test = builder
.build_with_websocket_server(&server)
.await
.expect("build websocket codex");
let submission_id = test
.codex
.submit(Op::UserInput {
environments: None,
items: vec![UserInput::Text {
text: "hello".into(),
text_elements: Vec::new(),
}],
final_output_json_schema: None,
responsesapi_client_metadata: None,
thread_settings: Default::default(),
})
.await
.expect("submission should emit relogin events");
let error_event = wait_for_event(&test.codex, |msg| matches!(msg, EventMsg::Error(_))).await;
let EventMsg::Error(error_event) = error_event else {
unreachable!();
};
assert!(
error_event
.message
.contains("Your ChatGPT session is no longer valid. Please sign in again."),
"unexpected error message for submission {submission_id}: {}",
error_event.message
);
wait_for_event(&test.codex, |msg| matches!(msg, EventMsg::TurnComplete(_))).await;
assert!(
!test.codex_home_path().join("auth.json").exists(),
"revoked managed ChatGPT auth should be removed"
);
assert_eq!(server.single_connection().len(), 2);
server.shutdown().await;
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn responses_websocket_revoked_managed_auth_retries_reloaded_auth() {
skip_if_no_network!();
let revoked_token_error = json!({
"type": "error",
"status": 401,
"error": {
"type": "token_revoked",
"message": "revoked"
}
});
let server = start_websocket_server(vec![
vec![
vec![
ev_response_created("resp-prewarm"),
ev_completed("resp-prewarm"),
],
vec![revoked_token_error],
],
vec![vec![
ev_response_created("resp-retry"),
ev_completed("resp-retry"),
]],
])
.await;
let codex_home = Arc::new(TempDir::new().expect("managed auth tempdir"));
write_managed_chatgpt_auth(codex_home.as_ref(), "revoked-access-token");
let auth = CodexAuth::from_auth_storage(
codex_home.path(),
AuthCredentialsStoreMode::File,
/*chatgpt_base_url*/ None,
)
.await
.expect("managed ChatGPT auth should load")
.expect("managed ChatGPT auth should exist");
let mut builder = test_codex()
.with_home(codex_home.clone())
.with_auth(auth)
.with_config(|config| {
config.model_provider.request_max_retries = Some(0);
config.model_provider.stream_max_retries = Some(0);
});
let test = builder
.build_with_websocket_server(&server)
.await
.expect("build websocket codex");
write_managed_chatgpt_auth(codex_home.as_ref(), "replacement-access-token");
test.submit_turn("hello")
.await
.expect("reloaded managed auth should retry the websocket turn");
let persisted_auth = CodexAuth::from_auth_storage(
codex_home.path(),
AuthCredentialsStoreMode::File,
/*chatgpt_base_url*/ None,
)
.await
.expect("replacement managed auth should load")
.expect("replacement managed auth should persist");
assert_eq!(
persisted_auth
.get_token()
.expect("replacement token should resolve"),
"replacement-access-token"
);
let connections = server.connections();
assert_eq!(connections.len(), 2);
assert_eq!(connections[0].len(), 2);
assert_eq!(connections[1].len(), 1);
let handshakes = server.handshakes();
assert_eq!(handshakes.len(), 2);
assert_eq!(
handshakes[0].header("authorization").as_deref(),
Some("Bearer revoked-access-token")
);
assert_eq!(
handshakes[1].header("authorization").as_deref(),
Some("Bearer replacement-access-token")
);
server.shutdown().await;
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn responses_websocket_revoked_managed_auth_prewarm_retries_reloaded_auth() {
skip_if_no_network!();
let revoked_token_error = json!({
"type": "error",
"status": 401,
"error": {
"type": "token_revoked",
"message": "revoked"
}
});
let server = start_websocket_server(vec![
vec![vec![revoked_token_error]],
vec![
vec![
ev_response_created("warm-retry"),
ev_completed("warm-retry"),
],
vec![ev_response_created("resp-1"), ev_completed("resp-1")],
],
])
.await;
let codex_home = Arc::new(TempDir::new().expect("managed auth tempdir"));
write_managed_chatgpt_auth(codex_home.as_ref(), "revoked-access-token");
let auth = CodexAuth::from_auth_storage(
codex_home.path(),
AuthCredentialsStoreMode::File,
/*chatgpt_base_url*/ None,
)
.await
.expect("managed ChatGPT auth should load")
.expect("managed ChatGPT auth should exist");
write_managed_chatgpt_auth(codex_home.as_ref(), "replacement-access-token");
let mut builder = test_codex()
.with_home(codex_home.clone())
.with_auth(auth)
.with_config(|config| {
config.model_provider.request_max_retries = Some(0);
config.model_provider.stream_max_retries = Some(0);
});
let test = builder
.build_with_websocket_server(&server)
.await
.expect("build websocket codex");
test.submit_turn("hello")
.await
.expect("reloaded managed auth should retry startup websocket prewarm");
let persisted_auth = CodexAuth::from_auth_storage(
codex_home.path(),
AuthCredentialsStoreMode::File,
/*chatgpt_base_url*/ None,
)
.await
.expect("replacement managed auth should load")
.expect("replacement managed auth should persist");
assert_eq!(
persisted_auth
.get_token()
.expect("replacement token should resolve"),
"replacement-access-token"
);
let connections = server.connections();
assert_eq!(connections.len(), 2);
assert_eq!(connections[0].len(), 1);
assert_eq!(connections[1].len(), 2);
let handshakes = server.handshakes();
assert_eq!(handshakes.len(), 2);
assert_eq!(
handshakes[0].header("authorization").as_deref(),
Some("Bearer revoked-access-token")
);
assert_eq!(
handshakes[1].header("authorization").as_deref(),
Some("Bearer replacement-access-token")
);
server.shutdown().await;
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn responses_websocket_revoked_managed_auth_compact_retries_reloaded_auth() {
skip_if_no_network!();
let revoked_token_error = json!({
"type": "error",
"status": 401,
"error": {
"type": "token_revoked",
"message": "revoked"
}
});
let server = start_websocket_server(vec![
vec![
vec![
ev_response_created("resp-prewarm"),
ev_completed("resp-prewarm"),
],
vec![ev_response_created("resp-1"), ev_completed("resp-1")],
vec![revoked_token_error],
],
vec![vec![
ev_response_created("compact-retry"),
ev_assistant_message("msg-compact", "COMPACT_SUMMARY"),
ev_completed("compact-retry"),
]],
])
.await;
let codex_home = Arc::new(TempDir::new().expect("managed auth tempdir"));
write_managed_chatgpt_auth(codex_home.as_ref(), "revoked-access-token");
let auth = CodexAuth::from_auth_storage(
codex_home.path(),
AuthCredentialsStoreMode::File,
/*chatgpt_base_url*/ None,
)
.await
.expect("managed ChatGPT auth should load")
.expect("managed ChatGPT auth should exist");
let mut builder = test_codex()
.with_home(codex_home.clone())
.with_auth(auth)
.with_config(|config| {
config.model_provider.request_max_retries = Some(0);
config.model_provider.stream_max_retries = Some(0);
});
let test = builder
.build_with_websocket_server(&server)
.await
.expect("build websocket codex");
test.submit_turn("hello")
.await
.expect("initial websocket turn should complete");
write_managed_chatgpt_auth(codex_home.as_ref(), "replacement-access-token");
test.codex
.submit(Op::Compact)
.await
.expect("compact submission should succeed");
wait_for_event(&test.codex, |msg| matches!(msg, EventMsg::TurnComplete(_))).await;
let persisted_auth = CodexAuth::from_auth_storage(
codex_home.path(),
AuthCredentialsStoreMode::File,
/*chatgpt_base_url*/ None,
)
.await
.expect("replacement managed auth should load")
.expect("replacement managed auth should persist");
assert_eq!(
persisted_auth
.get_token()
.expect("replacement token should resolve"),
"replacement-access-token"
);
let connections = server.connections();
assert_eq!(connections.len(), 2);
assert_eq!(connections[0].len(), 3);
assert_eq!(connections[1].len(), 1);
let handshakes = server.handshakes();
assert_eq!(handshakes.len(), 2);
assert_eq!(
handshakes[0].header("authorization").as_deref(),
Some("Bearer revoked-access-token")
);
assert_eq!(
handshakes[1].header("authorization").as_deref(),
Some("Bearer replacement-access-token")
);
server.shutdown().await;
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn responses_websocket_connection_limit_error_reconnects_and_completes() {
skip_if_no_network!();