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,