From 87b808bb570f01f4b6fc8485c5459052fac0e320 Mon Sep 17 00:00:00 2001 From: Anton Panasenko Date: Fri, 5 Jun 2026 22:54:57 -0700 Subject: [PATCH] fix(remote-control): preserve enrollment on generic websocket 404s (#26741) ## Why A remote-control WebSocket handshake can receive a generic HTTP 404 when an intermediary routes the request without preserving the WebSocket upgrade. Treating every 404 as proof that the remote app server is gone clears valid enrollment and causes repeated re-enrollment, new environment and server IDs, Habitat churn, and noisy `/server/enroll` traffic. ## What Changed - Clear enrollment only when a 404 JSON response explicitly contains `{"detail":"Remote app server not found"}`. - Preserve enrollment for empty, plain-text, malformed, or otherwise unrecognized 404 responses, return the transport error, and retry with the existing reconnect backoff. - Log the status, correlation headers (`request-id` or `x-oai-request-id`, plus `cf-ray`), and bounded/redacted response body for unrecognized 404s. - Cover both explicit missing-server re-enrollment and generic 404 enrollment preservation/reconnect behavior. ## Verification `just test -p codex-app-server-transport` passes all 114 tests on the rebased branch, including the targeted explicit and generic WebSocket 404 scenarios. Related issue: N/A --- .../src/transport/remote_control/tests.rs | 139 +++++++++++++++++- .../src/transport/remote_control/websocket.rs | 74 +++++++++- 2 files changed, 210 insertions(+), 3 deletions(-) diff --git a/codex-rs/app-server-transport/src/transport/remote_control/tests.rs b/codex-rs/app-server-transport/src/transport/remote_control/tests.rs index 69659e84a..1e653c296 100644 --- a/codex-rs/app-server-transport/src/transport/remote_control/tests.rs +++ b/codex-rs/app-server-transport/src/transport/remote_control/tests.rs @@ -1744,7 +1744,7 @@ async fn remote_control_http_mode_reenrolls_when_refresh_reports_stale_enrollmen } #[tokio::test] -async fn remote_control_http_mode_clears_stale_persisted_enrollment_after_404() { +async fn remote_control_http_mode_reenrolls_after_explicit_missing_server_404() { let listener = TcpListener::bind("127.0.0.1:0") .await .expect("listener should bind"); @@ -1831,7 +1831,12 @@ async fn remote_control_http_mode_clears_stale_persisted_enrollment_after_404() Some("env_stale"), ) .await; - respond_with_status(websocket_request.stream, "404 Not Found", "").await; + respond_with_status( + websocket_request.stream, + "404 Not Found", + &json!({"detail": "Remote app server not found"}).to_string(), + ) + .await; expect_remote_control_status( &mut status_rx, /*expected_status*/ None, @@ -1881,6 +1886,136 @@ async fn remote_control_http_mode_clears_stale_persisted_enrollment_after_404() let _ = remote_task.await; } +#[tokio::test] +async fn remote_control_http_mode_preserves_enrollment_after_generic_websocket_404() { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("listener should bind"); + let remote_control_url = remote_control_url_for_listener(&listener); + let codex_home = TempDir::new().expect("temp dir should create"); + let state_db = remote_control_state_runtime(&codex_home).await; + let remote_control_target = + normalize_remote_control_url(&remote_control_url).expect("target should parse"); + let stale_enrollment = RemoteControlEnrollment { + remote_control_target: remote_control_target.clone(), + account_id: "account_id".to_string(), + environment_id: "env_stale".to_string(), + server_id: "srv_e_stale".to_string(), + server_name: "stale-server".to_string(), + remote_control_token: None, + expires_at: None, + }; + update_persisted_remote_control_enrollment( + Some(state_db.as_ref()), + &remote_control_target, + "account_id", + /*app_server_client_name*/ None, + Some(&stale_enrollment), + ) + .await + .expect("stale enrollment should save"); + + let (transport_event_tx, _transport_event_rx) = + mpsc::channel::(CHANNEL_CAPACITY); + let shutdown_token = CancellationToken::new(); + let (remote_task, remote_handle) = start_remote_control( + RemoteControlStartConfig { + remote_control_url, + installation_id: TEST_INSTALLATION_ID.to_string(), + }, + Some(state_db.clone()), + remote_control_auth_manager_with_home(&codex_home), + transport_event_tx, + shutdown_token.clone(), + /*app_server_client_name_rx*/ None, + /*initial_enabled*/ true, + ) + .await + .expect("remote control should start"); + let mut status_rx = remote_handle.status_receiver(); + + let refresh_request = accept_http_request(&listener).await; + assert_eq!( + refresh_request.request_line, + "POST /backend-api/wham/remote/control/server/refresh HTTP/1.1" + ); + respond_with_json( + refresh_request.stream, + remote_control_server_token_response( + &stale_enrollment.server_id, + &stale_enrollment.environment_id, + TEST_REFRESHED_REMOTE_CONTROL_SERVER_TOKEN, + ), + ) + .await; + + let websocket_request = accept_http_request(&listener).await; + assert_eq!( + websocket_request.request_line, + "GET /backend-api/wham/remote/control/server HTTP/1.1" + ); + assert_eq!( + websocket_request.headers.get("x-codex-server-id"), + Some(&stale_enrollment.server_id) + ); + assert_eq!( + websocket_request.headers.get("authorization"), + Some(&format!( + "Bearer {TEST_REFRESHED_REMOTE_CONTROL_SERVER_TOKEN}" + )) + ); + expect_remote_control_status( + &mut status_rx, + /*expected_status*/ None, + Some("env_stale"), + ) + .await; + respond_with_status_and_headers( + websocket_request.stream, + "404 Not Found", + &[("x-request-id", "request-404"), ("cf-ray", "ray-404")], + "Not Found", + ) + .await; + + assert_eq!( + load_persisted_remote_control_enrollment( + Some(state_db.as_ref()), + &remote_control_target, + "account_id", + /*app_server_client_name*/ None, + ) + .await + .expect("stale enrollment should load"), + Some(stale_enrollment.clone()) + ); + + let (handshake_request, _websocket) = accept_remote_control_backend_connection(&listener).await; + assert_eq!( + handshake_request.headers.get("x-codex-server-id"), + Some(&stale_enrollment.server_id) + ); + assert_eq!( + handshake_request.headers.get("authorization"), + Some(&format!( + "Bearer {TEST_REFRESHED_REMOTE_CONTROL_SERVER_TOKEN}" + )) + ); + expect_remote_control_status_snapshot( + &mut status_rx, + RemoteControlStatusChangedNotification { + status: RemoteControlConnectionStatus::Connected, + server_name: test_server_name(), + installation_id: TEST_INSTALLATION_ID.to_string(), + environment_id: Some("env_stale".to_string()), + }, + ) + .await; + + shutdown_token.cancel(); + let _ = remote_task.await; +} + #[derive(Debug)] struct CapturedHttpRequest { stream: TcpStream, diff --git a/codex-rs/app-server-transport/src/transport/remote_control/websocket.rs b/codex-rs/app-server-transport/src/transport/remote_control/websocket.rs index 86b72849d..e0208b1f1 100644 --- a/codex-rs/app-server-transport/src/transport/remote_control/websocket.rs +++ b/codex-rs/app-server-transport/src/transport/remote_control/websocket.rs @@ -77,6 +77,7 @@ const REMOTE_CONTROL_WEBSOCKET_CONNECT_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30); const REMOTE_CONTROL_CONNECTION_SHUTDOWN_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(5); +const REMOTE_APP_SERVER_NOT_FOUND_DETAIL: &str = "Remote app server not found"; struct BoundedOutboundBuffer { buffer_by_stream: HashMap<(ClientId, StreamId), VecDeque>, @@ -1248,7 +1249,9 @@ pub(super) async fn connect_remote_control_websocket( Ok((websocket_stream, response)) => Ok((websocket_stream, response.map(|_| ()))), Err(err) => { match &err { - tungstenite::Error::Http(response) if response.status().as_u16() == 404 => { + tungstenite::Error::Http(response) + if websocket_response_reports_missing_remote_app_server(response) => + { info!( "remote control websocket returned HTTP 404; clearing stale enrollment before re-enrolling: websocket_url={}, account_id={}, server_id={}, environment_id={}", remote_control_target.websocket_url, @@ -1267,6 +1270,23 @@ pub(super) async fn connect_remote_control_websocket( ) .await; } + tungstenite::Error::Http(response) if response.status().as_u16() == 404 => { + let response_body = response + .body() + .as_deref() + .map(preview_remote_control_response_body) + .unwrap_or_else(|| "".to_string()); + warn!( + websocket_url = %remote_control_target.websocket_url, + account_id = %auth.account_id, + server_id = %enrollment.server_id, + environment_id = %enrollment.environment_id, + response_status = %response.status(), + response_headers = %format_headers(response.headers()), + response_body = %response_body, + "remote control websocket returned unrecognized HTTP 404; preserving enrollment before retry" + ); + } tungstenite::Error::Http(response) if matches!(response.status().as_u16(), 401 | 403) => { @@ -1434,6 +1454,18 @@ async fn prepare_remote_control_enrollment( Ok(auth) } +fn websocket_response_reports_missing_remote_app_server( + response: &tungstenite::http::Response>>, +) -> bool { + response.status().as_u16() == 404 + && response.body().as_deref().is_some_and(|body| { + serde_json::from_slice::(body).is_ok_and(|body| { + body.get("detail").and_then(serde_json::Value::as_str) + == Some(REMOTE_APP_SERVER_NOT_FOUND_DETAIL) + }) + }) +} + async fn clear_remote_control_enrollment( state_db: &StateRuntime, remote_control_target: &RemoteControlTarget, @@ -1673,6 +1705,46 @@ mod tests { assert_eq!(reconnect_attempt, 1); } + #[test] + fn websocket_404_only_reports_explicit_missing_remote_app_server() { + let cases = [ + ( + Some(br#"{"detail":"Remote app server not found"}"#.to_vec()), + true, + ), + ( + Some(br#" { "detail": "Remote app server not found", "extra": true } "#.to_vec()), + true, + ), + (Some(br#"{"detail":"Not Found"}"#.to_vec()), false), + (Some(b"Not Found".to_vec()), false), + (Some(b"{".to_vec()), false), + (Some(Vec::new()), false), + (None, false), + ]; + + for (body, expected) in cases { + let response = tungstenite::http::Response::builder() + .status(/*status*/ 404) + .body(body) + .expect("response should build"); + assert_eq!( + websocket_response_reports_missing_remote_app_server(&response), + expected + ); + } + + let response = tungstenite::http::Response::builder() + .status(/*status*/ 503) + .body(Some( + br#"{"detail":"Remote app server not found"}"#.to_vec(), + )) + .expect("response should build"); + assert!(!websocket_response_reports_missing_remote_app_server( + &response + )); + } + fn remote_control_status_channel() -> ( RemoteControlStatusPublisher, watch::Receiver,