diff --git a/codex-rs/app-server/tests/suite/v2/review.rs b/codex-rs/app-server/tests/suite/v2/review.rs index adbc26ece9..c15274644c 100644 --- a/codex-rs/app-server/tests/suite/v2/review.rs +++ b/codex-rs/app-server/tests/suite/v2/review.rs @@ -140,7 +140,6 @@ async fn review_start_runs_review_turn_and_emits_code_review_item() -> Result<() } #[tokio::test] -#[ignore = "TODO(owenlin0): flaky"] async fn review_start_exec_approval_item_id_matches_command_execution_item() -> Result<()> { let responses = vec![ create_shell_command_sse_response( @@ -183,33 +182,44 @@ async fn review_start_exec_approval_item_id_matches_command_execution_item() -> let ReviewStartResponse { turn, .. } = to_response::(review_resp)?; let turn_id = turn.id.clone(); - let server_req = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_request_message(), - ) - .await??; - let ServerRequest::CommandExecutionRequestApproval { request_id, params } = server_req else { - panic!("expected CommandExecutionRequestApproval request"); - }; - assert_eq!(params.item_id, "review-call-1"); - assert_eq!(params.turn_id, turn_id); - - let mut command_item_id = None; - for _ in 0..10 { - let item_started: JSONRPCNotification = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_notification_message("item/started"), - ) - .await??; - let started: ItemStartedNotification = - serde_json::from_value(item_started.params.expect("params must be present"))?; - if let ThreadItem::CommandExecution { id, .. } = started.item { - command_item_id = Some(id); - break; + let deadline = tokio::time::Instant::now() + DEFAULT_READ_TIMEOUT; + let mut approval_request: Option<(RequestId, String)> = None; + let mut command_item_id: Option = None; + while approval_request.is_none() || command_item_id.is_none() { + let remaining = deadline.saturating_duration_since(tokio::time::Instant::now()); + let message = timeout(remaining, mcp.read_next_message()).await??; + match message { + JSONRPCMessage::Request(request) => { + let server_req: ServerRequest = request.try_into()?; + let ServerRequest::CommandExecutionRequestApproval { request_id, params } = + server_req + else { + continue; + }; + assert_eq!(params.item_id, "review-call-1"); + assert_eq!(params.turn_id, turn_id); + approval_request = Some((request_id, params.item_id)); + } + JSONRPCMessage::Notification(notification) if notification.method == "item/started" => { + let started: ItemStartedNotification = + serde_json::from_value(notification.params.expect("params must be present"))?; + if started.turn_id != turn_id { + continue; + } + if let ThreadItem::CommandExecution { id, .. } = started.item { + command_item_id = Some(id); + } + } + _ => continue, } } - let command_item_id = command_item_id.expect("did not observe command execution item"); - assert_eq!(command_item_id, params.item_id); + + let (request_id, approval_item_id) = + approval_request.expect("did not observe command execution approval request"); + assert_eq!( + command_item_id.expect("did not observe command execution item"), + approval_item_id + ); mcp.send_response( request_id,