From 484668e0730bca8f71c65c248ee666defac15cbd Mon Sep 17 00:00:00 2001 From: Ahmed Ibrahim Date: Fri, 6 Mar 2026 20:28:25 -0800 Subject: [PATCH] codex: fix realtime startup context close race (#13593) --- .../core/tests/suite/realtime_conversation.rs | 113 ++++++++---------- 1 file changed, 49 insertions(+), 64 deletions(-) diff --git a/codex-rs/core/tests/suite/realtime_conversation.rs b/codex-rs/core/tests/suite/realtime_conversation.rs index 74b0eb1d5f..8521feb751 100644 --- a/codex-rs/core/tests/suite/realtime_conversation.rs +++ b/codex-rs/core/tests/suite/realtime_conversation.rs @@ -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)