codex: fix realtime startup context close race (#13593)

This commit is contained in:
Ahmed Ibrahim
2026-03-06 20:28:25 -08:00
parent ce981d77d2
commit 484668e073

View File

@@ -2,6 +2,7 @@ use anyhow::Context;
use anyhow::Result;
use chrono::Utc;
use codex_core::CodexAuth;
use codex_core::CodexThread;
use codex_core::auth::OPENAI_API_KEY_ENV_VAR;
use codex_protocol::ThreadId;
use codex_protocol::protocol::CodexErrorInfo;
@@ -38,6 +39,8 @@ 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.";
const REALTIME_NORMAL_CLOSE_MESSAGE: &str =
"failed to send realtime request: Connection closed normally";
fn websocket_request_text(
request: &core_test_support::responses::WebSocketRequest,
@@ -55,6 +58,32 @@ fn websocket_request_instructions(
.map(str::to_owned)
}
async fn wait_for_session_updated_allowing_clean_close(
codex: &CodexThread,
expected_session_id: &str,
) {
tokio::time::timeout(Duration::from_secs(3), async {
loop {
let event = codex
.next_event()
.await
.expect("realtime conversation event should arrive");
match event.msg {
EventMsg::RealtimeConversationRealtime(RealtimeConversationRealtimeEvent {
payload: RealtimeEvent::SessionUpdated { session_id, .. },
}) if session_id == expected_session_id => return,
EventMsg::Error(err) if err.message.contains(REALTIME_NORMAL_CLOSE_MESSAGE) => {}
EventMsg::Error(err) => panic!("conversation start failed: {err:?}"),
_ => {}
}
}
})
.await
.unwrap_or_else(|_| {
panic!("timeout waiting for session.updated event for {expected_session_id}")
});
}
async fn seed_recent_thread(
test: &TestCodex,
title: &str,
@@ -636,13 +665,10 @@ async fn conversation_uses_experimental_realtime_ws_startup_context_override() -
let server = start_websocket_server(vec![
vec![],
vec![
vec![json!({
"type": "session.updated",
"session": { "id": "sess_custom_context", "instructions": "prompt from config" }
})],
vec![],
],
vec![vec![json!({
"type": "session.updated",
"session": { "id": "sess_custom_context", "instructions": "prompt from config" }
})]],
])
.await;
@@ -670,15 +696,7 @@ 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:?}"));
wait_for_session_updated_allowing_clean_close(&test.codex, "sess_custom_context").await;
let startup_context_request = server.wait_for_request(1, 0).await;
let instructions = websocket_request_instructions(&startup_context_request)
@@ -698,13 +716,10 @@ async fn conversation_disables_realtime_startup_context_with_empty_override() ->
let server = start_websocket_server(vec![
vec![],
vec![
vec![json!({
"type": "session.updated",
"session": { "id": "sess_no_context", "instructions": "prompt from config" }
})],
vec![],
],
vec![vec![json!({
"type": "session.updated",
"session": { "id": "sess_no_context", "instructions": "prompt from config" }
})]],
])
.await;
@@ -731,15 +746,7 @@ 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:?}"));
wait_for_session_updated_allowing_clean_close(&test.codex, "sess_no_context").await;
let startup_context_request = server.wait_for_request(1, 0).await;
let instructions = websocket_request_instructions(&startup_context_request)
@@ -759,13 +766,10 @@ async fn conversation_start_injects_startup_context_from_thread_history() -> Res
let server = start_websocket_server(vec![
vec![],
vec![
vec![json!({
"type": "session.updated",
"session": { "id": "sess_context", "instructions": "backend prompt" }
})],
vec![],
],
vec![vec![json!({
"type": "session.updated",
"session": { "id": "sess_context", "instructions": "backend prompt" }
})]],
])
.await;
@@ -788,15 +792,7 @@ 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:?}"));
wait_for_session_updated_allowing_clean_close(&test.codex, "sess_context").await;
let startup_context_request = server.wait_for_request(1, 0).await;
let startup_context = websocket_request_instructions(&startup_context_request)
@@ -823,13 +819,10 @@ async fn conversation_startup_context_falls_back_to_workspace_map() -> Result<()
let server = start_websocket_server(vec![
vec![],
vec![
vec![json!({
"type": "session.updated",
"session": { "id": "sess_workspace", "instructions": "backend prompt" }
})],
vec![],
],
vec![vec![json!({
"type": "session.updated",
"session": { "id": "sess_workspace", "instructions": "backend prompt" }
})]],
])
.await;
@@ -845,15 +838,7 @@ 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:?}"));
wait_for_session_updated_allowing_clean_close(&test.codex, "sess_workspace").await;
let startup_context_request = server.wait_for_request(1, 0).await;
let startup_context = websocket_request_instructions(&startup_context_request)