From 781c183c3bf991ae938ffed7ac97d7c018c2d540 Mon Sep 17 00:00:00 2001 From: bkotsopoulos Date: Thu, 3 Sep 2026 21:22:00 +0000 Subject: [PATCH] Bound Noise handshakes by the exec server initialization timeout (#42623) ## What changed - Wait for the authenticated Noise handshake before sending the JSON-RPC `initialize` request. - Share the configured initialization timeout across the handshake and RPC, while preserving the configured duration in timeout diagnostics. - Record the handshake and initialization RPC as child spans of the existing initialization span, and terminate transport tasks when handshake readiness fails. ## Testing - Cover stalled handshake timeouts and cleanup of the underlying WebSocket. - Cover timeout reporting after the handshake has consumed part of the budget. GitOrigin-RevId: c3eb5f7cfedd8d155fde5b7614fc2eb399ce337c --- codex-rs/exec-server/src/client.rs | 82 ++++++-- codex-rs/exec-server/src/client_recovery.rs | 6 +- codex-rs/exec-server/src/client_transport.rs | 179 ++++++++++++++---- .../exec-server/src/client_transport_tests.rs | 154 ++++++++++++++- .../exec-server/src/noise_relay/harness.rs | 42 ++-- .../src/noise_relay/harness_tests.rs | 11 ++ codex-rs/exec-server/src/noise_relay/mod.rs | 2 +- codex-rs/exec-server/src/relay_noise_tests.rs | 7 +- codex-rs/exec-server/src/rpc.rs | 16 ++ 9 files changed, 433 insertions(+), 66 deletions(-) diff --git a/codex-rs/exec-server/src/client.rs b/codex-rs/exec-server/src/client.rs index 918d4467c9..f832d9be13 100644 --- a/codex-rs/exec-server/src/client.rs +++ b/codex-rs/exec-server/src/client.rs @@ -328,6 +328,16 @@ pub struct ExecServerClient { recovery_policy: RecoveryPolicy, } +/// State carried from a Noise readiness wait into the initialize RPC. +/// +/// The span preserves the existing logical initialize operation while +/// `timeout_for_error` keeps diagnostics tied to the caller's configured +/// budget after readiness has consumed part of that budget. +pub(crate) struct NoiseInitializeContext { + pub(crate) span: tracing::Span, + pub(crate) timeout_for_error: Duration, +} + struct ActiveProcessStart { inner: Arc, } @@ -721,23 +731,41 @@ impl ExecServerClient { &self, rpc_client: &RpcClient, options: ExecServerClientConnectOptions, + noise_context: Option, ) -> Result { let ExecServerClientConnectOptions { client_name, initialize_timeout, resume_session_id, } = options; + let timeout_for_error = noise_context + .as_ref() + .map_or(initialize_timeout, |context| context.timeout_for_error); timeout(initialize_timeout, async { - let response: InitializeResponse = rpc_client - .call( - INITIALIZE_METHOD, - &InitializeParams { - client_name, - resume_session_id, - }, - ) - .await?; + let params = InitializeParams { + client_name, + resume_session_id, + }; + let response: InitializeResponse = if let Some(noise_context) = noise_context { + // This is the one RPC whose wire method and trace operation + // intentionally differ: preserve the compatibility initialize + // parent while measuring only the actual RPC as initialize_rpc. + let initialize_rpc_span = tracing::info_span!( + parent: &noise_context.span, + "codex.exec_server.remote.initialize_rpc", + otel.kind = "client", + otel.name = "codex.exec_server.remote.initialize_rpc", + ); + let response = rpc_client + .call_untraced(INITIALIZE_METHOD, ¶ms) + .instrument(initialize_rpc_span) + .await; + drop(noise_context); + response? + } else { + rpc_client.call(INITIALIZE_METHOD, ¶ms).await? + }; let session_id = self .inner .session_id @@ -755,7 +783,7 @@ impl ExecServerClient { }) .await .map_err(|_| ExecServerError::InitializeTimedOut { - timeout: initialize_timeout, + timeout: timeout_for_error, })? } @@ -1108,6 +1136,36 @@ impl ExecServerClient { connection: JsonRpcConnection, options: ExecServerClientConnectOptions, reconnect_strategy: Option, + ) -> Result { + Self::connect_with_recovery_inner( + connection, + options, + reconnect_strategy, + /*noise_context*/ None, + ) + .await + } + + pub(crate) async fn connect_with_recovery_and_noise_context( + connection: JsonRpcConnection, + options: ExecServerClientConnectOptions, + reconnect_strategy: Option, + noise_context: NoiseInitializeContext, + ) -> Result { + Self::connect_with_recovery_inner( + connection, + options, + reconnect_strategy, + Some(noise_context), + ) + .await + } + + async fn connect_with_recovery_inner( + connection: JsonRpcConnection, + options: ExecServerClientConnectOptions, + reconnect_strategy: Option, + noise_context: Option, ) -> Result { let (rpc_client, events_rx) = RpcClient::new(connection); let rpc_client = Arc::new(rpc_client); @@ -1144,7 +1202,9 @@ impl ExecServerClient { // before initialize returns. Drain them immediately so a burst cannot // fill the bounded event channel and block the initialize response. client.spawn_rpc_reader(&rpc_client, events_rx); - let initialize_response = client.initialize_rpc(&rpc_client, options).await?; + let initialize_response = client + .initialize_rpc(&rpc_client, options, noise_context) + .await?; if let Some(info) = initialize_response.environment_info { assert!( client.inner.environment_info.set(info).is_ok(), diff --git a/codex-rs/exec-server/src/client_recovery.rs b/codex-rs/exec-server/src/client_recovery.rs index 8bb5a6e720..8e9a8c77b2 100644 --- a/codex-rs/exec-server/src/client_recovery.rs +++ b/codex-rs/exec-server/src/client_recovery.rs @@ -476,7 +476,7 @@ impl Inner { .as_ref() .ok_or_else(|| ExecServerError::Protocol("missing reconnect strategy".to_string()))?; let attempt = reconnect_strategy.resume(session_id).await?; - let (connection, options, attempt_permit) = attempt.into_parts(); + let (connection, options, attempt_permit, noise_context) = attempt.into_parts(); let (rpc_client, events_rx) = RpcClient::new(connection); let rpc_client = Arc::new(rpc_client); let client = ExecServerClient { @@ -488,7 +488,9 @@ impl Inner { // burst cannot fill the bounded event channel and block the initialize // response behind it. client.spawn_rpc_reader(&rpc_client, events_rx); - client.initialize_rpc(&rpc_client, options).await?; + client + .initialize_rpc(&rpc_client, options, noise_context) + .await?; self.recover_processes(&rpc_client).await?; Ok((rpc_client, attempt_permit)) diff --git a/codex-rs/exec-server/src/client_transport.rs b/codex-rs/exec-server/src/client_transport.rs index 0fd88bbbbc..b1c1ef09ee 100644 --- a/codex-rs/exec-server/src/client_transport.rs +++ b/codex-rs/exec-server/src/client_transport.rs @@ -23,6 +23,7 @@ use http::HeaderMap; use crate::ExecServerClient; use crate::ExecServerError; +use crate::client::NoiseInitializeContext; use crate::client::accepted::AcceptedConnectionSource; use crate::client::is_retryable_registry_error; use crate::client::registry_recovery_retry_delay; @@ -39,7 +40,7 @@ use crate::client_api::StdioExecServerConnectArgs; use crate::connection::JsonRpcConnection; use crate::noise_channel::NoiseChannelIdentity; use crate::noise_relay::NoiseHarnessConnectionArgs; -use crate::noise_relay::noise_harness_connection_from_websocket; +use crate::noise_relay::noise_harness_connection_from_websocket_with_readiness; use crate::noise_relay::noise_relay_websocket_config; use crate::relay::harness_connection_from_websocket; use crate::trace_context::current_rendezvous_headers; @@ -57,6 +58,19 @@ pub(crate) struct ReconnectAttempt { connection: JsonRpcConnection, options: ExecServerClientConnectOptions, attempt_permit: Option, + noise_context: Option, +} + +struct OpenNoiseRendezvousConnection { + connection: JsonRpcConnection, + options: ExecServerClientConnectOptions, + handshake_ready: tokio::sync::oneshot::Receiver<()>, +} + +struct ReadyNoiseRendezvousConnection { + connection: JsonRpcConnection, + options: ExecServerClientConnectOptions, + noise_context: NoiseInitializeContext, } impl ReconnectAttempt { @@ -68,6 +82,20 @@ impl ReconnectAttempt { connection, options, attempt_permit: None, + noise_context: None, + } + } + + fn with_noise_context( + connection: JsonRpcConnection, + options: ExecServerClientConnectOptions, + noise_context: NoiseInitializeContext, + ) -> Self { + Self { + connection, + options, + attempt_permit: None, + noise_context: Some(noise_context), } } @@ -80,6 +108,7 @@ impl ReconnectAttempt { connection, options, attempt_permit: Some(attempt_permit), + noise_context: None, } } @@ -89,8 +118,14 @@ impl ReconnectAttempt { JsonRpcConnection, ExecServerClientConnectOptions, Option, + Option, ) { - (self.connection, self.options, self.attempt_permit) + ( + self.connection, + self.options, + self.attempt_permit, + self.noise_context, + ) } } @@ -142,7 +177,7 @@ impl ExecServerReconnectStrategy { http_client_factory, } => { let bundle = provider.connect_bundle(identity.public_key()).await?; - let (connection, options) = ExecServerClient::open_noise_rendezvous_connection( + let opened = ExecServerClient::open_noise_rendezvous_connection( NoiseRendezvousConnectArgs { bundle, harness_identity: identity.clone(), @@ -154,7 +189,12 @@ impl ExecServerReconnectStrategy { }, ) .await?; - Ok(ReconnectAttempt::new(connection, options)) + let ready = ExecServerClient::finish_noise_rendezvous_connection(opened).await?; + Ok(ReconnectAttempt::with_noise_context( + ready.connection, + ready.options, + ready.noise_context, + )) } } } @@ -213,13 +253,12 @@ impl ExecServerClient { http_headers, ), ExecServerTransportParams::NoiseRendezvous { provider, identity } => { - let (connection, options, executor_public_key) = - Self::open_initial_noise_rendezvous_connection( - &provider, - &identity, - http_client_factory.clone(), - ) - .await?; + let (ready, executor_public_key) = Self::open_initial_noise_rendezvous_connection( + &provider, + &identity, + http_client_factory.clone(), + ) + .await?; let reconnect_strategy = ExecServerReconnectStrategy::NoiseRendezvous { executor_public_key, provider, @@ -229,8 +268,13 @@ impl ExecServerClient { initialize_timeout: DEFAULT_REMOTE_EXEC_SERVER_INITIALIZE_TIMEOUT, http_client_factory, }; - return Self::connect_with_recovery(connection, options, Some(reconnect_strategy)) - .await; + return Self::connect_with_recovery_and_noise_context( + ready.connection, + ready.options, + Some(reconnect_strategy), + ready.noise_context, + ) + .await; } ExecServerTransportParams::StdioCommand { command, @@ -265,14 +309,8 @@ impl ExecServerClient { provider: &Arc, identity: &NoiseChannelIdentity, http_client_factory: HttpClientFactory, - ) -> Result< - ( - JsonRpcConnection, - ExecServerClientConnectOptions, - crate::NoiseChannelPublicKey, - ), - ExecServerError, - > { + ) -> Result<(ReadyNoiseRendezvousConnection, crate::NoiseChannelPublicKey), ExecServerError> + { let open_connection = |bundle: NoiseRendezvousConnectBundle| { Self::open_noise_rendezvous_connection(NoiseRendezvousConnectArgs { bundle, @@ -344,8 +382,9 @@ impl ExecServerClient { result = connect_bundle().await; } result => { - return result - .map(|(connection, options)| (connection, options, executor_public_key)); + let opened = result?; + let ready = Self::finish_noise_rendezvous_connection(opened).await?; + return Ok((ready, executor_public_key)); } } } @@ -439,8 +478,15 @@ impl ExecServerClient { pub async fn connect_noise_rendezvous( args: NoiseRendezvousConnectArgs, ) -> Result { - let (connection, options) = Self::open_noise_rendezvous_connection(args).await?; - Self::connect(connection, options).await + let opened = Self::open_noise_rendezvous_connection(args).await?; + let ready = Self::finish_noise_rendezvous_connection(opened).await?; + Self::connect_with_recovery_and_noise_context( + ready.connection, + ready.options, + /*reconnect_strategy*/ None, + ready.noise_context, + ) + .await } #[tracing::instrument( @@ -453,9 +499,9 @@ impl ExecServerClient { executor_registration_id = %args.bundle.executor_registration_id, ) )] - pub(crate) async fn open_noise_rendezvous_connection( + async fn open_noise_rendezvous_connection( args: NoiseRendezvousConnectArgs, - ) -> Result<(JsonRpcConnection, ExecServerClientConnectOptions), ExecServerError> { + ) -> Result { ensure_rustls_crypto_provider(); // Keep the registry-issued URL, key, and authorization together for this // connection attempt. @@ -509,7 +555,7 @@ impl ExecServerClient { })?; let connection_label = format!("Noise exec-server rendezvous websocket {diagnostic_url}"); - let connection = noise_harness_connection_from_websocket( + let connection = noise_harness_connection_from_websocket_with_readiness( stream, NoiseHarnessConnectionArgs { connection_label, @@ -520,14 +566,83 @@ impl ExecServerClient { harness_key_authorization, }, ); - Ok(( - connection, - ExecServerClientConnectOptions { + Ok(OpenNoiseRendezvousConnection { + connection: connection.connection, + options: ExecServerClientConnectOptions { client_name, initialize_timeout, resume_session_id, }, - )) + handshake_ready: connection.handshake_ready, + }) + } + + #[tracing::instrument( + name = "codex.exec_server.remote.noise.handshake", + skip_all, + parent = initialize_span, + fields( + otel.kind = "client", + otel.name = "codex.exec_server.remote.noise.handshake", + ) + )] + async fn wait_for_noise_handshake( + handshake_ready: &mut tokio::sync::oneshot::Receiver<()>, + deadline: Instant, + initialize_timeout: Duration, + initialize_span: &tracing::Span, + ) -> Result<(), ExecServerError> { + match timeout_at(deadline, handshake_ready).await { + Ok(Ok(())) => Ok(()), + Ok(Err(_)) => Err(ExecServerError::Disconnected( + "Noise harness handshake failed before connection became ready".to_string(), + )), + Err(_) => Err(ExecServerError::InitializeTimedOut { + timeout: initialize_timeout, + }), + } + } + + async fn finish_noise_rendezvous_connection( + mut connection: OpenNoiseRendezvousConnection, + ) -> Result { + // Preserve the legacy initialize request span as the post-WebSocket + // startup parent while making its two child operations visible. + let initialize_timeout = connection.options.initialize_timeout; + let noise_context = NoiseInitializeContext { + span: tracing::info_span!( + "codex.exec_server.request", + otel.kind = "client", + otel.name = "initialize", + method = "initialize", + ), + timeout_for_error: initialize_timeout, + }; + let deadline = Instant::now() + initialize_timeout; + let readiness = Self::wait_for_noise_handshake( + &mut connection.handshake_ready, + deadline, + initialize_timeout, + &noise_context.span, + ) + .await; + if let Err(error) = readiness { + // Unlike the normal connect path, the connection has not reached + // RpcClient yet, so its Drop implementation cannot abort the + // transport task for us. + connection.connection.transport.terminate(); + for task in &connection.connection.task_handles { + task.abort(); + } + return Err(error); + } + let mut options = connection.options; + options.initialize_timeout = deadline.saturating_duration_since(Instant::now()); + Ok(ReadyNoiseRendezvousConnection { + connection: connection.connection, + options, + noise_context, + }) } pub(crate) async fn connect_stdio_command( diff --git a/codex-rs/exec-server/src/client_transport_tests.rs b/codex-rs/exec-server/src/client_transport_tests.rs index 362de81654..63345c227a 100644 --- a/codex-rs/exec-server/src/client_transport_tests.rs +++ b/codex-rs/exec-server/src/client_transport_tests.rs @@ -4,13 +4,20 @@ use std::sync::Arc; use std::sync::Mutex; use anyhow::Result; +use codex_exec_server_protocol::JSONRPCMessage; use futures::FutureExt; +use futures::SinkExt; +use futures::StreamExt; use futures::future::BoxFuture; use pretty_assertions::assert_eq; +use tokio::io::AsyncBufReadExt; use tokio::io::AsyncReadExt; use tokio::io::AsyncWriteExt; +use tokio::io::BufReader; +use tokio::io::duplex; use tokio::net::TcpListener; use tokio_tungstenite::accept_async; +use tokio_tungstenite::tungstenite::Message; use super::ExecServerClient; use super::ExecServerReconnectStrategy; @@ -20,10 +27,21 @@ use super::INITIAL_REGISTRY_REQUEST_TIMEOUT; use crate::ExecServerError; use crate::NoiseChannelIdentity; use crate::NoiseChannelPublicKey; +use crate::NoiseRendezvousConnectArgs; use crate::NoiseRendezvousConnectBundle; use crate::NoiseRendezvousConnectProvider; +use crate::client::NoiseInitializeContext; use crate::client_api::DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT; use crate::client_api::DEFAULT_REMOTE_EXEC_SERVER_INITIALIZE_TIMEOUT; +use crate::client_api::ExecServerClientConnectOptions; +use crate::connection::JsonRpcConnection; +use crate::noise_channel::PendingResponderHandshake; +use crate::noise_channel::noise_channel_prologue; +use crate::protocol::INITIALIZE_METHOD; +use crate::relay::RelayFrameBodyKind; +use crate::relay::decode_relay_message_frame; +use crate::relay::encode_relay_message_frame; +use crate::relay_proto::RelayMessageFrame; #[derive(Default)] struct SequenceNoiseConnectProvider { @@ -85,7 +103,7 @@ impl SequenceNoiseConnectProvider { ), ) .await - .map(|(connection, options, _)| (connection, options)) + .map(|(ready, _)| (ready.connection, ready.options)) } } @@ -132,6 +150,110 @@ fn registry_error(status: http::StatusCode, code: &str) -> ExecServerError { } } +#[tokio::test] +async fn noise_handshake_uses_initialize_timeout() -> Result<()> { + 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?; + // Drain the frames sent before the harness waits for the responder, + // then verify that a timed-out readiness wait closes the socket. + assert!(websocket.next().await.is_some()); + assert!(websocket.next().await.is_some()); + let closed = + tokio::time::timeout(std::time::Duration::from_secs(1), websocket.next()).await?; + assert!( + matches!(closed, None | Some(Ok(Message::Close(_))) | Some(Err(_))), + "timed-out Noise handshake must close its websocket" + ); + anyhow::Ok(()) + }); + let initialize_timeout = std::time::Duration::from_millis(1); + let opened = ExecServerClient::open_noise_rendezvous_connection(NoiseRendezvousConnectArgs { + bundle: NoiseRendezvousConnectBundle { + websocket_url, + environment_id: "environment".to_string(), + executor_registration_id: "registration".to_string(), + executor_public_key: NoiseChannelIdentity::generate()?.public_key(), + harness_key_authorization: "authorization".to_string(), + }, + harness_identity: NoiseChannelIdentity::generate()?, + client_name: "test".to_string(), + connect_timeout: DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT, + 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("stalled Noise handshake must time out"); + assert!(matches!( + error, + ExecServerError::InitializeTimedOut { timeout } if timeout == initialize_timeout + )); + + server.await??; + Ok(()) +} + +#[tokio::test] +async fn deferred_initialize_timeout_reports_configured_budget() { + let (client_stdin, server_reader) = duplex(1 << 20); + let (server_writer, client_stdout) = duplex(1 << 20); + let server = tokio::spawn(async move { + let _server_writer = server_writer; + let mut lines = BufReader::new(server_reader).lines(); + let line = lines + .next_line() + .await + .expect("initialize read should succeed") + .expect("initialize request should arrive"); + let request: JSONRPCMessage = + serde_json::from_str(&line).expect("initialize request should parse"); + assert!( + matches!( + request, + JSONRPCMessage::Request(ref request) if request.method == INITIALIZE_METHOD + ), + "expected initialize request, got {request:?}" + ); + futures::future::pending::<()>().await; + }); + let configured_timeout = std::time::Duration::from_secs(10); + let error = ExecServerClient::connect_with_recovery_and_noise_context( + JsonRpcConnection::from_stdio( + client_stdout, + client_stdin, + "timeout-test-client".to_string(), + ), + ExecServerClientConnectOptions { + client_name: "timeout-test-client".to_string(), + initialize_timeout: std::time::Duration::from_millis(1), + resume_session_id: None, + }, + /*reconnect_strategy*/ None, + NoiseInitializeContext { + span: tracing::info_span!("codex.exec_server.request"), + timeout_for_error: configured_timeout, + }, + ) + .await + .err() + .expect("initialize RPC must time out"); + assert!(matches!( + error, + ExecServerError::InitializeTimedOut { timeout } if timeout == configured_timeout + )); + server.abort(); + let _ = server.await; +} + #[tokio::test(start_paused = true)] async fn initial_noise_connection_bounds_offline_retries() -> Result<()> { let sequence = Arc::new(SequenceNoiseConnectProvider::default()); @@ -382,14 +504,40 @@ async fn initial_noise_connection_refreshes_bundle_after_exhausting_initial_retr }); let accepted_listener = TcpListener::bind("127.0.0.1:0").await?; let accepted_url = format!("ws://{}", accepted_listener.local_addr()?); + let executor_identity = NoiseChannelIdentity::generate()?; + let executor_public_key = executor_identity.public_key(); let accepted_server = tokio::spawn(async move { let (socket, _) = accepted_listener.accept().await?; - let _websocket = accept_async(socket).await?; + let mut websocket = accept_async(socket).await?; + let Message::Binary(resume_payload) = websocket.next().await.unwrap()? else { + anyhow::bail!("expected Noise relay resume frame"); + }; + let resume = decode_relay_message_frame(resume_payload.as_ref())?; + assert_eq!(resume.validate()?, RelayFrameBodyKind::Resume); + let Message::Binary(handshake_payload) = websocket.next().await.unwrap()? else { + anyhow::bail!("expected Noise relay handshake frame"); + }; + let handshake = decode_relay_message_frame(handshake_payload.as_ref())?; + let stream_id = handshake.stream_id.clone(); + let prologue = noise_channel_prologue("environment", "registration", &stream_id); + let pending = PendingResponderHandshake::read_request( + &executor_identity, + &prologue, + &handshake.into_handshake_payload()?, + )?; + let (_transport, response) = pending.complete()?; + websocket + .send(Message::Binary( + encode_relay_message_frame(&RelayMessageFrame::handshake(stream_id, response)) + .into(), + )) + .await?; anyhow::Ok(()) }); let sequence = Arc::new(SequenceNoiseConnectProvider::default()); let unauthorized_bundle = test_bundle(unauthorized_url.clone())?; - let accepted_bundle = test_bundle(accepted_url.clone())?; + let mut accepted_bundle = test_bundle(accepted_url.clone())?; + accepted_bundle.executor_public_key = executor_public_key; sequence.push_response(async { tokio::time::pause(); Err(registry_error( diff --git a/codex-rs/exec-server/src/noise_relay/harness.rs b/codex-rs/exec-server/src/noise_relay/harness.rs index 472d9c8184..793f7c1e12 100644 --- a/codex-rs/exec-server/src/noise_relay/harness.rs +++ b/codex-rs/exec-server/src/noise_relay/harness.rs @@ -12,6 +12,7 @@ use futures::SinkExt; use futures::Stream; use futures::StreamExt; use tokio::sync::mpsc; +use tokio::sync::oneshot; use tokio::sync::watch; use tokio_tungstenite::tungstenite::Message; use tracing::Instrument; @@ -59,6 +60,13 @@ pub(crate) struct NoiseHarnessConnectionArgs { pub(crate) harness_key_authorization: String, } +/// One Noise-backed JSON-RPC connection plus a signal that its authenticated +/// transport is ready for application messages. +pub(crate) struct NoiseHarnessConnection { + pub(crate) connection: JsonRpcConnection, + pub(crate) handshake_ready: oneshot::Receiver<()>, +} + // Reset frames are cleartext relay control and are not authenticated by Noise. // Preserve the availability signal while replacing attacker-controlled reason // text before it reaches disconnect diagnostics. @@ -66,17 +74,15 @@ const NOISE_RELAY_RESET_DISCONNECT_REASON: &str = "Noise relay stream reset"; // Give a Pong already queued behind data a bounded chance to reach the reader. const MAX_FRAMES_DRAINED_AFTER_PONG_DEADLINE: usize = 32; -/// Adapt one harness rendezvous websocket into an authenticated JSON-RPC connection. +/// Adapt one harness rendezvous websocket and expose when hybrid IK completes. /// -/// The returned connection is not usable until the background task completes -/// hybrid IK against the registry-pinned exec-server key. Rendezvous can see -/// stream metadata and ciphertext, but never JSON-RPC plaintext or either -/// endpoint's private key. Failures close the connection rather than falling -/// back to plaintext. -pub(crate) fn noise_harness_connection_from_websocket( +/// Callers that send application messages immediately after opening the +/// websocket must await handshake_ready first. Dropping that receiver keeps +/// the legacy fire-and-forget behavior for existing connection owners. +pub(crate) fn noise_harness_connection_from_websocket_with_readiness( stream: T, args: NoiseHarnessConnectionArgs, -) -> JsonRpcConnection +) -> NoiseHarnessConnection where T: Sink + Stream> + Unpin + Send + 'static, E: std::fmt::Display + Send + 'static, @@ -93,6 +99,7 @@ where let (outgoing_tx, mut outgoing_rx) = mpsc::channel(CHANNEL_CAPACITY); let (incoming_tx, incoming_rx) = mpsc::channel(CHANNEL_CAPACITY); let (disconnected_tx, disconnected_rx) = watch::channel(false); + let (handshake_ready_tx, handshake_ready_rx) = oneshot::channel(); let stream_span = tracing::debug_span!("noise_relay.stream", noise_side = "harness",); debug!( environment_id, @@ -100,6 +107,7 @@ where ); let websocket_task = tokio::spawn(async move { + let mut handshake_ready_tx = Some(handshake_ready_tx); let mut websocket = stream; // Bind the Noise transcript to the exact environment registration and @@ -221,6 +229,9 @@ where noise_outcome = "ok", "Noise harness handshake completed" ); + if let Some(handshake_ready_tx) = handshake_ready_tx.take() { + let _ = handshake_ready_tx.send(()); + } break transport; } Err(error) => { @@ -525,12 +536,15 @@ where } .instrument(stream_span)); - JsonRpcConnection { - outgoing_tx, - incoming_rx, - disconnected_rx, - task_handles: vec![websocket_task], - transport: JsonRpcTransport::Plain, + NoiseHarnessConnection { + connection: JsonRpcConnection { + outgoing_tx, + incoming_rx, + disconnected_rx, + task_handles: vec![websocket_task], + transport: JsonRpcTransport::Plain, + }, + handshake_ready: handshake_ready_rx, } } diff --git a/codex-rs/exec-server/src/noise_relay/harness_tests.rs b/codex-rs/exec-server/src/noise_relay/harness_tests.rs index 16df1cbf0f..5530941ca0 100644 --- a/codex-rs/exec-server/src/noise_relay/harness_tests.rs +++ b/codex-rs/exec-server/src/noise_relay/harness_tests.rs @@ -29,6 +29,17 @@ use super::*; use crate::connection::JsonRpcConnectionEvent; use crate::noise_channel::PendingResponderHandshake; +fn noise_harness_connection_from_websocket( + stream: T, + args: NoiseHarnessConnectionArgs, +) -> JsonRpcConnection +where + T: Sink + Stream> + Unpin + Send + 'static, + E: std::fmt::Display + Send + 'static, +{ + noise_harness_connection_from_websocket_with_readiness(stream, args).connection +} + const ENVIRONMENT_ID: &str = "environment-1"; const EXECUTOR_REGISTRATION_ID: &str = "registration-1"; diff --git a/codex-rs/exec-server/src/noise_relay/mod.rs b/codex-rs/exec-server/src/noise_relay/mod.rs index 12bbdd6f13..dc6030e13b 100644 --- a/codex-rs/exec-server/src/noise_relay/mod.rs +++ b/codex-rs/exec-server/src/noise_relay/mod.rs @@ -9,7 +9,7 @@ use tokio_tungstenite::tungstenite::protocol::WebSocketConfig; use crate::ExecServerError; pub(crate) use harness::NoiseHarnessConnectionArgs; -pub(crate) use harness::noise_harness_connection_from_websocket; +pub(crate) use harness::noise_harness_connection_from_websocket_with_readiness; pub(crate) const NOISE_RELAY_RESET_REASON: &str = "noise_relay_protocol_error"; diff --git a/codex-rs/exec-server/src/relay_noise_tests.rs b/codex-rs/exec-server/src/relay_noise_tests.rs index ff402b316d..b92d22b9bf 100644 --- a/codex-rs/exec-server/src/relay_noise_tests.rs +++ b/codex-rs/exec-server/src/relay_noise_tests.rs @@ -31,7 +31,7 @@ use crate::noise_channel::NoiseChannelIdentity; use crate::noise_channel::NoiseChannelPublicKey; use crate::noise_channel::noise_channel_prologue; use crate::noise_relay::NoiseHarnessConnectionArgs; -use crate::noise_relay::noise_harness_connection_from_websocket; +use crate::noise_relay::noise_harness_connection_from_websocket_with_readiness; use crate::noise_relay::stream_handler::NoiseOutboundMessage; use crate::noise_relay::stream_handler::NoiseStreamConnection; use crate::noise_relay::stream_handler::NoiseStreamHandler; @@ -208,7 +208,7 @@ async fn processor_exit_resets_noise_harness_stream() -> Result<()> { release, }, ))); - let mut connection = noise_harness_connection_from_websocket( + let mut connection = noise_harness_connection_from_websocket_with_readiness( harness_websocket, NoiseHarnessConnectionArgs { connection_label: "processor exit test".to_string(), @@ -218,7 +218,8 @@ async fn processor_exit_resets_noise_harness_stream() -> Result<()> { responder_public_key: identity.public_key(), harness_key_authorization: "authorization".to_string(), }, - ); + ) + .connection; // Valid JSON reaches the processor; an unsolicited response closes it and // aborts its writer. The physical relay must still deliver the reset. connection diff --git a/codex-rs/exec-server/src/rpc.rs b/codex-rs/exec-server/src/rpc.rs index b6eb30adba..998413756a 100644 --- a/codex-rs/exec-server/src/rpc.rs +++ b/codex-rs/exec-server/src/rpc.rs @@ -539,6 +539,22 @@ impl RpcClient { ) )] pub(crate) async fn call(&self, method: &str, params: &P) -> Result + where + P: Serialize, + T: DeserializeOwned, + { + self.call_untraced(method, params).await + } + + /// Send one request without creating the standard request span. + /// + /// Callers use this only when they install a more precise request span + /// around the same wire operation. + pub(crate) async fn call_untraced( + &self, + method: &str, + params: &P, + ) -> Result where P: Serialize, T: DeserializeOwned,