From 88165dcd0786c8446837eed9ea8cb5d19a4dd3af Mon Sep 17 00:00:00 2001 From: Ahmed Ibrahim Date: Tue, 9 Jun 2026 18:32:26 +0000 Subject: [PATCH] preserve deferred initialization boundaries --- codex-rs/codex-mcp/src/connection_manager.rs | 4 +- codex-rs/core/src/session/deferred_init.rs | 16 +++- codex-rs/core/src/session/session.rs | 2 + codex-rs/core/src/session/tests.rs | 2 + codex-rs/core/src/session/turn_context.rs | 4 +- codex-rs/core/tests/suite/agent_websocket.rs | 73 ++++++------------- .../core/tests/suite/client_websockets.rs | 18 +---- .../core/tests/suite/websocket_fallback.rs | 14 ++-- 8 files changed, 53 insertions(+), 80 deletions(-) diff --git a/codex-rs/codex-mcp/src/connection_manager.rs b/codex-rs/codex-mcp/src/connection_manager.rs index 28b92090e3..163853e8ea 100644 --- a/codex-rs/codex-mcp/src/connection_manager.rs +++ b/codex-rs/codex-mcp/src/connection_manager.rs @@ -265,7 +265,9 @@ impl McpConnectionManager { ) .await; if let Some(event) = empty_startup_event { - let _ = tx_event.send(event).await; + tokio::spawn(async move { + let _ = tx_event.send(event).await; + }); } (manager, cancel_token) } diff --git a/codex-rs/core/src/session/deferred_init.rs b/codex-rs/core/src/session/deferred_init.rs index 73bded8c46..55f1b43aea 100644 --- a/codex-rs/core/src/session/deferred_init.rs +++ b/codex-rs/core/src/session/deferred_init.rs @@ -94,7 +94,6 @@ impl Session { self.deferred_initialization_ready .get_or_init(|| async { self.ensure_model_catalog_ready().await; - self.initialize_plugins_and_skills().await; let initial_history = self.initial_history.lock().await.take(); if let Some(initial_history) = initial_history { self.record_initial_history(initial_history).await; @@ -103,6 +102,19 @@ impl Session { .await; } + async fn ensure_plugin_initialization_ready(&self) { + self.plugin_initialization_ready + .get_or_init(|| self.initialize_plugins_and_skills()) + .await; + } + + pub(crate) async fn ensure_turn_initialization_ready(&self) { + tokio::join!( + self.ensure_deferred_initialization_ready(), + self.ensure_plugin_initialization_ready(), + ); + } + #[expect( clippy::await_holding_invalid_type, reason = "the session-owned MCP manager must remain exclusively locked while adding servers" @@ -199,7 +211,7 @@ impl Session { reason = "the prewarm gate must stay locked until the handle is installed" )] pub(crate) async fn finish_deferred_initialization(self: &Arc) { - self.ensure_deferred_initialization_ready().await; + self.ensure_turn_initialization_ready().await; let mut prewarm_pending = self.model_catalog_prewarm_pending.lock().await; if self.shutdown_requested.load(Ordering::SeqCst) || !*prewarm_pending { return; diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 1c8d776a03..3948ba720e 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -46,6 +46,7 @@ pub(crate) struct Session { pub(crate) late_mcp_tools: LateToolRegistry, pub(super) model_catalog_ready: OnceCell<()>, pub(super) deferred_initialization_ready: OnceCell<()>, + pub(super) plugin_initialization_ready: OnceCell<()>, pub(super) initial_history: Mutex>, pub(super) model_catalog_refresh_enabled: bool, pub(super) initial_default_model: Option, @@ -1036,6 +1037,7 @@ impl Session { late_mcp_tools: LateToolRegistry::default(), model_catalog_ready: OnceCell::new(), deferred_initialization_ready: OnceCell::new(), + plugin_initialization_ready: OnceCell::new(), initial_history: Mutex::new(Some(initial_history.clone())), model_catalog_refresh_enabled, initial_default_model, diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 6e40a2b0d2..eda9e69f69 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -4927,6 +4927,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { late_mcp_tools: crate::tools::registry::LateToolRegistry::default(), model_catalog_ready: tokio::sync::OnceCell::new(), deferred_initialization_ready: tokio::sync::OnceCell::new(), + plugin_initialization_ready: tokio::sync::OnceCell::new(), initial_history: Mutex::new(None), model_catalog_refresh_enabled: false, initial_default_model: None, @@ -7004,6 +7005,7 @@ where late_mcp_tools: crate::tools::registry::LateToolRegistry::default(), model_catalog_ready: tokio::sync::OnceCell::new(), deferred_initialization_ready: tokio::sync::OnceCell::new(), + plugin_initialization_ready: tokio::sync::OnceCell::new(), initial_history: Mutex::new(None), model_catalog_refresh_enabled: false, initial_default_model: None, diff --git a/codex-rs/core/src/session/turn_context.rs b/codex-rs/core/src/session/turn_context.rs index bfae0d4c94..37f638d3ab 100644 --- a/codex-rs/core/src/session/turn_context.rs +++ b/codex-rs/core/src/session/turn_context.rs @@ -583,7 +583,7 @@ impl Session { updates: SessionSettingsUpdate, ) -> CodexResult> { self.claim_model_catalog_prewarm().await; - self.ensure_deferred_initialization_ready().await; + self.ensure_turn_initialization_ready().await; let notify_config_contributors = !self.services.extensions.config_contributors().is_empty(); let update_result: CodexResult<_> = { let mut state = self.state.lock().await; @@ -837,7 +837,7 @@ impl Session { pub(crate) async fn new_inference_turn_with_sub_id(&self, sub_id: String) -> Arc { self.claim_model_catalog_prewarm().await; - self.ensure_deferred_initialization_ready().await; + self.ensure_turn_initialization_ready().await; self.new_default_turn_with_sub_id(sub_id).await } diff --git a/codex-rs/core/tests/suite/agent_websocket.rs b/codex-rs/core/tests/suite/agent_websocket.rs index c405e35995..fc1d5c39e3 100644 --- a/codex-rs/core/tests/suite/agent_websocket.rs +++ b/codex-rs/core/tests/suite/agent_websocket.rs @@ -67,17 +67,14 @@ async fn websocket_test_codex_shell_chain() -> Result<()> { } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn websocket_first_turn_uses_startup_prewarm_and_create() -> Result<()> { +async fn websocket_immediate_first_turn_skips_unfinished_startup_prewarm() -> Result<()> { skip_if_no_network!(Ok(())); - let server = start_websocket_server(vec![vec![ - vec![ev_response_created("warm-1"), ev_completed("warm-1")], - vec![ - ev_response_created("resp-1"), - ev_assistant_message("msg-1", "hello"), - ev_completed("resp-1"), - ], - ]]) + let server = start_websocket_server(vec![vec![vec![ + ev_response_created("resp-1"), + ev_assistant_message("msg-1", "hello"), + ev_completed("resp-1"), + ]]]) .await; let mut builder = test_codex(); @@ -87,24 +84,11 @@ async fn websocket_first_turn_uses_startup_prewarm_and_create() -> Result<()> { assert_eq!(server.handshakes().len(), 1); let connection = server.single_connection(); - assert_eq!(connection.len(), 2); - let warmup = connection + assert_eq!(connection.len(), 1); + let turn = connection .first() - .expect("missing warmup request") + .expect("missing turn request") .body_json(); - let turn = connection.get(1).expect("missing turn request").body_json(); - assert_eq!(warmup["type"].as_str(), Some("response.create")); - assert_eq!(warmup["generate"].as_bool(), Some(false)); - let warmup_metadata: Value = serde_json::from_str( - warmup["client_metadata"]["x-codex-turn-metadata"] - .as_str() - .expect("warmup turn metadata"), - )?; - assert_eq!(warmup_metadata["request_kind"].as_str(), Some("prewarm")); - assert_eq!( - warmup_metadata["window_id"].as_str(), - warmup["client_metadata"]["x-codex-window-id"].as_str() - ); assert!( turn["tools"] .as_array() @@ -124,18 +108,15 @@ async fn websocket_first_turn_uses_startup_prewarm_and_create() -> Result<()> { } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn websocket_first_turn_handles_handshake_delay_with_startup_prewarm() -> Result<()> { +async fn websocket_immediate_first_turn_handles_handshake_delay_without_prewarm() -> Result<()> { skip_if_no_network!(Ok(())); let server = start_websocket_server_with_headers(vec![WebSocketConnectionConfig { - requests: vec![ - vec![ev_response_created("warm-1"), ev_completed("warm-1")], - vec![ - ev_response_created("resp-1"), - ev_assistant_message("msg-1", "hello"), - ev_completed("resp-1"), - ], - ], + requests: vec![vec![ + ev_response_created("resp-1"), + ev_assistant_message("msg-1", "hello"), + ev_completed("resp-1"), + ]], response_headers: Vec::new(), // Delay handshake so turn processing must tolerate websocket startup latency. accept_delay: Some(Duration::from_millis(150)), @@ -150,14 +131,11 @@ async fn websocket_first_turn_handles_handshake_delay_with_startup_prewarm() -> assert_eq!(server.handshakes().len(), 1); let connection = server.single_connection(); - assert_eq!(connection.len(), 2); - let warmup = connection + assert_eq!(connection.len(), 1); + let turn = connection .first() - .expect("missing warmup request") + .expect("missing turn request") .body_json(); - let turn = connection.get(1).expect("missing turn request").body_json(); - assert_eq!(warmup["type"].as_str(), Some("response.create")); - assert_eq!(warmup["generate"].as_bool(), Some(false)); assert!( turn["tools"] .as_array() @@ -176,7 +154,6 @@ async fn websocket_v2_test_codex_shell_chain() -> Result<()> { let call_id = "shell-command-call"; let server = start_websocket_server(vec![vec![ - vec![ev_response_created("warm-1"), ev_completed("warm-1")], vec![ ev_response_created("resp-1"), ev_shell_command_call(call_id, "echo websocket"), @@ -202,25 +179,19 @@ async fn websocket_v2_test_codex_shell_chain() -> Result<()> { .await?; let connection = server.single_connection(); - assert_eq!(connection.len(), 3); + assert_eq!(connection.len(), 2); - let warmup = connection - .first() - .expect("missing warmup request") - .body_json(); let first_turn = connection - .get(1) + .first() .expect("missing first turn request") .body_json(); let second_turn = connection - .get(2) + .get(1) .expect("missing second turn request") .body_json(); - assert_eq!(warmup["type"].as_str(), Some("response.create")); - assert_eq!(warmup["generate"].as_bool(), Some(false)); assert_eq!(first_turn["type"].as_str(), Some("response.create")); - assert_eq!(first_turn["previous_response_id"].as_str(), Some("warm-1")); + assert_eq!(first_turn.get("previous_response_id"), None); assert!( first_turn .get("input") diff --git a/codex-rs/core/tests/suite/client_websockets.rs b/codex-rs/core/tests/suite/client_websockets.rs index 4dc674ec30..b791c82dcf 100755 --- a/codex-rs/core/tests/suite/client_websockets.rs +++ b/codex-rs/core/tests/suite/client_websockets.rs @@ -1211,14 +1211,7 @@ async fn responses_websocket_usage_limit_error_emits_rate_limit_event() { } }); - let server = start_websocket_server(vec![vec![ - vec![ - ev_response_created("resp-prewarm"), - ev_completed("resp-prewarm"), - ], - vec![usage_limit_error], - ]]) - .await; + let server = start_websocket_server(vec![vec![vec![usage_limit_error]]]).await; let mut builder = test_codex().with_config(|config| { config.model_provider.request_max_retries = Some(0); config.model_provider.stream_max_retries = Some(0); @@ -1301,14 +1294,7 @@ async fn responses_websocket_invalid_request_error_with_status_is_forwarded() { } }); - let server = start_websocket_server(vec![vec![ - vec![ - ev_response_created("resp-prewarm"), - ev_completed("resp-prewarm"), - ], - vec![invalid_request_error], - ]]) - .await; + let server = start_websocket_server(vec![vec![vec![invalid_request_error]]]).await; let mut builder = test_codex().with_config(|config| { config.model_provider.request_max_retries = Some(0); config.model_provider.stream_max_retries = Some(0); diff --git a/codex-rs/core/tests/suite/websocket_fallback.rs b/codex-rs/core/tests/suite/websocket_fallback.rs index cdf77694c2..5e0729040a 100644 --- a/codex-rs/core/tests/suite/websocket_fallback.rs +++ b/codex-rs/core/tests/suite/websocket_fallback.rs @@ -113,10 +113,9 @@ async fn websocket_fallback_switches_to_http_after_retries_exhausted() -> Result .filter(|req| req.method == Method::POST && req.url.path().ends_with("/responses")) .count(); - // Deferred request prewarm is attempted at startup. - // The first turn then makes 3 websocket stream attempts (initial try + 2 retries), - // after which fallback activates and the request is replayed over HTTP. - assert_eq!(websocket_attempts, 4); + // The immediate first turn claims the deferred prewarm, then makes 3 websocket stream + // attempts (initial try + 2 retries) before replaying the request over HTTP. + assert_eq!(websocket_attempts, 3); assert_eq!(http_attempts, 1); assert_eq!(response_mock.requests().len(), 1); @@ -244,11 +243,10 @@ async fn websocket_fallback_is_sticky_across_turns() -> Result<()> { .filter(|req| req.method == Method::POST && req.url.path().ends_with("/responses")) .count(); - // WebSocket attempts all happen on the first turn: - // 1 deferred request prewarm attempt (startup) + 3 stream attempts - // (initial try + 2 retries) before fallback. + // All 3 websocket attempts happen on the immediate first turn (initial try + 2 retries) + // before fallback; the unfinished deferred prewarm is claimed without a request. // Fallback is sticky, so the second turn stays on HTTP and adds no websocket attempts. - assert_eq!(websocket_attempts, 4); + assert_eq!(websocket_attempts, 3); assert_eq!(http_attempts, 2); assert_eq!(response_mock.requests().len(), 2);