From 2354f861e7e0dcf6086a0cb65e656fd4f8f303f5 Mon Sep 17 00:00:00 2001 From: Ahmed Ibrahim Date: Wed, 4 Mar 2026 13:51:00 -0800 Subject: [PATCH] Fix flaky app-server/core tests and remove temporary logging Co-authored-by: Codex --- .../app-server/src/bespoke_event_handling.rs | 31 +++++++++++++++++-- codex-rs/app-server/src/thread_state.rs | 1 + .../app-server/tests/common/mcp_process.rs | 10 +----- .../app-server/tests/suite/v2/initialize.rs | 4 +-- .../tests/suite/v2/request_user_input.rs | 3 +- codex-rs/app-server/tests/suite/v2/review.rs | 5 +-- .../tests/suite/v2/turn_interrupt.rs | 21 +++++++++---- .../app-server/tests/suite/v2/turn_start.rs | 9 ++++-- .../tests/suite/v2/turn_start_zsh_fork.rs | 16 +++++++--- codex-rs/core/tests/suite/approvals.rs | 17 +++++----- .../core/tests/suite/shell_serialization.rs | 2 +- 11 files changed, 82 insertions(+), 37 deletions(-) diff --git a/codex-rs/app-server/src/bespoke_event_handling.rs b/codex-rs/app-server/src/bespoke_event_handling.rs index 4d005706db..55c96aa632 100644 --- a/codex-rs/app-server/src/bespoke_event_handling.rs +++ b/codex-rs/app-server/src/bespoke_event_handling.rs @@ -188,8 +188,9 @@ pub(crate) async fn apply_bespoke_event_handling( } = event; match msg { EventMsg::TurnStarted(payload) => { - // While not technically necessary as it was already done on TurnComplete, be extra cautios and abort any pending server requests. - outgoing.abort_pending_server_requests().await; + // Do not abort pending server requests here. In practice, approval requests can be + // emitted very close to turn-start handling and event ordering may vary by platform. + // Aborting on TurnStarted can race with those approval requests and drop callbacks. thread_watch_manager .note_turn_started(&conversation_id.to_string()) .await; @@ -213,6 +214,9 @@ pub(crate) async fn apply_bespoke_event_handling( } } EventMsg::TurnComplete(_ev) => { + if should_ignore_turn_complete(&event_turn_id, &thread_state).await { + return; + } // All per-thread requests are bound to a turn, so abort them. outgoing.abort_pending_server_requests().await; let turn_failed = thread_state.lock().await.turn_summary.last_error.is_some(); @@ -1423,6 +1427,7 @@ pub(crate) async fn apply_bespoke_event_handling( outgoing.abort_pending_server_requests().await; let pending = { let mut state = thread_state.lock().await; + state.interrupted_turn_ids.insert(event_turn_id.clone()); std::mem::take(&mut state.pending_interrupts) }; if !pending.is_empty() { @@ -1716,6 +1721,14 @@ async fn find_and_remove_turn_summary( std::mem::take(&mut state.turn_summary) } +async fn should_ignore_turn_complete( + event_turn_id: &str, + thread_state: &Arc>, +) -> bool { + let mut state = thread_state.lock().await; + state.interrupted_turn_ids.remove(event_turn_id) +} + async fn handle_turn_complete( conversation_id: ThreadId, event_turn_id: String, @@ -2493,6 +2506,20 @@ mod tests { Ok(()) } + #[tokio::test] + async fn test_should_ignore_turn_complete_for_interrupted_turn() -> Result<()> { + let thread_state = new_thread_state(); + let turn_id = "interrupt_then_complete".to_string(); + { + let mut state = thread_state.lock().await; + state.interrupted_turn_ids.insert(turn_id.clone()); + } + + assert!(should_ignore_turn_complete(&turn_id, &thread_state).await); + assert!(!should_ignore_turn_complete(&turn_id, &thread_state).await); + Ok(()) + } + #[tokio::test] async fn test_handle_turn_interrupted_emits_interrupted_with_error() -> Result<()> { let conversation_id = ThreadId::new(); diff --git a/codex-rs/app-server/src/thread_state.rs b/codex-rs/app-server/src/thread_state.rs index 915cc0fd77..d9fe8565b7 100644 --- a/codex-rs/app-server/src/thread_state.rs +++ b/codex-rs/app-server/src/thread_state.rs @@ -51,6 +51,7 @@ pub(crate) struct TurnSummary { #[derive(Default)] pub(crate) struct ThreadState { pub(crate) pending_interrupts: PendingInterruptQueue, + pub(crate) interrupted_turn_ids: HashSet, pub(crate) pending_rollbacks: Option, pub(crate) turn_summary: TurnSummary, pub(crate) cancel_tx: Option>, diff --git a/codex-rs/app-server/tests/common/mcp_process.rs b/codex-rs/app-server/tests/common/mcp_process.rs index 753816b8d2..1c7e74c0d5 100644 --- a/codex-rs/app-server/tests/common/mcp_process.rs +++ b/codex-rs/app-server/tests/common/mcp_process.rs @@ -101,7 +101,7 @@ impl McpProcess { cmd.stderr(Stdio::piped()); cmd.current_dir(codex_home); cmd.env("CODEX_HOME", codex_home); - cmd.env("RUST_LOG", "info"); + cmd.env("RUST_LOG", "warn"); cmd.env_remove(CODEX_INTERNAL_ORIGINATOR_OVERRIDE_ENV_VAR); for (k, v) in env_overrides { @@ -795,7 +795,6 @@ impl McpProcess { } async fn send_jsonrpc_message(&mut self, message: JSONRPCMessage) -> anyhow::Result<()> { - eprintln!("writing message to stdin: {message:?}"); let Some(stdin) = self.stdin.as_mut() else { anyhow::bail!("mcp stdin closed"); }; @@ -810,13 +809,10 @@ impl McpProcess { let mut line = String::new(); self.stdout.read_line(&mut line).await?; let message = serde_json::from_str::(&line)?; - eprintln!("read message from stdout: {message:?}"); Ok(message) } pub async fn read_stream_until_request_message(&mut self) -> anyhow::Result { - eprintln!("in read_stream_until_request_message()"); - let message = self .read_stream_until_message(|message| matches!(message, JSONRPCMessage::Request(_))) .await?; @@ -833,8 +829,6 @@ impl McpProcess { &mut self, request_id: RequestId, ) -> anyhow::Result { - eprintln!("in read_stream_until_response_message({request_id:?})"); - let message = self .read_stream_until_message(|message| { Self::message_request_id(message) == Some(&request_id) @@ -867,8 +861,6 @@ impl McpProcess { &mut self, method: &str, ) -> anyhow::Result { - eprintln!("in read_stream_until_notification_message({method})"); - let message = self .read_stream_until_message(|message| { matches!( diff --git a/codex-rs/app-server/tests/suite/v2/initialize.rs b/codex-rs/app-server/tests/suite/v2/initialize.rs index de437c5bf5..b1f7144148 100644 --- a/codex-rs/app-server/tests/suite/v2/initialize.rs +++ b/codex-rs/app-server/tests/suite/v2/initialize.rs @@ -18,11 +18,11 @@ use core_test_support::fs_wait; use pretty_assertions::assert_eq; use serde_json::Value; use std::path::Path; -use std::time::Duration; use tempfile::TempDir; use tokio::time::timeout; const DEFAULT_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10); +const NOTIFY_FILE_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(20); #[tokio::test] async fn initialize_uses_client_info_name_as_originator() -> Result<()> { @@ -261,7 +261,7 @@ tmp_path.replace(payload_path) ) .await??; - fs_wait::wait_for_path_exists(¬ify_file, Duration::from_secs(5)).await?; + fs_wait::wait_for_path_exists(¬ify_file, NOTIFY_FILE_TIMEOUT).await?; let payload_raw = tokio::fs::read_to_string(¬ify_file).await?; let payload: Value = serde_json::from_str(&payload_raw)?; assert_eq!(payload["client"], "xcode"); diff --git a/codex-rs/app-server/tests/suite/v2/request_user_input.rs b/codex-rs/app-server/tests/suite/v2/request_user_input.rs index f77ddfb4f7..83de2fbcb0 100644 --- a/codex-rs/app-server/tests/suite/v2/request_user_input.rs +++ b/codex-rs/app-server/tests/suite/v2/request_user_input.rs @@ -21,6 +21,7 @@ use codex_protocol::openai_models::ReasoningEffort; use tokio::time::timeout; const DEFAULT_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10); +const STARTUP_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30); #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn request_user_input_round_trip() -> Result<()> { @@ -33,7 +34,7 @@ async fn request_user_input_round_trip() -> Result<()> { create_config_toml(codex_home.path(), &server.uri())?; let mut mcp = McpProcess::new(codex_home.path()).await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + timeout(STARTUP_TIMEOUT, mcp.initialize()).await??; let thread_start_id = mcp .send_thread_start_request(ThreadStartParams { diff --git a/codex-rs/app-server/tests/suite/v2/review.rs b/codex-rs/app-server/tests/suite/v2/review.rs index 6febbaf52a..1620152d40 100644 --- a/codex-rs/app-server/tests/suite/v2/review.rs +++ b/codex-rs/app-server/tests/suite/v2/review.rs @@ -31,6 +31,7 @@ use tempfile::TempDir; use tokio::time::timeout; const DEFAULT_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10); +const STARTUP_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30); const INVALID_REQUEST_ERROR_CODE: i64 = -32600; #[tokio::test] @@ -232,7 +233,7 @@ async fn review_start_rejects_empty_base_branch() -> Result<()> { create_config_toml(codex_home.path(), &server.uri())?; let mut mcp = McpProcess::new(codex_home.path()).await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + timeout(STARTUP_TIMEOUT, mcp.initialize()).await??; let thread_id = start_default_thread(&mut mcp).await?; let request_id = mcp @@ -340,7 +341,7 @@ async fn review_start_rejects_empty_commit_sha() -> Result<()> { create_config_toml(codex_home.path(), &server.uri())?; let mut mcp = McpProcess::new(codex_home.path()).await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + timeout(STARTUP_TIMEOUT, mcp.initialize()).await??; let thread_id = start_default_thread(&mut mcp).await?; let request_id = mcp diff --git a/codex-rs/app-server/tests/suite/v2/turn_interrupt.rs b/codex-rs/app-server/tests/suite/v2/turn_interrupt.rs index 8f944dca07..38125e1804 100644 --- a/codex-rs/app-server/tests/suite/v2/turn_interrupt.rs +++ b/codex-rs/app-server/tests/suite/v2/turn_interrupt.rs @@ -3,6 +3,7 @@ use anyhow::Result; use app_test_support::McpProcess; use app_test_support::create_mock_responses_server_sequence; +use app_test_support::create_mock_responses_server_sequence_unchecked; use app_test_support::create_shell_command_sse_response; use app_test_support::to_response; use codex_app_server_protocol::JSONRPCNotification; @@ -19,6 +20,7 @@ use codex_app_server_protocol::TurnStartParams; use codex_app_server_protocol::TurnStartResponse; use codex_app_server_protocol::TurnStatus; use codex_app_server_protocol::UserInput as V2UserInput; +use core_test_support::responses; use tempfile::TempDir; use tokio::time::timeout; @@ -139,12 +141,19 @@ async fn turn_interrupt_resolves_pending_command_approval_request() -> Result<() let working_directory = tmp.path().join("workdir"); std::fs::create_dir(&working_directory)?; - let server = create_mock_responses_server_sequence(vec![create_shell_command_sse_response( - shell_command.clone(), - Some(&working_directory), - Some(10_000), - "call_sleep_approval", - )?]) + let no_op_response = responses::sse(vec![ + responses::ev_response_created("resp-2"), + responses::ev_completed("resp-2"), + ]); + let server = create_mock_responses_server_sequence_unchecked(vec![ + create_shell_command_sse_response( + shell_command.clone(), + Some(&working_directory), + Some(10_000), + "call_sleep_approval", + )?, + no_op_response, + ]) .await; create_config_toml(&codex_home, &server.uri(), "untrusted")?; diff --git a/codex-rs/app-server/tests/suite/v2/turn_start.rs b/codex-rs/app-server/tests/suite/v2/turn_start.rs index 15b309ddf4..6897f6ca74 100644 --- a/codex-rs/app-server/tests/suite/v2/turn_start.rs +++ b/codex-rs/app-server/tests/suite/v2/turn_start.rs @@ -1423,11 +1423,13 @@ async fn turn_start_updates_sandbox_and_cwd_between_turns_v2() -> Result<()> { collaboration_mode: None, }) .await?; - timeout( + let second_turn_resp: JSONRPCResponse = timeout( DEFAULT_READ_TIMEOUT, mcp.read_stream_until_response_message(RequestId::Integer(second_turn)), ) .await??; + let TurnStartResponse { turn, .. } = to_response::(second_turn_resp)?; + let second_turn_id = turn.id; let command_exec_item = timeout(DEFAULT_READ_TIMEOUT, async { loop { @@ -1440,7 +1442,10 @@ async fn turn_start_updates_sandbox_and_cwd_between_turns_v2() -> Result<()> { .expect("item/started params"); let item_started: ItemStartedNotification = serde_json::from_value(params).expect("deserialize item/started notification"); - if matches!(item_started.item, ThreadItem::CommandExecution { .. }) { + if item_started.turn_id == second_turn_id + && let ThreadItem::CommandExecution { id, .. } = &item_started.item + && id == "call-second" + { return Ok::(item_started.item); } } diff --git a/codex-rs/app-server/tests/suite/v2/turn_start_zsh_fork.rs b/codex-rs/app-server/tests/suite/v2/turn_start_zsh_fork.rs index bdbb283432..e6225abda4 100644 --- a/codex-rs/app-server/tests/suite/v2/turn_start_zsh_fork.rs +++ b/codex-rs/app-server/tests/suite/v2/turn_start_zsh_fork.rs @@ -547,11 +547,17 @@ async fn turn_start_shell_zsh_fork_subcommand_decline_marks_parent_declined_v2() let second_file_str = second_file.to_string_lossy().into_owned(); let parent_shell_hint = format!("&& {}", &first_file_str); while target_decision_index < target_decisions.len() || !saw_parent_approval { - let server_req = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_request_message(), - ) - .await??; + let wait_timeout = if target_decision_index < target_decisions.len() { + DEFAULT_READ_TIMEOUT + } else { + std::time::Duration::from_secs(2) + }; + let server_req = match timeout(wait_timeout, mcp.read_stream_until_request_message()).await + { + Ok(server_req) => server_req?, + Err(_) if target_decision_index >= target_decisions.len() => break, + Err(error) => return Err(error.into()), + }; let ServerRequest::CommandExecutionRequestApproval { request_id, params } = server_req else { panic!("expected CommandExecutionRequestApproval request"); diff --git a/codex-rs/core/tests/suite/approvals.rs b/codex-rs/core/tests/suite/approvals.rs index 9ca159965f..8be2cda365 100644 --- a/codex-rs/core/tests/suite/approvals.rs +++ b/codex-rs/core/tests/suite/approvals.rs @@ -653,14 +653,17 @@ async fn expect_patch_approval( } async fn wait_for_completion_without_approval(test: &TestCodex) { - let event = wait_for_event(&test.codex, |event| { - matches!( - event, - EventMsg::ExecApprovalRequest(_) | EventMsg::TurnComplete(_) - ) - }) + let event = wait_for_event_with_timeout( + &test.codex, + |event| { + matches!( + event, + EventMsg::ExecApprovalRequest(_) | EventMsg::TurnComplete(_) + ) + }, + std::time::Duration::from_secs(10), + ) .await; - match event { EventMsg::TurnComplete(_) => {} EventMsg::ExecApprovalRequest(event) => { diff --git a/codex-rs/core/tests/suite/shell_serialization.rs b/codex-rs/core/tests/suite/shell_serialization.rs index 3e5942da6a..018807761c 100644 --- a/codex-rs/core/tests/suite/shell_serialization.rs +++ b/codex-rs/core/tests/suite/shell_serialization.rs @@ -841,7 +841,7 @@ async fn shell_command_output_is_not_truncated_over_10k_bytes() -> Result<()> { let call_id = "shell-command"; let args = json!({ "command": "perl -e 'print \"1\" x 10001'", - "timeout_ms": 1000, + "timeout_ms": 5000, }); let responses = vec![ sse(vec![