codex: stabilize flaky tests on PR #13593

This commit is contained in:
Ahmed Ibrahim
2026-03-08 20:52:27 -07:00
parent 05332b0e96
commit af618c8dd2
38 changed files with 1267 additions and 725 deletions
@@ -29,6 +29,7 @@ use core_test_support::wait_for_event_match;
use pretty_assertions::assert_eq;
use serde_json::Value;
use serde_json::json;
use serial_test::serial;
use std::ffi::OsString;
use std::fs;
use std::time::Duration;
@@ -37,7 +38,6 @@ use tokio::sync::oneshot;
const STARTUP_CONTEXT_HEADER: &str = "Startup context from Codex.";
const MEMORY_PROMPT_PHRASE: &str =
"You have access to a memory folder with guidance from prior runs.";
fn websocket_request_text(
request: &core_test_support::responses::WebSocketRequest,
) -> Option<String> {
@@ -54,6 +54,34 @@ fn websocket_request_instructions(
.map(str::to_owned)
}
async fn wait_for_matching_websocket_request<F>(
server: &core_test_support::responses::WebSocketTestServer,
description: &str,
predicate: F,
) -> core_test_support::responses::WebSocketRequest
where
F: Fn(&core_test_support::responses::WebSocketRequest) -> bool,
{
let deadline = tokio::time::Instant::now() + Duration::from_secs(10);
loop {
if let Some(request) = server
.connections()
.iter()
.flat_map(|connection| connection.iter())
.find(|request| predicate(request))
.cloned()
{
return request;
}
assert!(
tokio::time::Instant::now() < deadline,
"timed out waiting for {description}"
);
tokio::time::sleep(Duration::from_millis(10)).await;
}
}
async fn seed_recent_thread(
test: &TestCodex,
title: &str,
@@ -230,6 +258,7 @@ async fn conversation_start_audio_text_close_round_trip() -> Result<()> {
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial(openai_api_key_env)]
async fn conversation_start_uses_openai_env_key_fallback_with_chatgpt_auth() -> Result<()> {
skip_if_no_network!(Ok(()));
@@ -665,17 +694,12 @@ async fn conversation_uses_experimental_realtime_ws_startup_context_override() -
}))
.await?;
wait_for_event_match(&test.codex, |msg| match msg {
EventMsg::RealtimeConversationRealtime(RealtimeConversationRealtimeEvent {
payload: RealtimeEvent::SessionUpdated { session_id, .. },
}) if session_id == "sess_custom_context" => Some(Ok(())),
EventMsg::Error(err) => Some(Err(err.clone())),
_ => None,
})
.await
.unwrap_or_else(|err: ErrorEvent| panic!("conversation start failed: {err:?}"));
let startup_context_request = server.wait_for_request(1, 0).await;
let startup_context_request = wait_for_matching_websocket_request(
&server,
"startup context request with instructions",
|request| websocket_request_instructions(request).is_some(),
)
.await;
let instructions = websocket_request_instructions(&startup_context_request)
.expect("custom startup context request should contain instructions");
@@ -723,17 +747,12 @@ async fn conversation_disables_realtime_startup_context_with_empty_override() ->
}))
.await?;
wait_for_event_match(&test.codex, |msg| match msg {
EventMsg::RealtimeConversationRealtime(RealtimeConversationRealtimeEvent {
payload: RealtimeEvent::SessionUpdated { session_id, .. },
}) if session_id == "sess_no_context" => Some(Ok(())),
EventMsg::Error(err) => Some(Err(err.clone())),
_ => None,
})
.await
.unwrap_or_else(|err: ErrorEvent| panic!("conversation start failed: {err:?}"));
let startup_context_request = server.wait_for_request(1, 0).await;
let startup_context_request = wait_for_matching_websocket_request(
&server,
"startup context disable request with instructions",
|request| websocket_request_instructions(request).is_some(),
)
.await;
let instructions = websocket_request_instructions(&startup_context_request)
.expect("startup context disable request should contain instructions");
@@ -777,17 +796,12 @@ async fn conversation_start_injects_startup_context_from_thread_history() -> Res
}))
.await?;
wait_for_event_match(&test.codex, |msg| match msg {
EventMsg::RealtimeConversationRealtime(RealtimeConversationRealtimeEvent {
payload: RealtimeEvent::SessionUpdated { session_id, .. },
}) if session_id == "sess_context" => Some(Ok(())),
EventMsg::Error(err) => Some(Err(err.clone())),
_ => None,
})
.await
.unwrap_or_else(|err: ErrorEvent| panic!("conversation start failed: {err:?}"));
let startup_context_request = server.wait_for_request(1, 0).await;
let startup_context_request = wait_for_matching_websocket_request(
&server,
"startup context request with instructions",
|request| websocket_request_instructions(request).is_some(),
)
.await;
let startup_context = websocket_request_instructions(&startup_context_request)
.expect("startup context request should contain instructions");
@@ -831,17 +845,12 @@ async fn conversation_startup_context_falls_back_to_workspace_map() -> Result<()
}))
.await?;
wait_for_event_match(&test.codex, |msg| match msg {
EventMsg::RealtimeConversationRealtime(RealtimeConversationRealtimeEvent {
payload: RealtimeEvent::SessionUpdated { session_id, .. },
}) if session_id == "sess_workspace" => Some(Ok(())),
EventMsg::Error(err) => Some(Err(err.clone())),
_ => None,
})
.await
.unwrap_or_else(|err: ErrorEvent| panic!("conversation start failed: {err:?}"));
let startup_context_request = server.wait_for_request(1, 0).await;
let startup_context_request = wait_for_matching_websocket_request(
&server,
"workspace-map startup context request with instructions",
|request| websocket_request_instructions(request).is_some(),
)
.await;
let startup_context = websocket_request_instructions(&startup_context_request)
.expect("startup context request should contain instructions");
@@ -883,17 +892,12 @@ async fn conversation_startup_context_is_truncated_and_sent_once_per_start() ->
}))
.await?;
wait_for_event_match(&test.codex, |msg| match msg {
EventMsg::RealtimeConversationRealtime(RealtimeConversationRealtimeEvent {
payload: RealtimeEvent::SessionUpdated { session_id, .. },
}) if session_id == "sess_truncated" => Some(Ok(())),
EventMsg::Error(err) => Some(Err(err.clone())),
_ => None,
})
.await
.unwrap_or_else(|err: ErrorEvent| panic!("conversation start failed: {err:?}"));
let startup_context_request = server.wait_for_request(1, 0).await;
let startup_context_request = wait_for_matching_websocket_request(
&server,
"truncated startup context request with instructions",
|request| websocket_request_instructions(request).is_some(),
)
.await;
let startup_context = websocket_request_instructions(&startup_context_request)
.expect("startup context request should contain instructions");
assert!(startup_context.contains(STARTUP_CONTEXT_HEADER));
@@ -905,7 +909,11 @@ async fn conversation_startup_context_is_truncated_and_sent_once_per_start() ->
}))
.await?;
let explicit_text_request = server.wait_for_request(1, 1).await;
let explicit_text_request =
wait_for_matching_websocket_request(&server, "explicit realtime text request", |request| {
websocket_request_text(request).as_deref() == Some("hello")
})
.await;
assert_eq!(
websocket_request_text(&explicit_text_request),
Some("hello".to_string())
+51 -16
View File
@@ -4,6 +4,7 @@ use std::ffi::OsString;
use std::fs;
use std::net::TcpListener;
use std::path::Path;
use std::sync::Arc;
use std::time::Duration;
use std::time::SystemTime;
use std::time::UNIX_EPOCH;
@@ -36,11 +37,13 @@ use core_test_support::skip_if_no_network;
use core_test_support::stdio_server_bin;
use core_test_support::test_codex::test_codex;
use core_test_support::wait_for_event;
use core_test_support::wait_for_event_with_timeout;
use reqwest::Client;
use reqwest::StatusCode;
use serde_json::Value;
use serde_json::json;
use serial_test::serial;
use tempfile::tempdir;
use tokio::net::TcpStream;
use tokio::process::Child;
use tokio::process::Command;
use tokio::time::Instant;
@@ -263,7 +266,7 @@ async fn stdio_image_responses_round_trip() -> anyhow::Result<()> {
let tools_ready_deadline = Instant::now() + Duration::from_secs(30);
loop {
fixture.codex.submit(Op::ListMcpTools).await?;
let list_event = core_test_support::wait_for_event_with_timeout(
let list_event = wait_for_event_with_timeout(
&fixture.codex,
|ev| matches!(ev, EventMsg::McpListToolsResponse(_)),
Duration::from_secs(10),
@@ -851,10 +854,8 @@ async fn streamable_http_tool_call_round_trip() -> anyhow::Result<()> {
Ok(())
}
/// This test writes to a fallback credentials file in CODEX_HOME.
/// Ideally, we wouldn't need to serialize the test but it's much more cumbersome to wire CODEX_HOME through the code.
#[serial(codex_home)]
#[test]
#[serial(mcp_test_value)]
fn streamable_http_with_oauth_round_trip() -> anyhow::Result<()> {
const TEST_STACK_SIZE_BYTES: usize = 8 * 1024 * 1024;
@@ -936,8 +937,8 @@ async fn streamable_http_with_oauth_round_trip_impl() -> anyhow::Result<()> {
wait_for_streamable_http_server(&mut http_server_child, &bind_addr, Duration::from_secs(5))
.await?;
let temp_home = tempdir()?;
let _guard = EnvVarGuard::set("CODEX_HOME", temp_home.path().as_os_str());
let temp_home = Arc::new(tempdir()?);
let _codex_home_guard = EnvVarGuard::set("CODEX_HOME", temp_home.path().as_os_str());
write_fallback_oauth_tokens(
temp_home.path(),
server_name,
@@ -948,10 +949,10 @@ async fn streamable_http_with_oauth_round_trip_impl() -> anyhow::Result<()> {
)?;
let fixture = test_codex()
.with_home(temp_home.clone())
.with_config(move |config| {
// This test seeds OAuth tokens in CODEX_HOME/.credentials.json and
// validates file-backed OAuth loading. Force file mode so Linux
// keyring backend quirks do not affect this test.
// Keep OAuth credentials isolated to this test home because Bazel
// runs the full core suite in one process.
config.mcp_oauth_credentials_store_mode = serde_json::from_value(json!("file"))
.expect("`file` should deserialize as OAuthCredentialsStoreMode");
let mut servers = config.mcp_servers.get().clone();
@@ -984,6 +985,31 @@ async fn streamable_http_with_oauth_round_trip_impl() -> anyhow::Result<()> {
.await?;
let session_model = fixture.session_configured.model.clone();
let tools_ready_deadline = Instant::now() + Duration::from_secs(30);
loop {
fixture.codex.submit(Op::ListMcpTools).await?;
let list_event = wait_for_event_with_timeout(
&fixture.codex,
|ev| matches!(ev, EventMsg::McpListToolsResponse(_)),
Duration::from_secs(10),
)
.await;
let EventMsg::McpListToolsResponse(tool_list) = list_event else {
unreachable!("event guard guarantees McpListToolsResponse");
};
if tool_list.tools.contains_key(&tool_name) {
break;
}
let available_tools: Vec<&str> = tool_list.tools.keys().map(String::as_str).collect();
if Instant::now() >= tools_ready_deadline {
panic!(
"timed out waiting for MCP tool {tool_name} to become available; discovered tools: {available_tools:?}"
);
}
sleep(Duration::from_millis(200)).await;
}
fixture
.codex
.submit(Op::UserTurn {
@@ -1078,7 +1104,8 @@ async fn wait_for_streamable_http_server(
timeout: Duration,
) -> anyhow::Result<()> {
let deadline = Instant::now() + timeout;
let metadata_url = format!("http://{address}/.well-known/oauth-authorization-server/mcp");
let client = Client::builder().no_proxy().build()?;
loop {
if let Some(status) = server_child.try_wait()? {
return Err(anyhow::anyhow!(
@@ -1090,22 +1117,30 @@ async fn wait_for_streamable_http_server(
if remaining.is_zero() {
return Err(anyhow::anyhow!(
"timed out waiting for streamable HTTP server at {address}: deadline reached"
"timed out waiting for streamable HTTP server metadata at {metadata_url}: deadline reached"
));
}
match tokio::time::timeout(remaining, TcpStream::connect(address)).await {
Ok(Ok(_)) => return Ok(()),
match tokio::time::timeout(remaining, client.get(&metadata_url).send()).await {
Ok(Ok(response)) if response.status() == StatusCode::OK => return Ok(()),
Ok(Ok(response)) => {
if Instant::now() >= deadline {
return Err(anyhow::anyhow!(
"timed out waiting for streamable HTTP server metadata at {metadata_url}: HTTP {}",
response.status()
));
}
}
Ok(Err(error)) => {
if Instant::now() >= deadline {
return Err(anyhow::anyhow!(
"timed out waiting for streamable HTTP server at {address}: {error}"
"timed out waiting for streamable HTTP server metadata at {metadata_url}: {error}"
));
}
}
Err(_) => {
return Err(anyhow::anyhow!(
"timed out waiting for streamable HTTP server at {address}: connect call timed out"
"timed out waiting for streamable HTTP server metadata at {metadata_url}: request timed out"
));
}
}
@@ -6,63 +6,40 @@ use codex_core::WireApi;
use codex_protocol::protocol::EventMsg;
use codex_protocol::protocol::Op;
use codex_protocol::user_input::UserInput;
use codex_utils_cargo_bin::find_resource;
use core_test_support::load_sse_fixture;
use core_test_support::responses::ev_completed;
use core_test_support::responses::ev_response_created;
use core_test_support::responses::sse;
use core_test_support::responses;
use core_test_support::skip_if_no_network;
use core_test_support::streaming_sse::StreamingSseChunk;
use core_test_support::streaming_sse::start_streaming_sse_server;
use core_test_support::test_codex::TestCodex;
use core_test_support::test_codex::test_codex;
use core_test_support::wait_for_event;
use wiremock::Mock;
use wiremock::MockServer;
use wiremock::Request;
use wiremock::Respond;
use wiremock::ResponseTemplate;
use wiremock::matchers::method;
use wiremock::matchers::path;
fn sse_incomplete() -> String {
load_sse_fixture("tests/fixtures/incomplete_sse.json")
let fixture = find_resource!("tests/fixtures/incomplete_sse.json")
.unwrap_or_else(|err| panic!("failed to resolve incomplete_sse fixture: {err}"));
load_sse_fixture(fixture)
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn retries_on_early_close() {
skip_if_no_network!();
let server = MockServer::start().await;
let incomplete_sse = sse_incomplete();
let completed_sse = responses::sse_completed("resp_ok");
struct SeqResponder;
impl Respond for SeqResponder {
fn respond(&self, _: &Request) -> ResponseTemplate {
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering;
static CALLS: AtomicUsize = AtomicUsize::new(0);
let n = CALLS.fetch_add(1, Ordering::SeqCst);
if n == 0 {
ResponseTemplate::new(200)
.insert_header("content-type", "text/event-stream")
.set_body_raw(sse_incomplete(), "text/event-stream")
} else {
ResponseTemplate::new(200)
.insert_header("content-type", "text/event-stream")
.set_body_raw(
sse(vec![
ev_response_created("resp_ok"),
ev_completed("resp_ok"),
]),
"text/event-stream",
)
}
}
}
Mock::given(method("POST"))
.and(path("/v1/responses"))
.respond_with(SeqResponder {})
.expect(2)
.mount(&server)
.await;
let (server, _) = start_streaming_sse_server(vec![
vec![StreamingSseChunk {
gate: None,
body: incomplete_sse,
}],
vec![StreamingSseChunk {
gate: None,
body: completed_sse,
}],
])
.await;
// Configure retry behavior explicitly to avoid mutating process-wide
// environment variables.
@@ -92,7 +69,7 @@ async fn retries_on_early_close() {
.with_config(move |config| {
config.model_provider = model_provider;
})
.build(&server)
.build_with_streaming_server(&server)
.await
.unwrap();
@@ -109,4 +86,13 @@ async fn retries_on_early_close() {
// Wait until TurnComplete (should succeed after retry).
wait_for_event(&codex, |event| matches!(event, EventMsg::TurnComplete(_))).await;
let requests = server.requests().await;
assert_eq!(
requests.len(),
2,
"expected retry after incomplete SSE stream"
);
server.shutdown().await;
}