diff --git a/codex-rs/app-server-transport/Cargo.toml b/codex-rs/app-server-transport/Cargo.toml index 175890962..7ad944492 100644 --- a/codex-rs/app-server-transport/Cargo.toml +++ b/codex-rs/app-server-transport/Cargo.toml @@ -57,3 +57,4 @@ chrono = { workspace = true } codex-config = { workspace = true } pretty_assertions = { workspace = true } tempfile = { workspace = true } +tokio = { workspace = true, features = ["test-util"] } diff --git a/codex-rs/app-server-transport/src/transport/remote_control/client_tracker.rs b/codex-rs/app-server-transport/src/transport/remote_control/client_tracker.rs index 22e5e74ed..aef2bfd25 100644 --- a/codex-rs/app-server-transport/src/transport/remote_control/client_tracker.rs +++ b/codex-rs/app-server-transport/src/transport/remote_control/client_tracker.rs @@ -692,14 +692,14 @@ mod tests { } } - #[tokio::test] + #[tokio::test(start_paused = true)] async fn initialize_timeout_closes_open_connection() { let (server_event_tx, _server_event_rx) = mpsc::channel(CHANNEL_CAPACITY); let (transport_event_tx, mut transport_event_rx) = mpsc::channel(1); let shutdown_token = CancellationToken::new(); let client_tracker = ClientTracker::new(server_event_tx, transport_event_tx, &shutdown_token); - let mut handle_message = tokio::spawn(async move { + let handle_message = tokio::spawn(async move { let mut client_tracker = client_tracker; client_tracker .handle_message(initialize_envelope_with_stream_id( @@ -709,13 +709,13 @@ mod tests { .await }); - assert!( - timeout(Duration::from_millis(50), &mut handle_message) - .await - .expect("initialize timeout rollback should not wait for close delivery") - .expect("handle message task should not panic") - .is_err() - ); + tokio::task::yield_now().await; + tokio::time::advance( + REMOTE_CONTROL_TRANSPORT_EVENT_SEND_TIMEOUT + Duration::from_millis(1), + ) + .await; + + assert!(handle_message.await.expect("handle message task").is_err()); let connection_id = match transport_event_rx.recv().await.expect("open event") { TransportEvent::ConnectionOpened { connection_id, .. } => connection_id, other => panic!("expected connection opened, got {other:?}"),