Make login token-refresh tests wait for auth publication (#57248)

## Why

The same-owner token-refresh tests need the original backend response to remain pending until `getAuthStatus` publishes refreshed credentials. A fixed response delay does not guarantee that ordering.

## What changed

Gate responses in the `getAuthStatus` test cases until a request-specific auth reload log confirms the credential change. Emit that log after updating the cached auth and generation, and include `auth_changed_for_refresh` as a structured field.

Extend the JSON log test helper to match events with predicates, allowing tests to identify the refresh by its request span.

GitOrigin-RevId: 698f3be499ea3c035f082c501836dae4178adaef
This commit is contained in:
goodman-oai
2026-10-26 20:53:10 +00:00
committed by copyberry
parent 12b2b58e63
commit 8f71a8657e
4 changed files with 105 additions and 29 deletions
@@ -28,31 +28,15 @@ impl JsonLogCapture {
pub(crate) async fn wait_for_event(
&self,
event_name: &str,
description: &str,
predicate: impl Fn(&Value) -> bool,
timeout: Duration,
) -> Result<Value> {
let mut events = self
.wait_for_events(event_name, /*count*/ 1, timeout)
.await?;
Ok(events.remove(0))
}
pub(crate) async fn wait_for_events(
&self,
event_name: &str,
count: usize,
timeout: Duration,
) -> Result<Vec<Value>> {
let result = tokio::time::timeout(timeout, async {
loop {
let updated = self.updated.notified();
let events = self
.events()?
.into_iter()
.filter(|event| event["fields"]["event.name"].as_str() == Some(event_name))
.collect::<Vec<_>>();
if events.len() >= count {
return Ok(events);
if let Some(event) = self.events()?.into_iter().find(&predicate) {
return Ok(event);
}
updated.await;
}
@@ -67,7 +51,7 @@ impl JsonLogCapture {
.unwrap_or_else(std::sync::PoisonError::into_inner)
.join("\n");
anyhow::bail!(
"timed out waiting for {count} JSON log event(s) named `{event_name}`; captured stderr:\n{lines}"
"timed out waiting for a JSON log event {description}; captured stderr:\n{lines}"
)
}
}
@@ -250,7 +250,24 @@ impl TestAppServer {
event_name: &str,
timeout: Duration,
) -> anyhow::Result<serde_json::Value> {
self.json_logs.wait_for_event(event_name, timeout).await
self.wait_for_matching_json_log_event(
&format!("named `{event_name}`"),
|event| event["fields"]["event.name"].as_str() == Some(event_name),
timeout,
)
.await
}
/// Waits for a JSON stderr event matching a predicate, including its request span.
pub async fn wait_for_matching_json_log_event(
&self,
description: &str,
predicate: impl Fn(&serde_json::Value) -> bool,
timeout: Duration,
) -> anyhow::Result<serde_json::Value> {
self.json_logs
.wait_for_event(description, predicate, timeout)
.await
}
async fn new_with_program_env_and_args(
+78 -6
View File
@@ -9,6 +9,12 @@ use app_test_support::DEFAULT_CLIENT_NAME;
use app_test_support::encode_id_token;
use app_test_support::write_chatgpt_auth;
use app_test_support::write_models_cache;
use axum::Router;
use axum::body::Body;
use axum::http::HeaderMap;
use axum::http::Uri;
use axum::http::header::HOST;
use axum::routing::get;
use chrono::Duration as ChronoDuration;
use chrono::Utc;
use codex_app_server_protocol::Account;
@@ -61,8 +67,11 @@ use std::sync::Arc;
use std::time::Duration;
use tempfile::TempDir;
use test_case::test_case;
use tokio::net::TcpListener;
use tokio::sync::Notify;
use tokio::time::timeout;
use tokio_util::sync::CancellationToken;
use tokio_util::task::AbortOnDropHandle;
use url::Url;
use wiremock::Mock;
use wiremock::MockServer;
@@ -2296,11 +2305,53 @@ async fn login_survives_same_owner_token_refresh(
) -> Result<()> {
let codex_home = TempDir::new()?;
let backend = MockServer::start().await;
let release_response = CancellationToken::new();
let (chatgpt_base_url, _response_gate) = match refresh_trigger {
LoginRefreshTrigger::AuthStatus => {
let listener = TcpListener::bind("127.0.0.1:0").await?;
let base_url = format!("http://{}/backend-api", listener.local_addr()?);
let upstream = backend.uri();
let gated_path = delayed_path.to_string();
let release = release_response.clone();
let client = HttpClientBuilder::new().build_direct()?;
// Keep Wiremock's matching and request recording, but hold the response
// asynchronously so the token refresh can use the same mock backend.
let router = Router::new().fallback(get(move |uri: Uri, mut headers: HeaderMap| {
let upstream = upstream.clone();
let gated_path = gated_path.clone();
let release = release.clone();
let client = client.clone();
async move {
let hold_response = uri.path() == gated_path
&& headers
.get("authorization")
.is_some_and(|value| value == "Bearer access-token-123");
headers.remove(HOST);
let response = client
.get(format!("{upstream}{uri}"))
.headers(headers)
.send()
.await
.expect("mock backend response");
if hold_response {
release.cancelled().await;
}
response.into_http_response().map(Body::new)
}
}));
let task =
AbortOnDropHandle::new(tokio::spawn(
async move { axum::serve(listener, router).await },
));
(base_url, Some(task))
}
LoginRefreshTrigger::UnauthorizedConfig => (format!("{}/backend-api", backend.uri()), None),
};
create_config_toml(
codex_home.path(),
CreateConfigTomlParams {
requires_openai_auth: Some(true),
chatgpt_base_url: Some(format!("{}/backend-api", backend.uri())),
chatgpt_base_url: Some(chatgpt_base_url),
..Default::default()
},
)?;
@@ -2352,11 +2403,11 @@ async fn login_survives_same_owner_token_refresh(
.and(header("authorization", "Bearer access-token-123"))
.respond_with(move |_: &wiremock::Request| {
request_started.notify_one();
ResponseTemplate::new(match refresh_trigger {
LoginRefreshTrigger::AuthStatus => 200,
LoginRefreshTrigger::UnauthorizedConfig => 401,
})
.set_delay(Duration::from_secs(/*secs*/ 2))
match refresh_trigger {
LoginRefreshTrigger::AuthStatus => ResponseTemplate::new(200),
LoginRefreshTrigger::UnauthorizedConfig => ResponseTemplate::new(401)
.set_delay(Duration::from_secs(/*secs*/ 2)),
}
.set_body_json(response.clone())
})
.with_priority(1)
@@ -2368,6 +2419,7 @@ async fn login_survives_same_owner_token_refresh(
let issuer = backend.uri();
let refresh_url = format!("{issuer}/oauth/refresh");
let mut server = TestAppServer::builder()
.with_json_logging("warn,codex_login::auth::manager=info,codex_app_server=info")
.with_codex_home(codex_home.path())
.with_env_overrides(&[
("OPENAI_API_KEY", None),
@@ -2404,6 +2456,26 @@ async fn login_survives_same_owner_token_refresh(
refresh_token: Some(true),
})
.await?;
// getAuthStatus reloads config before replying, which can wait for
// this held response. The reload log records the credential change
// after publishing the refreshed auth cache and its generation.
let request_id = request.to_string();
server
.wait_for_matching_json_log_event(
"auth cache refresh in getAuthStatus",
|event| {
event["fields"]["auth_changed_for_refresh"] == true
&& event["spans"].as_array().is_some_and(|spans| {
spans.iter().any(|span| {
span["rpc.method"] == "getAuthStatus"
&& span["rpc.request_id"] == request_id
})
})
},
DEFAULT_READ_TIMEOUT,
)
.await?;
release_response.cancel();
let refreshed: GetAuthStatusResponse =
timeout(DEFAULT_READ_TIMEOUT, server.read_response(request)).await??;
assert_eq!(
+4 -1
View File
@@ -2918,7 +2918,6 @@ impl AuthManager {
if auth_changed_for_refresh {
guard.permanent_refresh_failure = None;
}
tracing::info!("Reloaded auth, changed: {changed}");
guard.auth = new_auth;
if auth_changed_for_refresh {
self.auth_change_state_tx.send_modify(|state| {
@@ -2935,6 +2934,10 @@ impl AuthManager {
});
self.auth_change_tx.send_modify(|revision| *revision += 1);
}
tracing::info!(
auth_changed_for_refresh,
"Reloaded auth, changed: {changed}"
);
changed
}