Preserve disconnect reasons in Noise handshake errors (#54637)

## Why

Noise handshake failures before connection readiness returned a generic error even when the transport had reported a specific disconnect reason. No `RpcClient` was reading the connection event channel at that point.

## What changed

Read the pending disconnect event when a Noise rendezvous handshake fails and append its reason to the handshake error when available.

## Testing

Add a regression test covering unexpected text frames and relay stream resets, checking that each produces the expected transport reason in the handshake error.

GitOrigin-RevId: 9b8ef9c8dcba75b46d51dc67784af46d4ca74b85
This commit is contained in:
virajrai-oai
2026-10-16 15:43:21 +00:00
committed by copyberry
parent b5115cdd66
commit 34420d4264
2 changed files with 73 additions and 1 deletions
+13 -1
View File
@@ -45,6 +45,7 @@ use crate::client_api::RemoteExecServerConnectArgs;
use crate::client_api::StdioExecServerCommand;
use crate::client_api::StdioExecServerConnectArgs;
use crate::connection::JsonRpcConnection;
use crate::connection::JsonRpcConnectionEvent;
use crate::noise_channel::NoiseChannelIdentity;
use crate::noise_relay::NoiseHarnessConnectionArgs;
use crate::noise_relay::noise_harness_connection_from_websocket_with_readiness;
@@ -754,7 +755,18 @@ impl ExecServerClient {
&noise_context.span,
)
.await;
if let Err(error) = readiness {
if let Err(mut error) = readiness {
// Before readiness, the transport reports the reason through the
// connection event channel, but no RpcClient is reading it yet.
if matches!(error, ExecServerError::Disconnected(_))
&& let Ok(JsonRpcConnectionEvent::Disconnected {
reason: Some(reason),
}) = connection.connection.incoming_rx.try_recv()
{
error = ExecServerError::Disconnected(format!(
"Noise harness handshake failed before connection became ready: {reason}"
));
}
// Unlike the normal connect path, the connection has not reached
// RpcClient yet, so its Drop implementation cannot abort the
// transport task for us.
@@ -175,6 +175,66 @@ fn registry_error(status: http::StatusCode, code: &str) -> ExecServerError {
}
}
#[tokio::test]
async fn noise_handshake_reports_disconnect_reason() -> Result<()> {
for (send_reset, expected_reason) in [
(
false,
"Noise relay transport expects binary protobuf frames",
),
(true, "Noise relay stream reset"),
] {
let listener = TcpListener::bind("127.0.0.1:0").await?;
let websocket_url = format!("ws://{}", listener.local_addr()?);
let server = tokio::spawn(async move {
let (socket, _) = listener.accept().await?;
let mut websocket = accept_async(socket).await?;
let Message::Binary(resume) = websocket.next().await.unwrap()? else {
anyhow::bail!("expected resume frame");
};
let stream_id = decode_relay_message_frame(&resume)?.stream_id;
let _handshake = websocket.next().await.unwrap()?;
let response = if send_reset {
let frame = RelayMessageFrame::reset(stream_id, "untrusted reason".to_string());
Message::Binary(encode_relay_message_frame(&frame).into())
} else {
Message::Text("untrusted text".into())
};
websocket.send(response).await?;
anyhow::Ok(())
});
let opened =
ExecServerClient::open_noise_rendezvous_connection(NoiseRendezvousConnectArgs {
bundle: test_bundle(websocket_url)?,
harness_identity: NoiseChannelIdentity::generate()?,
client_name: "test".to_string(),
connect_timeout: DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT,
initialize_timeout: DEFAULT_REMOTE_EXEC_SERVER_INITIALIZE_TIMEOUT,
resume_session_id: None,
http_client_factory: codex_http_client::HttpClientFactory::new(
codex_http_client::OutboundProxyPolicy::ReqwestDefault,
),
})
.await?;
let error = ExecServerClient::finish_noise_rendezvous_connection(opened)
.await
.err()
.expect("invalid handshake response must fail");
let ExecServerError::Disconnected(reason) = error else {
anyhow::bail!("expected disconnect, got {error:?}");
};
assert_eq!(
reason,
format!(
"Noise harness handshake failed before connection became ready: {expected_reason}"
)
);
server.await??;
}
Ok(())
}
#[tokio::test]
async fn noise_handshake_uses_initialize_timeout() -> Result<()> {
let listener = TcpListener::bind("127.0.0.1:0").await?;