mirror of
https://github.com/openai/codex.git
synced 2026-10-27 18:48:37 +00:00
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:
@@ -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(
|
||||
|
||||
@@ -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!(
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user