openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-11-05 20:10:33 +00:00
Code Issues Packages Projects Releases Wiki Activity

Bound Noise handshakes by the exec server initialization timeout (#42623)

Browse Source 
## 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
This commit is contained in:
bkotsopoulos
2026-09-03 22:16:19 +00:00
committed by copyberry
parent 1f7b99922a
commit 781c183c3b
9 changed files with 433 additions and 66 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/exec-server/src/client.rs
+71 -11
View File
@@ -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<Inner>,
}
@@ -721,23 +731,41 @@ impl ExecServerClient {
&self,
rpc_client: &RpcClient,
options: ExecServerClientConnectOptions,
noise_context: Option<NoiseInitializeContext>,
) -> Result<InitializeResponse, ExecServerError> {
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, &params)
.instrument(initialize_rpc_span)
.await;
drop(noise_context);
response?
} else {
rpc_client.call(INITIALIZE_METHOD, &params).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<ExecServerReconnectStrategy>,
) -> Result<Self, ExecServerError> {
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<ExecServerReconnectStrategy>,
noise_context: NoiseInitializeContext,
) -> Result<Self, ExecServerError> {
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<ExecServerReconnectStrategy>,
noise_context: Option<NoiseInitializeContext>,
) -> Result<Self, ExecServerError> {
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(),
codex-rs/exec-server/src/client_recovery.rs
+4 -2
View File
@@ -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))
codex-rs/exec-server/src/client_transport.rs
+147 -32
View File
@@ -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<OwnedSemaphorePermit>,
noise_context: Option<NoiseInitializeContext>,
}
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<OwnedSemaphorePermit>,
Option<NoiseInitializeContext>,
) {
(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<dyn NoiseRendezvousConnectProvider>,
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<Self, ExecServerError> {
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<OpenNoiseRendezvousConnection, ExecServerError> {
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<ReadyNoiseRendezvousConnection, ExecServerError> {
// 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(
codex-rs/exec-server/src/client_transport_tests.rs
+151 -3
View File
@@ -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(
codex-rs/exec-server/src/noise_relay/harness.rs
+28 -14
View File
@@ -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<T, E>(
/// 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<T, E>(
stream: T,
args: NoiseHarnessConnectionArgs,
) -> JsonRpcConnection
) -> NoiseHarnessConnection
where
T: Sink<Message, Error = E> + Stream<Item = Result<Message, E>> + 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,
}
}
codex-rs/exec-server/src/noise_relay/harness_tests.rs
+11
View File
@@ -29,6 +29,17 @@ use super::*;
use crate::connection::JsonRpcConnectionEvent;
use crate::noise_channel::PendingResponderHandshake;
fn noise_harness_connection_from_websocket<T, E>(
stream: T,
args: NoiseHarnessConnectionArgs,
) -> JsonRpcConnection
where
T: Sink<Message, Error = E> + Stream<Item = Result<Message, E>> + 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";
codex-rs/exec-server/src/noise_relay/mod.rs
+1 -1
View File
@@ -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";
codex-rs/exec-server/src/relay_noise_tests.rs
+4 -3
View File
@@ -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
codex-rs/exec-server/src/rpc.rs
+16
View File
@@ -539,6 +539,22 @@ impl RpcClient {
)
)]
pub(crate) async fn call<P, T>(&self, method: &str, params: &P) -> Result<T, RpcCallError>
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<P, T>(
&self,
method: &str,
params: &P,
) -> Result<T, RpcCallError>
where
P: Serialize,
T: DeserializeOwned,