openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-11-12 21:20:34 +00:00
Code Issues Packages Projects Releases Wiki Activity

preserve deferred initialization boundaries

Browse Source
This commit is contained in:
Ahmed Ibrahim
2026-06-09 18:32:26 +00:00
parent 73cf5d1d18
commit 88165dcd07
8 changed files with 53 additions and 80 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/codex-mcp/src/connection_manager.rs
+3 -1
View File
@@ -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)
}
codex-rs/core/src/session/deferred_init.rs
+14 -2
View File
@@ -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>) {
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;
codex-rs/core/src/session/session.rs
+2
View File
@@ -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<Option<InitialHistory>>,
pub(super) model_catalog_refresh_enabled: bool,
pub(super) initial_default_model: Option<String>,
@@ -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,
codex-rs/core/src/session/tests.rs
+2
View File
@@ -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,
codex-rs/core/src/session/turn_context.rs
+2 -2
View File
@@ -583,7 +583,7 @@ impl Session {
updates: SessionSettingsUpdate,
) -> CodexResult<Arc<TurnContext>> {
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<TurnContext> {
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
}
codex-rs/core/tests/suite/agent_websocket.rs
+22 -51
View File
@@ -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")
codex-rs/core/tests/suite/client_websockets.rs
+2 -16
View File
@@ -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);
codex-rs/core/tests/suite/websocket_fallback.rs
+6 -8
View File
@@ -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);