agent/codex
1
0
Fork 0
mirror of https://github.com/pchuan98/codex.git synced 2026-07-01 00:31:56 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(remote-control): preserve enrollment on generic websocket 404s (#26741)

Browse Source 
## 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
This commit is contained in:
Anton Panasenko
2026-06-05 22:54:57 -07:00
committed by GitHub
Unverified
parent 3ea9e98333
commit 87b808bb57
2 changed files with 210 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/app-server-transport/src/transport/remote_control/tests.rs
+137 -2
View File
@@ -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::<TransportEvent>(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,
codex-rs/app-server-transport/src/transport/remote_control/websocket.rs
+73 -1
View File
@@ -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<ServerEnvelope>>,
@@ -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(|| "<missing>".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<Option<Vec<u8>>>,
) -> bool {
response.status().as_u16() == 404
&& response.body().as_deref().is_some_and(|body| {
serde_json::from_slice::<serde_json::Value>(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<RemoteControlStatusChangedNotification>,