From 042e61726d28d79cad6f307b2f3ab085861c2212 Mon Sep 17 00:00:00 2001 From: richardopenai Date: Wed, 1 Jul 2026 14:15:34 -0700 Subject: [PATCH 1/3] [codex] bound Rendezvous WebSocket liveness (#30643) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary - require a Pong within 60 seconds for established Noise Rendezvous WebSockets on both the harness and executor - bound steady-state WebSocket writes and harness event delivery so backpressure cannot mask the deadline - classify executor disconnects with bounded reasons and feed them into the existing reconnect metric and structured log - cover silent peers, responsive peers, continuous non-Pong traffic, and local application backpressure ## Why The existing periodic Pings did not track Pongs, so a half-open or blackholed connection could remain stuck until the operating system's TCP timeout. This adds the smallest explicit liveness contract without new spans, RTT histograms, feature flags, or TCP diagnostics. ## Testing - `just test -p codex-exec-server` on devbox `richard-6` — 300 passed, 2 skipped - `just fix -p codex-exec-server` - `just fmt` - independent correctness, performance/security, and YAGNI reviews — no findings --- codex-rs/exec-server/src/lib.rs | 1 + .../exec-server/src/noise_relay/harness.rs | 314 ++++++++++--- .../src/noise_relay/harness_tests.rs | 442 ++++++++++++++++++ codex-rs/exec-server/src/relay.rs | 102 +++- codex-rs/exec-server/src/relay_noise_tests.rs | 81 ++++ codex-rs/exec-server/src/remote.rs | 12 +- .../src/websocket_pong_watchdog.rs | 44 ++ .../src/websocket_pong_watchdog_tests.rs | 34 ++ 8 files changed, 953 insertions(+), 77 deletions(-) create mode 100644 codex-rs/exec-server/src/noise_relay/harness_tests.rs create mode 100644 codex-rs/exec-server/src/websocket_pong_watchdog.rs create mode 100644 codex-rs/exec-server/src/websocket_pong_watchdog_tests.rs diff --git a/codex-rs/exec-server/src/lib.rs b/codex-rs/exec-server/src/lib.rs index 0068c2cc1a..882c83230c 100644 --- a/codex-rs/exec-server/src/lib.rs +++ b/codex-rs/exec-server/src/lib.rs @@ -29,6 +29,7 @@ mod sandboxed_file_system; mod server; mod telemetry; mod trace_context; +mod websocket_pong_watchdog; use codex_exec_server_protocol as protocol; diff --git a/codex-rs/exec-server/src/noise_relay/harness.rs b/codex-rs/exec-server/src/noise_relay/harness.rs index 1989525224..9969a3e130 100644 --- a/codex-rs/exec-server/src/noise_relay/harness.rs +++ b/codex-rs/exec-server/src/noise_relay/harness.rs @@ -6,6 +6,7 @@ //! normal `JsonRpcConnection`. Outbound JSON-RPC is framed and split into Noise //! records; inbound records are reordered before decryption and reassembly. +use futures::FutureExt; use futures::Sink; use futures::SinkExt; use futures::Stream; @@ -24,6 +25,7 @@ use crate::connection::CHANNEL_CAPACITY; use crate::connection::JsonRpcConnection; use crate::connection::JsonRpcConnectionEvent; use crate::connection::JsonRpcTransport; +use crate::connection::WEBSOCKET_KEEPALIVE_INTERVAL; use crate::noise_channel::InitiatorHandshake; use crate::noise_channel::NoiseChannelIdentity; use crate::noise_channel::NoiseChannelPublicKey; @@ -39,6 +41,9 @@ use crate::relay::decode_relay_message_frame; use crate::relay::encode_relay_message_frame; use crate::relay_proto::RelayData; use crate::relay_proto::RelayMessageFrame; +use crate::websocket_pong_watchdog::WEBSOCKET_PONG_TIMEOUT; +use crate::websocket_pong_watchdog::WEBSOCKET_PONG_TIMEOUT_REASON; +use crate::websocket_pong_watchdog::WebSocketPongWatchdog; /// Values that bind one harness websocket to the intended executor registration. /// @@ -58,6 +63,8 @@ pub(crate) struct NoiseHarnessConnectionArgs { // Preserve the availability signal while replacing attacker-controlled reason // text before it reaches disconnect diagnostics. 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. /// @@ -112,8 +119,7 @@ where &incoming_tx, &disconnected_tx, format!("failed to start Noise relay handshake: {error}"), - ) - .await; + ); return; } }; @@ -147,8 +153,7 @@ where &incoming_tx, &disconnected_tx, "Noise relay websocket ended during handshake".to_string(), - ) - .await; + ); return; }; let message = match incoming_message { @@ -158,8 +163,7 @@ where &incoming_tx, &disconnected_tx, "Noise relay websocket received close frame during handshake".to_string(), - ) - .await; + ); return; } Ok(Message::Ping(_) | Message::Pong(_) | Message::Frame(_)) => continue, @@ -168,8 +172,7 @@ where &incoming_tx, &disconnected_tx, "Noise relay transport expects binary protobuf frames".to_string(), - ) - .await; + ); return; } Err(error) => { @@ -179,8 +182,7 @@ where format!( "failed to read Noise relay websocket from {connection_label}: {error}" ), - ) - .await; + ); return; } }; @@ -191,8 +193,7 @@ where &incoming_tx, &disconnected_tx, format!("failed to parse Noise relay frame: {error}"), - ) - .await; + ); return; } }; @@ -209,8 +210,7 @@ where &incoming_tx, &disconnected_tx, format!("invalid Noise relay handshake response: {error}"), - ) - .await; + ); return; } }; @@ -228,8 +228,7 @@ where &incoming_tx, &disconnected_tx, format!("Noise relay handshake failed: {error}"), - ) - .await; + ); return; } } @@ -239,8 +238,7 @@ where &incoming_tx, &disconnected_tx, NOISE_RELAY_RESET_DISCONNECT_REASON.to_string(), - ) - .await; + ); return; } Ok( @@ -253,8 +251,7 @@ where &incoming_tx, &disconnected_tx, "Noise relay received data before handshake completion".to_string(), - ) - .await; + ); return; } } @@ -263,57 +260,177 @@ where // After the handshake, each relay sequence maps to exactly one Noise // transport record. Outbound records are encrypted once; inbound // records are reordered and deduplicated before decryption. + let mut websocket = websocket.peekable(); let mut next_outbound_seq = 0u32; let mut inbound_ciphertexts = OrderedCiphertextFrames::default(); let mut inbound_decoder = JsonRpcMessageDecoder::default(); + let mut keepalive = tokio::time::interval_at( + tokio::time::Instant::now() + WEBSOCKET_KEEPALIVE_INTERVAL, + WEBSOCKET_KEEPALIVE_INTERVAL, + ); + keepalive.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + let mut pong_watchdog = WebSocketPongWatchdog::new(WEBSOCKET_PONG_TIMEOUT); + let pong_deadline = tokio::time::sleep(WEBSOCKET_PONG_TIMEOUT); + tokio::pin!(pong_deadline); + // Keep one framed message as a cursor. Sending one Noise record per loop + // creates a scheduling point for keepalive and inbound control frames + // without splitting the WebSocket reader and writer. + let mut pending_outbound: Option<(Vec, usize)> = None; + let mut force_incoming = false; + let mut frames_drained_after_pong_deadline = 0usize; 'relay: loop { + // Consume a due tick before the always-ready record arm below can win + // another select iteration and postpone the keepalive. + if pong_watchdog.deadline().is_none() + && keepalive.tick().now_or_never().is_some() + { + if let Err(error) = send_keepalive_ping( + &mut websocket, + &mut pong_watchdog, + pong_deadline.as_mut(), + ) + .await + { + warn!("failed to write Noise relay keepalive ping: {error}"); + break; + } + frames_drained_after_pong_deadline = 0; + continue; + } + + let pong_deadline_expired = pong_watchdog + .deadline() + .is_some_and(|deadline| tokio::time::Instant::now() >= deadline); + // After expiry, inspect only frames already queued. Forcing the peeked + // item through next() makes the 32-frame grace deterministic. + if pong_deadline_expired && !force_incoming { + if frames_drained_after_pong_deadline + < MAX_FRAMES_DRAINED_AFTER_PONG_DEADLINE + && std::pin::Pin::new(&mut websocket) + .peek() + .now_or_never() + .is_some() + { + force_incoming = true; + } else { + warn!( + noise_reason = WEBSOCKET_PONG_TIMEOUT_REASON, + "Noise harness rendezvous websocket disconnected" + ); + send_disconnected( + &incoming_tx, + &disconnected_tx, + WEBSOCKET_PONG_TIMEOUT_REASON.to_string(), + ); + return; + } + } + + // While a Pong is outstanding, drain already-queued inbound traffic + // before the next fragment so a queued Pong cannot sit behind writes. + if !force_incoming + && pong_watchdog.deadline().is_some() + && pending_outbound.is_some() + && std::pin::Pin::new(&mut websocket) + .peek() + .now_or_never() + .is_some() + { + force_incoming = true; + } + tokio::select! { - maybe_message = outgoing_rx.recv() => { + maybe_message = outgoing_rx.recv(), if pending_outbound.is_none() && !force_incoming && !pong_deadline_expired => { let Some(message) = maybe_message else { break; }; - let framed = match frame_jsonrpc_message(&message) { - Ok(framed) => framed, + pending_outbound = Some(match frame_jsonrpc_message(&message) { + Ok(framed) => (framed, 0), Err(error) => { warn!("failed to frame JSON-RPC payload for Noise relay: {error}"); break; } + }); + } + _ = std::future::ready(()), if pending_outbound.is_some() && !force_incoming && !pong_deadline_expired => { + let seq = match take_next_sequence(&mut next_outbound_seq) { + Ok(seq) => seq, + Err(error) => { + warn!("Noise relay sequence exhausted: {error}"); + break 'relay; + } }; - for plaintext_record in framed.chunks(NOISE_RECORD_PLAINTEXT_LEN) { - let seq = match take_next_sequence(&mut next_outbound_seq) { - Ok(seq) => seq, - Err(error) => { - warn!("Noise relay sequence exhausted: {error}"); - break 'relay; - } + let (ciphertext, next_offset, message_complete) = { + let Some((framed, offset)) = pending_outbound.as_ref() else { + continue; }; - let ciphertext = match transport.encrypt(plaintext_record) { + let next_offset = (*offset + NOISE_RECORD_PLAINTEXT_LEN).min(framed.len()); + let ciphertext = match transport.encrypt(&framed[*offset..next_offset]) { Ok(ciphertext) => ciphertext, Err(error) => { warn!("failed to encrypt JSON-RPC payload for Noise relay: {error}"); break 'relay; } }; - let frame = RelayMessageFrame::data(stream_id.clone(), seq, ciphertext); - if let Err(error) = websocket - .send(Message::Binary(encode_relay_message_frame(&frame).into())) - .await - { - warn!("failed to write Noise relay websocket: {error}"); - break 'relay; - } + (ciphertext, next_offset, next_offset == framed.len()) + }; + let frame = RelayMessageFrame::data(stream_id.clone(), seq, ciphertext); + // A Pong can arrive after the readiness check while this write owns the + // combined sink and stream. A single bounded record can therefore hit the + // deadline and disconnect with that Pong queued. Treat that as write + // backpressure; this loop yields only between records. + if let Err(error) = send_websocket_message( + &mut websocket, + Message::Binary(encode_relay_message_frame(&frame).into()), + pong_watchdog.write_deadline(tokio::time::Instant::now()), + ) + .await + { + warn!("failed to write Noise relay websocket: {error}"); + break 'relay; + } + if message_complete { + pending_outbound = None; + } else if let Some((_framed, offset)) = pending_outbound.as_mut() { + *offset = next_offset; } } + _ = &mut pong_deadline, if pong_watchdog.deadline().is_some() && !force_incoming => { + continue; + } + _ = keepalive.tick(), if pong_watchdog.deadline().is_none() => { + if let Err(error) = send_keepalive_ping( + &mut websocket, + &mut pong_watchdog, + pong_deadline.as_mut(), + ) + .await + { + warn!("failed to write Noise relay keepalive ping: {error}"); + break; + } + frames_drained_after_pong_deadline = 0; + } incoming_message = websocket.next() => { + force_incoming = false; let Some(incoming_message) = incoming_message else { break; }; + // Count each completed read after expiry. If only the deadline arm + // advanced this counter, reads won by a simultaneously ready incoming + // arm would not count toward the 32-frame cap. + if pong_watchdog + .deadline() + .is_some_and(|deadline| tokio::time::Instant::now() >= deadline) + { + frames_drained_after_pong_deadline += 1; + } match incoming_message { Ok(Message::Binary(payload)) => { let frame = match decode_relay_message_frame(payload.as_ref()) { Ok(frame) => frame, Err(error) => { - send_malformed(&incoming_tx, error.to_string()).await; + send_malformed(&incoming_tx, error.to_string()); break; } }; @@ -325,7 +442,7 @@ where let data = match frame.into_data() { Ok(data) => data, Err(error) => { - send_malformed(&incoming_tx, error.to_string()).await; + send_malformed(&incoming_tx, error.to_string()); break; } }; @@ -334,22 +451,26 @@ where &mut transport, &mut inbound_decoder, data, + pong_watchdog.write_deadline(tokio::time::Instant::now()), &incoming_tx, ) .await { - send_malformed(&incoming_tx, error.to_string()).await; + if matches!(error, ExecServerError::Closed) { + break; + } + send_malformed(&incoming_tx, error.to_string()); break; } } Ok(RelayFrameBodyKind::Reset) => { - let _ = incoming_tx - .send(JsonRpcConnectionEvent::Disconnected { + let _ = incoming_tx.try_send( + JsonRpcConnectionEvent::Disconnected { reason: Some( NOISE_RELAY_RESET_DISCONNECT_REASON.to_string(), ), - }) - .await; + }, + ); break; } Ok( @@ -361,20 +482,22 @@ where send_malformed( &incoming_tx, "Noise relay received invalid post-handshake frame".to_string(), - ) - .await; + ); break; } } } Ok(Message::Close(_)) => break, - Ok(Message::Ping(_) | Message::Pong(_) | Message::Frame(_)) => {} + Ok(Message::Pong(_)) => { + pong_watchdog.received_pong(); + frames_drained_after_pong_deadline = 0; + } + Ok(Message::Ping(_) | Message::Frame(_)) => {} Ok(Message::Text(_)) => { send_malformed( &incoming_tx, "Noise relay transport expects binary protobuf frames".to_string(), - ) - .await; + ); break; } Err(error) => { @@ -398,6 +521,46 @@ where } } +async fn send_websocket_message( + websocket: &mut T, + message: Message, + deadline: tokio::time::Instant, +) -> Result<(), String> +where + T: Sink + Unpin, + E: std::fmt::Display, +{ + match tokio::time::timeout_at(deadline, websocket.send(message)).await { + Ok(Ok(())) => Ok(()), + Ok(Err(error)) => Err(error.to_string()), + Err(_) => Err("websocket write timed out".to_string()), + } +} + +async fn send_keepalive_ping( + websocket: &mut T, + pong_watchdog: &mut WebSocketPongWatchdog, + pong_deadline: std::pin::Pin<&mut tokio::time::Sleep>, +) -> Result<(), String> +where + T: Sink + Unpin, + E: std::fmt::Display, +{ + send_websocket_message( + websocket, + Message::Ping(Vec::new().into()), + pong_watchdog.write_deadline(tokio::time::Instant::now()), + ) + .await?; + // Start the response clock after the Ping flushes; waiting for sink capacity + // is governed by the write deadline above. + pong_watchdog.ping_sent(tokio::time::Instant::now()); + if let Some(deadline) = pong_watchdog.deadline() { + pong_deadline.reset(deadline); + } + Ok(()) +} + /// Order and decrypt one relay frame, then emit any complete JSON-RPC messages. /// Relay records and JSON-RPC messages do not share boundaries, so reassembly /// happens after decryption. @@ -406,6 +569,7 @@ async fn receive_data( transport: &mut NoiseTransport, decoder: &mut JsonRpcMessageDecoder, data: RelayData, + delivery_deadline: tokio::time::Instant, incoming_tx: &mpsc::Sender, ) -> Result<(), ExecServerError> { // Ordering must happen before decryption because Noise transport nonces are @@ -419,30 +583,50 @@ async fn receive_data( // The authenticated byte stream can carry partial or multiple JSON-RPC // messages; emit only complete, successfully parsed messages. for message in decoder.push(&plaintext)? { - incoming_tx - .send(JsonRpcConnectionEvent::Message(message)) - .await - .map_err(|_| ExecServerError::Closed)?; + send_incoming_event( + incoming_tx, + JsonRpcConnectionEvent::Message(message), + delivery_deadline, + ) + .await?; } } Ok(()) } -async fn send_malformed(incoming_tx: &mpsc::Sender, reason: String) { - let _ = incoming_tx - .send(JsonRpcConnectionEvent::MalformedMessage { reason }) - .await; +async fn send_incoming_event( + incoming_tx: &mpsc::Sender, + event: JsonRpcConnectionEvent, + deadline: tokio::time::Instant, +) -> Result<(), ExecServerError> { + match tokio::time::timeout_at(deadline, incoming_tx.send(event)).await { + Ok(Ok(())) => Ok(()), + Ok(Err(_)) => Err(ExecServerError::Closed), + Err(_) => { + warn!( + noise_reason = "application_backpressure", + "Noise harness application event delivery timed out" + ); + Err(ExecServerError::Closed) + } + } } -async fn send_disconnected( +fn send_malformed(incoming_tx: &mpsc::Sender, reason: String) { + let _ = incoming_tx.try_send(JsonRpcConnectionEvent::MalformedMessage { reason }); +} + +fn send_disconnected( incoming_tx: &mpsc::Sender, disconnected_tx: &watch::Sender, reason: String, ) { let _ = disconnected_tx.send(true); - let _ = incoming_tx - .send(JsonRpcConnectionEvent::Disconnected { - reason: Some(reason), - }) - .await; + let _ = incoming_tx.try_send(JsonRpcConnectionEvent::Disconnected { + reason: Some(reason), + }); } + +#[cfg(test)] +#[path = "harness_tests.rs"] +mod tests; diff --git a/codex-rs/exec-server/src/noise_relay/harness_tests.rs b/codex-rs/exec-server/src/noise_relay/harness_tests.rs new file mode 100644 index 0000000000..ce26a8ca8c --- /dev/null +++ b/codex-rs/exec-server/src/noise_relay/harness_tests.rs @@ -0,0 +1,442 @@ +use std::pin::Pin; +use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::task::Context as TaskContext; +use std::task::Poll; +use std::time::Duration; + +use anyhow::Context; +use anyhow::Result; +use codex_exec_server_protocol::JSONRPCMessage; +use codex_exec_server_protocol::JSONRPCRequest; +use codex_exec_server_protocol::RequestId; +use futures::Sink; +use futures::SinkExt; +use futures::StreamExt; +use futures::channel::mpsc as futures_mpsc; +use pretty_assertions::assert_eq; +use tokio::net::TcpListener; +use tokio::sync::mpsc; +use tokio::time::Instant; +use tokio::time::timeout; +use tokio_tungstenite::accept_async; +use tokio_tungstenite::connect_async; +use tokio_tungstenite::tungstenite::Message; + +use super::*; +use crate::connection::JsonRpcConnectionEvent; +use crate::noise_channel::PendingResponderHandshake; + +const ENVIRONMENT_ID: &str = "environment-1"; +const EXECUTOR_REGISTRATION_ID: &str = "registration-1"; + +#[tokio::test] +async fn fragmented_writes_yield_to_keepalive_and_queued_pong() -> Result<()> { + let (connection, mut control, mut outbound_rx) = connected_controlled_harness().await?; + + connection + .outgoing_tx + .send(JSONRPCMessage::Request(JSONRPCRequest { + id: RequestId::Integer(1), + method: "large".to_string(), + params: Some(serde_json::json!({ + "payload": "x".repeat(NOISE_RECORD_PLAINTEXT_LEN * 3), + })), + trace: None, + })) + .await?; + + control.wait_for_blocked_write(/*expected*/ 1).await?; + tokio::time::sleep(WEBSOCKET_KEEPALIVE_INTERVAL + Duration::from_millis(10)).await; + control.grant_writes(/*count*/ 1); + let first_data = read_outbound_data(&mut outbound_rx).await?; + assert_eq!(first_data.seq, 0); + + control.wait_for_blocked_write(/*expected*/ 2).await?; + control.grant_writes(/*count*/ 1); + let Message::Ping(ping_payload) = timeout(Duration::from_secs(1), outbound_rx.next()) + .await? + .context("harness closed before sending keepalive")? + else { + anyhow::bail!("expected keepalive between fragmented writes"); + }; + + control.wait_for_blocked_write(/*expected*/ 3).await?; + control.send_inbound(Message::Pong(ping_payload))?; + tokio::time::sleep(WEBSOCKET_KEEPALIVE_INTERVAL + Duration::from_millis(10)).await; + control.grant_writes(/*count*/ 1); + let second_data = read_outbound_data(&mut outbound_rx).await?; + assert_eq!(second_data.seq, 1); + + control.wait_for_blocked_write(/*expected*/ 4).await?; + control.grant_writes(/*count*/ 1); + let next_message = timeout(Duration::from_secs(1), outbound_rx.next()) + .await? + .context("harness closed after receiving queued Pong")?; + assert!(matches!(next_message, Message::Ping(_))); + + for task in &connection.task_handles { + task.abort(); + } + Ok(()) +} + +#[tokio::test(flavor = "current_thread")] +async fn post_deadline_drain_stops_before_frame_33() -> Result<()> { + let (mut connection, mut control, mut outbound_rx) = connected_controlled_harness().await?; + + control.wait_for_blocked_write(/*expected*/ 1).await?; + control.grant_writes(/*count*/ 1); + let Message::Ping(ping_payload) = timeout(Duration::from_secs(1), outbound_rx.next()) + .await? + .context("harness closed before sending keepalive")? + else { + anyhow::bail!("expected keepalive ping"); + }; + let reads_before_deadline = control.inbound_reads(); + + let unrelated_frame = + encode_relay_message_frame(&RelayMessageFrame::resume("unrelated-stream".to_string())); + for _ in 0..MAX_FRAMES_DRAINED_AFTER_PONG_DEADLINE { + control.send_inbound(Message::Binary(unrelated_frame.clone().into()))?; + } + control.send_inbound(Message::Pong(ping_payload))?; + + // Keep the current-thread runtime from consuming the queued frames until the + // Pong deadline and every frame are ready together. + std::thread::sleep(WEBSOCKET_PONG_TIMEOUT + Duration::from_millis(10)); + + let event = timeout(Duration::from_secs(1), connection.incoming_rx.recv()).await?; + let Some(JsonRpcConnectionEvent::Disconnected { reason }) = event else { + anyhow::bail!("expected Pong timeout, got {event:?}"); + }; + assert_eq!(reason.as_deref(), Some(WEBSOCKET_PONG_TIMEOUT_REASON)); + assert_eq!( + control.inbound_reads() - reads_before_deadline, + MAX_FRAMES_DRAINED_AFTER_PONG_DEADLINE + ); + Ok(()) +} + +#[tokio::test] +async fn pong_keeps_harness_alive_until_peer_stops_responding() -> Result<()> { + let listener = TcpListener::bind("127.0.0.1:0").await?; + let websocket_url = format!("ws://{}", listener.local_addr()?); + let harness_connection = tokio::spawn(connect_async(websocket_url)); + let (socket, _peer_addr) = listener.accept().await?; + let mut executor_websocket = accept_async(socket).await?; + let (harness_websocket, _response) = harness_connection.await??; + + let executor_identity = NoiseChannelIdentity::generate()?; + let mut connection = noise_harness_connection_from_websocket( + harness_websocket, + NoiseHarnessConnectionArgs { + connection_label: "test rendezvous".to_string(), + environment_id: ENVIRONMENT_ID.to_string(), + executor_registration_id: EXECUTOR_REGISTRATION_ID.to_string(), + identity: NoiseChannelIdentity::generate()?, + responder_public_key: executor_identity.public_key(), + harness_key_authorization: "authorization".to_string(), + }, + ); + + let resume_message = timeout(Duration::from_secs(1), executor_websocket.next()) + .await? + .context("harness closed before sending resume")??; + let Message::Binary(resume_payload) = resume_message else { + anyhow::bail!("expected resume frame, got {resume_message:?}"); + }; + let resume = decode_relay_message_frame(resume_payload.as_ref())?; + assert_eq!(resume.validate()?, RelayFrameBodyKind::Resume); + + let handshake_message = timeout(Duration::from_secs(1), executor_websocket.next()) + .await? + .context("harness closed before sending handshake")??; + let Message::Binary(handshake_payload) = handshake_message else { + anyhow::bail!("expected handshake frame, got {handshake_message:?}"); + }; + let handshake = decode_relay_message_frame(handshake_payload.as_ref())?; + assert_eq!(handshake.stream_id, resume.stream_id); + let stream_id = handshake.stream_id.clone(); + let prologue = + noise_channel_prologue(ENVIRONMENT_ID, EXECUTOR_REGISTRATION_ID, stream_id.as_str()); + let pending = PendingResponderHandshake::read_request( + &executor_identity, + &prologue, + &handshake.into_handshake_payload()?, + )?; + let (_transport, response) = pending.complete()?; + let response = RelayMessageFrame::handshake(stream_id, response); + executor_websocket + .send(Message::Binary( + encode_relay_message_frame(&response).into(), + )) + .await?; + + let mut pings = 0; + while pings < 6 { + let message = timeout(Duration::from_secs(1), executor_websocket.next()) + .await? + .context("harness disconnected before six keepalive pings")??; + match message { + Message::Ping(payload) => { + executor_websocket.send(Message::Pong(payload)).await?; + pings += 1; + } + Message::Pong(_) | Message::Frame(_) => {} + message => anyhow::bail!("expected keepalive ping, got {message:?}"), + } + } + + // Keep non-Pong traffic flowing after responses stop. It must not defeat + // the bounded grace for a Pong already queued behind data. + let unrelated_frame = + encode_relay_message_frame(&RelayMessageFrame::resume("unrelated-stream".to_string())); + let traffic_task = tokio::spawn(async move { + loop { + if executor_websocket + .send(Message::Binary(unrelated_frame.clone().into())) + .await + .is_err() + { + break; + } + tokio::time::sleep(Duration::from_millis(5)).await; + } + }); + let event = timeout(Duration::from_secs(1), connection.incoming_rx.recv()).await?; + traffic_task.abort(); + let _ = traffic_task.await; + let Some(JsonRpcConnectionEvent::Disconnected { reason }) = event else { + anyhow::bail!("expected pong timeout, got {event:?}"); + }; + assert_eq!(reason.as_deref(), Some(WEBSOCKET_PONG_TIMEOUT_REASON)); + Ok(()) +} + +#[tokio::test] +async fn application_event_delivery_is_bounded() -> Result<()> { + let (incoming_tx, _incoming_rx) = mpsc::channel(1); + incoming_tx + .send(JsonRpcConnectionEvent::MalformedMessage { + reason: "fill queue".to_string(), + }) + .await?; + + let result = send_incoming_event( + &incoming_tx, + JsonRpcConnectionEvent::MalformedMessage { + reason: "blocked event".to_string(), + }, + Instant::now() + Duration::from_millis(10), + ) + .await; + + assert!(matches!(result, Err(ExecServerError::Closed))); + Ok(()) +} + +async fn read_outbound_data( + outbound_rx: &mut futures_mpsc::UnboundedReceiver, +) -> Result { + let Message::Binary(payload) = timeout(Duration::from_secs(1), outbound_rx.next()) + .await? + .context("harness closed before sending data")? + else { + anyhow::bail!("expected relay data frame"); + }; + let frame = decode_relay_message_frame(payload.as_ref())?; + assert_eq!(frame.validate()?, RelayFrameBodyKind::Data); + frame.into_data().map_err(anyhow::Error::from) +} + +async fn connected_controlled_harness() -> Result<( + JsonRpcConnection, + ControlledWebSocketHandle, + futures_mpsc::UnboundedReceiver, +)> { + let (websocket, control, mut outbound_rx) = ControlledWebSocket::new(/*write_permits*/ 2); + let executor_identity = NoiseChannelIdentity::generate()?; + let connection = noise_harness_connection_from_websocket( + websocket, + NoiseHarnessConnectionArgs { + connection_label: "test rendezvous".to_string(), + environment_id: ENVIRONMENT_ID.to_string(), + executor_registration_id: EXECUTOR_REGISTRATION_ID.to_string(), + identity: NoiseChannelIdentity::generate()?, + responder_public_key: executor_identity.public_key(), + harness_key_authorization: "authorization".to_string(), + }, + ); + + let Message::Binary(resume_payload) = timeout(Duration::from_secs(1), outbound_rx.next()) + .await? + .context("harness closed before sending resume")? + else { + anyhow::bail!("expected resume frame"); + }; + let resume = decode_relay_message_frame(resume_payload.as_ref())?; + let Message::Binary(handshake_payload) = timeout(Duration::from_secs(1), outbound_rx.next()) + .await? + .context("harness closed before sending handshake")? + else { + anyhow::bail!("expected handshake frame"); + }; + let handshake = decode_relay_message_frame(handshake_payload.as_ref())?; + let stream_id = handshake.stream_id.clone(); + assert_eq!(stream_id, resume.stream_id); + let prologue = + noise_channel_prologue(ENVIRONMENT_ID, EXECUTOR_REGISTRATION_ID, stream_id.as_str()); + let pending = PendingResponderHandshake::read_request( + &executor_identity, + &prologue, + &handshake.into_handshake_payload()?, + )?; + let (_transport, response) = pending.complete()?; + control.send_inbound(Message::Binary( + encode_relay_message_frame(&RelayMessageFrame::handshake(stream_id, response)).into(), + ))?; + Ok((connection, control, outbound_rx)) +} + +struct ControlledWebSocket { + inbound_rx: futures_mpsc::UnboundedReceiver>, + outbound_tx: futures_mpsc::UnboundedSender, + write_permit_rx: futures_mpsc::UnboundedReceiver<()>, + blocked_write_tx: futures_mpsc::UnboundedSender, + write_waiting: bool, + blocked_writes: usize, + inbound_reads: Arc, +} + +struct ControlledWebSocketHandle { + inbound_tx: futures_mpsc::UnboundedSender>, + write_permit_tx: futures_mpsc::UnboundedSender<()>, + blocked_write_rx: futures_mpsc::UnboundedReceiver, + inbound_reads: Arc, +} + +impl ControlledWebSocket { + fn new( + write_permits: usize, + ) -> ( + Self, + ControlledWebSocketHandle, + futures_mpsc::UnboundedReceiver, + ) { + let (inbound_tx, inbound_rx) = futures_mpsc::unbounded(); + let (outbound_tx, outbound_rx) = futures_mpsc::unbounded(); + let (write_permit_tx, write_permit_rx) = futures_mpsc::unbounded(); + let (blocked_write_tx, blocked_write_rx) = futures_mpsc::unbounded(); + for _ in 0..write_permits { + write_permit_tx + .unbounded_send(()) + .expect("test write permit receiver should stay open"); + } + let inbound_reads = Arc::new(AtomicUsize::new(0)); + ( + Self { + inbound_rx, + outbound_tx, + write_permit_rx, + blocked_write_tx, + write_waiting: false, + blocked_writes: 0, + inbound_reads: Arc::clone(&inbound_reads), + }, + ControlledWebSocketHandle { + inbound_tx, + write_permit_tx, + blocked_write_rx, + inbound_reads, + }, + outbound_rx, + ) + } +} + +impl ControlledWebSocketHandle { + fn send_inbound(&self, message: Message) -> Result<()> { + self.inbound_tx + .unbounded_send(Ok(message)) + .map_err(anyhow::Error::from) + } + + fn grant_writes(&self, count: usize) { + for _ in 0..count { + self.write_permit_tx + .unbounded_send(()) + .expect("test write permit receiver should stay open"); + } + } + + fn inbound_reads(&self) -> usize { + self.inbound_reads.load(Ordering::Acquire) + } + + async fn wait_for_blocked_write(&mut self, expected: usize) -> Result<()> { + let actual = timeout(Duration::from_secs(1), self.blocked_write_rx.next()) + .await? + .context("websocket closed before blocking the expected write")?; + assert_eq!(actual, expected); + Ok(()) + } +} + +impl Sink for ControlledWebSocket { + type Error = std::convert::Infallible; + + fn poll_ready(self: Pin<&mut Self>, cx: &mut TaskContext<'_>) -> Poll> { + let this = self.get_mut(); + match Pin::new(&mut this.write_permit_rx).poll_next(cx) { + Poll::Ready(Some(())) => { + this.write_waiting = false; + Poll::Ready(Ok(())) + } + Poll::Ready(None) | Poll::Pending => { + if !this.write_waiting { + this.write_waiting = true; + this.blocked_writes += 1; + this.blocked_write_tx + .unbounded_send(this.blocked_writes) + .expect("test blocked-write receiver should stay open"); + } + Poll::Pending + } + } + } + + fn start_send(self: Pin<&mut Self>, item: Message) -> Result<(), Self::Error> { + self.outbound_tx + .unbounded_send(item) + .expect("test outbound receiver should stay open"); + Ok(()) + } + + fn poll_flush( + self: Pin<&mut Self>, + _cx: &mut TaskContext<'_>, + ) -> Poll> { + Poll::Ready(Ok(())) + } + + fn poll_close( + self: Pin<&mut Self>, + _cx: &mut TaskContext<'_>, + ) -> Poll> { + Poll::Ready(Ok(())) + } +} + +impl futures::Stream for ControlledWebSocket { + type Item = Result; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut TaskContext<'_>) -> Poll> { + let result = Pin::new(&mut self.inbound_rx).poll_next(cx); + if matches!(result, Poll::Ready(Some(_))) { + self.inbound_reads.fetch_add(1, Ordering::Release); + } + result + } +} diff --git a/codex-rs/exec-server/src/relay.rs b/codex-rs/exec-server/src/relay.rs index d7e3d9c2a7..b8d477f5db 100644 --- a/codex-rs/exec-server/src/relay.rs +++ b/codex-rs/exec-server/src/relay.rs @@ -41,6 +41,9 @@ use crate::relay_proto::RelayReset; use crate::relay_proto::RelayResume; use crate::relay_proto::relay_message_frame; use crate::server::ConnectionProcessor; +use crate::websocket_pong_watchdog::WEBSOCKET_PONG_TIMEOUT; +use crate::websocket_pong_watchdog::WEBSOCKET_PONG_TIMEOUT_REASON; +use crate::websocket_pong_watchdog::WebSocketPongWatchdog; const RELAY_MESSAGE_FRAME_VERSION: u32 = 1; const MAX_ACTIVE_NOISE_RELAY_STREAMS: usize = 128; @@ -49,6 +52,27 @@ const MAX_HARNESS_KEY_AUTHORIZATION_BYTES: usize = 4096; const MAX_PENDING_HANDSHAKE_VALIDATIONS: usize = 32; const HARNESS_KEY_VALIDATION_TIMEOUT: Duration = Duration::from_secs(10); +#[derive(Debug, Clone, Copy, Eq, PartialEq)] +pub(crate) enum RendezvousDisconnectReason { + PeerClose, + ReadError, + WriteError, + PongTimeout, + LocalShutdown, +} + +impl RendezvousDisconnectReason { + pub(crate) fn as_str(self) -> &'static str { + match self { + Self::PeerClose => "peer_close", + Self::ReadError => "read_error", + Self::WriteError => "write_error", + Self::PongTimeout => WEBSOCKET_PONG_TIMEOUT_REASON, + Self::LocalShutdown => "local_shutdown", + } + } +} + #[derive(Debug, Clone, Copy, Eq, PartialEq)] pub(crate) enum RelayFrameBodyKind { Data, @@ -448,7 +472,8 @@ pub(crate) async fn run_multiplexed_environment( executor_registration_id: String, identity: NoiseChannelIdentity, validator: V, -) where +) -> RendezvousDisconnectReason +where S: AsyncRead + AsyncWrite + Unpin + Send + 'static, V: HarnessKeyValidator + Clone + 'static, { @@ -461,6 +486,7 @@ pub(crate) async fn run_multiplexed_environment( mpsc::channel::>(CHANNEL_CAPACITY); let (closed_stream_tx, mut closed_stream_rx) = mpsc::channel::(MAX_ACTIVE_NOISE_RELAY_STREAMS); + let (pong_tx, mut pong_rx) = mpsc::channel(1); // Use a separate writer so this loop never waits on the channel it drains. let mut physical_writer_task = tokio::spawn(async move { let mut keepalive = tokio::time::interval_at( @@ -468,19 +494,61 @@ pub(crate) async fn run_multiplexed_environment( WEBSOCKET_KEEPALIVE_INTERVAL, ); keepalive.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + let mut pong_watchdog = WebSocketPongWatchdog::new(WEBSOCKET_PONG_TIMEOUT); + let pong_deadline = tokio::time::sleep(WEBSOCKET_PONG_TIMEOUT); + tokio::pin!(pong_deadline); loop { let message = tokio::select! { + pong = pong_rx.recv() => { + let Some(()) = pong else { + break RendezvousDisconnectReason::LocalShutdown; + }; + pong_watchdog.received_pong(); + continue; + } + _ = &mut pong_deadline, if pong_watchdog.deadline().is_some() => { + match pong_rx.try_recv() { + Ok(()) => { + pong_watchdog.received_pong(); + continue; + } + Err(tokio::sync::mpsc::error::TryRecvError::Empty) => { + break RendezvousDisconnectReason::PongTimeout; + } + Err(tokio::sync::mpsc::error::TryRecvError::Disconnected) => { + break RendezvousDisconnectReason::LocalShutdown; + } + } + } + _ = keepalive.tick(), if pong_watchdog.deadline().is_none() => { + Message::Ping(Vec::new().into()) + } encoded = physical_outgoing_rx.recv() => { let Some(encoded) = encoded else { - break; + break RendezvousDisconnectReason::LocalShutdown; }; Message::Binary(encoded.into()) } - _ = keepalive.tick() => Message::Ping(Vec::new().into()), }; - if let Err(error) = websocket_sink.send(message).await { - warn!("Noise multiplexed environment websocket write failed: {error}"); - break; + let is_keepalive_ping = matches!(message, Message::Ping(_)); + let write_deadline = pong_watchdog.write_deadline(tokio::time::Instant::now()); + match tokio::time::timeout_at(write_deadline, websocket_sink.send(message)).await { + Ok(Ok(())) => { + if is_keepalive_ping { + pong_watchdog.ping_sent(tokio::time::Instant::now()); + if let Some(deadline) = pong_watchdog.deadline() { + pong_deadline.as_mut().reset(deadline); + } + } + } + Ok(Err(error)) => { + warn!("Noise multiplexed environment websocket write failed: {error}"); + break RendezvousDisconnectReason::WriteError; + } + Err(_) => { + warn!("Noise multiplexed environment websocket write timed out"); + break RendezvousDisconnectReason::WriteError; + } } } }); @@ -489,13 +557,18 @@ pub(crate) async fn run_multiplexed_environment( let mut validation_tasks: JoinSet = JoinSet::new(); let mut failed_handshakes = 0usize; let mut next_validation_id = 0u64; + let mut disconnect_reason = RendezvousDisconnectReason::LocalShutdown; loop { // Registry calls run separately so a slow check does not block the relay. let frame = tokio::select! { writer_result = &mut physical_writer_task => { - if let Err(error) = writer_result { - warn!("Noise multiplexed environment websocket writer failed: {error}"); + match writer_result { + Ok(reason) => disconnect_reason = reason, + Err(error) => { + warn!("Noise multiplexed environment websocket writer failed: {error}"); + disconnect_reason = RendezvousDisconnectReason::LocalShutdown; + } } break; } @@ -621,14 +694,22 @@ pub(crate) async fn run_multiplexed_environment( continue; } }, - Some(Ok(Message::Close(_))) | None => break, - Some(Ok(Message::Ping(_) | Message::Pong(_) | Message::Frame(_))) => continue, + Some(Ok(Message::Close(_))) | None => { + disconnect_reason = RendezvousDisconnectReason::PeerClose; + break; + } + Some(Ok(Message::Pong(_))) => { + let _ = pong_tx.try_send(()); + continue; + } + Some(Ok(Message::Ping(_) | Message::Frame(_))) => continue, Some(Ok(Message::Text(_))) => { warn!("dropping non-binary Noise relay frame from harness"); continue; } Some(Err(error)) => { debug!("Noise multiplexed environment websocket read failed: {error}"); + disconnect_reason = RendezvousDisconnectReason::ReadError; break; } } @@ -802,6 +883,7 @@ pub(crate) async fn run_multiplexed_environment( physical_writer_task.abort(); let _ = physical_writer_task.await; } + disconnect_reason } /// Charge one failed authenticated-channel attempt to this physical relay. diff --git a/codex-rs/exec-server/src/relay_noise_tests.rs b/codex-rs/exec-server/src/relay_noise_tests.rs index 0020efe347..704c48d9a1 100644 --- a/codex-rs/exec-server/src/relay_noise_tests.rs +++ b/codex-rs/exec-server/src/relay_noise_tests.rs @@ -17,6 +17,7 @@ use tokio_tungstenite::tungstenite::Message; use super::HarnessKeyValidator; use super::MAX_FAILED_NOISE_HANDSHAKES; use super::MAX_HARNESS_KEY_AUTHORIZATION_BYTES; +use super::RendezvousDisconnectReason; use super::run_multiplexed_environment; use crate::ExecServerError; use crate::ExecServerRuntimePaths; @@ -33,6 +34,86 @@ use crate::server::ConnectionProcessor; const ENVIRONMENT_ID: &str = "environment-1"; const EXECUTOR_REGISTRATION_ID: &str = "registration-1"; +#[tokio::test] +async fn missing_pong_disconnects_physical_relay() -> Result<()> { + let listener = TcpListener::bind("127.0.0.1:0").await?; + let websocket_url = format!("ws://{}", listener.local_addr()?); + let harness_connection = tokio::spawn(connect_async(websocket_url)); + let (socket, _peer_addr) = listener.accept().await?; + let environment_websocket = accept_async(socket).await?; + let (_harness_websocket, _response) = harness_connection.await??; + + let environment_task = tokio::spawn(run_multiplexed_environment( + environment_websocket, + ConnectionProcessor::new(ExecServerRuntimePaths::new( + std::env::current_exe()?, + /*codex_linux_sandbox_exe*/ None, + )?), + ENVIRONMENT_ID.to_string(), + EXECUTOR_REGISTRATION_ID.to_string(), + NoiseChannelIdentity::generate()?, + BlockingValidator { + calls: Arc::new(AtomicUsize::new(0)), + release: Arc::new(Notify::new()), + }, + )); + + assert_eq!( + timeout(Duration::from_secs(1), environment_task).await??, + RendezvousDisconnectReason::PongTimeout + ); + Ok(()) +} + +#[tokio::test] +async fn pong_keeps_physical_relay_connected() -> Result<()> { + let listener = TcpListener::bind("127.0.0.1:0").await?; + let websocket_url = format!("ws://{}", listener.local_addr()?); + let harness_connection = tokio::spawn(connect_async(websocket_url)); + let (socket, _peer_addr) = listener.accept().await?; + let environment_websocket = accept_async(socket).await?; + let (mut harness_websocket, _response) = harness_connection.await??; + + let environment_task = tokio::spawn(run_multiplexed_environment( + environment_websocket, + ConnectionProcessor::new(ExecServerRuntimePaths::new( + std::env::current_exe()?, + /*codex_linux_sandbox_exe*/ None, + )?), + ENVIRONMENT_ID.to_string(), + EXECUTOR_REGISTRATION_ID.to_string(), + NoiseChannelIdentity::generate()?, + BlockingValidator { + calls: Arc::new(AtomicUsize::new(0)), + release: Arc::new(Notify::new()), + }, + )); + + timeout(Duration::from_secs(1), async { + let mut pings = 0; + while pings < 6 { + match harness_websocket.next().await { + Some(Ok(Message::Ping(payload))) => { + harness_websocket.send(Message::Pong(payload)).await?; + pings += 1; + } + Some(Ok(Message::Pong(_) | Message::Frame(_))) => {} + Some(Ok(message)) => anyhow::bail!("expected keepalive ping, got {message:?}"), + Some(Err(error)) => return Err(error.into()), + None => anyhow::bail!("environment disconnected before six keepalive pings"), + } + } + Ok::<_, anyhow::Error>(()) + }) + .await??; + harness_websocket.close(None).await?; + assert_eq!( + timeout(Duration::from_secs(1), environment_task).await??, + RendezvousDisconnectReason::PeerClose + ); + Ok(()) +} + #[derive(Clone)] struct BlockingValidator { calls: Arc, diff --git a/codex-rs/exec-server/src/remote.rs b/codex-rs/exec-server/src/remote.rs index 58fcd8c654..8870e0efe9 100644 --- a/codex-rs/exec-server/src/remote.rs +++ b/codex-rs/exec-server/src/remote.rs @@ -492,7 +492,7 @@ pub async fn run_remote_environment( noise_outcome = "ok", "Noise executor connected to rendezvous" ); - run_multiplexed_environment( + let disconnect_reason = run_multiplexed_environment( websocket, processor.clone(), response.environment_id.clone(), @@ -505,7 +505,15 @@ pub async fn run_remote_environment( }, ) .await; - config.telemetry.remote_reconnect("disconnected"); + info!( + noise_event = "rendezvous_connection", + noise_outcome = "disconnected", + noise_reason = disconnect_reason.as_str(), + "Noise executor disconnected from rendezvous" + ); + config + .telemetry + .remote_reconnect(disconnect_reason.as_str()); } Err(error) => { let registration_rejected = matches!( diff --git a/codex-rs/exec-server/src/websocket_pong_watchdog.rs b/codex-rs/exec-server/src/websocket_pong_watchdog.rs new file mode 100644 index 0000000000..b656772b3d --- /dev/null +++ b/codex-rs/exec-server/src/websocket_pong_watchdog.rs @@ -0,0 +1,44 @@ +use std::time::Duration; + +use tokio::time::Instant; + +#[cfg(test)] +pub(crate) const WEBSOCKET_PONG_TIMEOUT: Duration = Duration::from_millis(100); +#[cfg(not(test))] +pub(crate) const WEBSOCKET_PONG_TIMEOUT: Duration = Duration::from_secs(60); +pub(crate) const WEBSOCKET_PONG_TIMEOUT_REASON: &str = "pong_timeout"; + +/// Tracks whether a WebSocket peer has acknowledged a keepalive ping. +pub(crate) struct WebSocketPongWatchdog { + timeout: Duration, + deadline: Option, +} + +impl WebSocketPongWatchdog { + pub(crate) fn new(timeout: Duration) -> Self { + Self { + timeout, + deadline: None, + } + } + + pub(crate) fn ping_sent(&mut self, now: Instant) { + self.deadline.get_or_insert(now + self.timeout); + } + + pub(crate) fn deadline(&self) -> Option { + self.deadline + } + + pub(crate) fn write_deadline(&self, now: Instant) -> Instant { + self.deadline.unwrap_or(now + self.timeout) + } + + pub(crate) fn received_pong(&mut self) { + self.deadline = None; + } +} + +#[cfg(test)] +#[path = "websocket_pong_watchdog_tests.rs"] +mod tests; diff --git a/codex-rs/exec-server/src/websocket_pong_watchdog_tests.rs b/codex-rs/exec-server/src/websocket_pong_watchdog_tests.rs new file mode 100644 index 0000000000..2b15a0d952 --- /dev/null +++ b/codex-rs/exec-server/src/websocket_pong_watchdog_tests.rs @@ -0,0 +1,34 @@ +use std::time::Duration; + +use pretty_assertions::assert_eq; +use tokio::time::Instant; + +use super::WebSocketPongWatchdog; + +#[test] +fn repeated_ping_does_not_extend_deadline() { + let started_at = Instant::now(); + let timeout = Duration::from_secs(2); + let mut watchdog = WebSocketPongWatchdog::new(timeout); + + watchdog.ping_sent(started_at); + watchdog.ping_sent(started_at + Duration::from_secs(1)); + assert_eq!( + watchdog.write_deadline(started_at + Duration::from_secs(1)), + started_at + timeout + ); + assert_eq!(watchdog.deadline(), Some(started_at + timeout)); +} + +#[test] +fn pong_starts_a_fresh_deadline() { + let started_at = Instant::now(); + let timeout = Duration::from_secs(2); + let mut watchdog = WebSocketPongWatchdog::new(timeout); + + watchdog.ping_sent(started_at); + watchdog.received_pong(); + assert_eq!(watchdog.deadline(), None); + watchdog.ping_sent(started_at + timeout); + assert_eq!(watchdog.deadline(), Some(started_at + timeout + timeout)); +} From a98a21798c3301cfbeb6c323d6c0f6a804e08a57 Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Wed, 1 Jul 2026 17:08:11 -0700 Subject: [PATCH 2/3] Consolidate multi-agent v2 communication sends (#30867) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Why Multi-agent v2 communications currently use separate outbound paths: direct messages, follow-up tasks, and completion results go through `send_inter_agent_communication`, while a spawn's initial message goes through the generic input submission path. That split makes it difficult to add complete communication lifecycle logging in one place. This refactor makes `submit_inter_agent_communication` the common sink for those paths, preparing the follow-up observability work discussed in [#30516](https://github.com/openai/codex/pull/30516). ## What changed - Routed all current outbound `InterAgentCommunication` paths in `AgentControl`—direct messages, follow-up tasks, completion results, and multi-agent v2 spawn initial messages—through `submit_inter_agent_communication`. - Centralized the actual submission and last-task-message bookkeeping there, providing one place for the follow-up PR to instrument communication creation and successful enqueue. - Left non-communication input handling and the multi-agent v1 spawn flow unchanged. ## Testing - `just test -p codex-core 'agent::control::tests::'` (51 passed) - `just test -p codex-core 'suite::subagent_notifications::encrypted_multi_agent_v2_spawn_sends_agent_message_to_child'` (passed) --- [//]: # (BEGIN SAPLING FOOTER) Stack created with [Sapling](https://sapling-scm.com). Best reviewed with [ReviewStack](https://reviewstack.dev/openai/codex/pull/30867). * #30872 * __->__ #30867 --- codex-rs/core/src/agent/control.rs | 34 +++++++++++++++++++++--------- 1 file changed, 24 insertions(+), 10 deletions(-) diff --git a/codex-rs/core/src/agent/control.rs b/codex-rs/core/src/agent/control.rs index 49bbab0f87..8de6443402 100644 --- a/codex-rs/core/src/agent/control.rs +++ b/codex-rs/core/src/agent/control.rs @@ -154,12 +154,13 @@ impl AgentControl { state: &Arc, initial_operation: Op, ) -> CodexResult { - let last_task_message = match &initial_operation { - Op::InterAgentCommunication { communication } => { - last_task_message_from_communication(communication) - } - _ => non_empty_task_message(render_input_preview(&initial_operation)), - }; + if let Op::InterAgentCommunication { communication } = initial_operation { + return self + .submit_inter_agent_communication(agent_id, state, communication) + .await; + } + + let last_task_message = non_empty_task_message(render_input_preview(&initial_operation)); let result = self .handle_thread_request_result( agent_id, @@ -182,13 +183,26 @@ impl AgentControl { &self, agent_id: ThreadId, communication: InterAgentCommunication, + ) -> CodexResult { + self.send_input(agent_id, Op::InterAgentCommunication { communication }) + .await + } + + async fn submit_inter_agent_communication( + &self, + agent_id: ThreadId, + state: &Arc, + communication: InterAgentCommunication, ) -> CodexResult { let last_task_message = last_task_message_from_communication(&communication); - let state = self.upgrade()?; - let op = Op::InterAgentCommunication { communication }; - self.ensure_execution_capacity_for_op(agent_id, &op).await?; let result = self - .handle_thread_request_result(agent_id, &state, state.send_op(agent_id, op).await) + .handle_thread_request_result( + agent_id, + state, + state + .send_op(agent_id, Op::InterAgentCommunication { communication }) + .await, + ) .await; if result.is_ok() { match last_task_message { From 64280aa283f91491bbcb9e622b1a554ced7d2f3f Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Wed, 1 Jul 2026 17:08:27 -0700 Subject: [PATCH 3/3] Log multi-agent communication lifecycle --- codex-rs/core/src/agent/control.rs | 97 ++++++++++++------- codex-rs/core/src/agent/control/execution.rs | 11 ++- codex-rs/core/src/agent/control/spawn.rs | 62 ++++++++++-- codex-rs/core/src/agent/control_tests.rs | 32 +++--- codex-rs/core/src/agent_communication.rs | 79 +++++++++++++++ codex-rs/core/src/lib.rs | 1 + codex-rs/core/src/session/handlers.rs | 1 + codex-rs/core/src/session/mod.rs | 6 +- .../core/src/tools/handlers/agent_jobs.rs | 2 +- .../src/tools/handlers/multi_agents_common.rs | 8 +- .../src/tools/handlers/multi_agents_tests.rs | 18 ++-- .../src/tools/handlers/multi_agents_v2.rs | 1 - .../handlers/multi_agents_v2/message_tool.rs | 11 ++- .../tools/handlers/multi_agents_v2/spawn.rs | 54 +++++------ .../tests/suite/subagent_notifications.rs | 56 ++++++++++- codex-rs/feedback/src/lib.rs | 29 +++++- codex-rs/state/src/log_db.rs | 1 + codex-rs/state/src/log_db_filter_tests.rs | 12 +-- 18 files changed, 367 insertions(+), 114 deletions(-) create mode 100644 codex-rs/core/src/agent_communication.rs diff --git a/codex-rs/core/src/agent/control.rs b/codex-rs/core/src/agent/control.rs index 8de6443402..7b13c0257d 100644 --- a/codex-rs/core/src/agent/control.rs +++ b/codex-rs/core/src/agent/control.rs @@ -4,6 +4,8 @@ use crate::agent::registry::AgentRegistry; use crate::agent::role::DEFAULT_ROLE_NAME; use crate::agent::role::resolve_role_config; use crate::agent::status::is_final; +use crate::agent_communication::AgentCommunicationContext; +use crate::agent_communication::AgentCommunicationKind; use crate::codex_thread::ThreadConfigSnapshot; use crate::config::Config; use crate::config::RolloutBudgetConfig; @@ -139,12 +141,12 @@ impl AgentControl { pub(crate) async fn send_input( &self, agent_id: ThreadId, - initial_operation: Op, + input: Vec, ) -> CodexResult { let state = self.upgrade()?; - self.ensure_execution_capacity_for_op(agent_id, &initial_operation) + self.ensure_execution_capacity_for_turn_start(agent_id, true) .await?; - self.send_input_after_capacity_check(agent_id, &state, initial_operation) + self.send_input_after_capacity_check(agent_id, &state, input) .await } @@ -152,20 +154,14 @@ impl AgentControl { &self, agent_id: ThreadId, state: &Arc, - initial_operation: Op, + input: Vec, ) -> CodexResult { - if let Op::InterAgentCommunication { communication } = initial_operation { - return self - .submit_inter_agent_communication(agent_id, state, communication) - .await; - } - - let last_task_message = non_empty_task_message(render_input_preview(&initial_operation)); + let last_task_message = non_empty_task_message(render_input_preview(&input)); let result = self .handle_thread_request_result( agent_id, state, - state.send_op(agent_id, initial_operation).await, + state.send_op(agent_id, input.into()).await, ) .await; if result.is_ok() { @@ -183,8 +179,28 @@ impl AgentControl { &self, agent_id: ThreadId, communication: InterAgentCommunication, + agent_communication_context: AgentCommunicationContext, ) -> CodexResult { - self.send_input(agent_id, Op::InterAgentCommunication { communication }) + let state = self.upgrade()?; + self.ensure_execution_capacity_for_turn_start(agent_id, communication.trigger_turn) + .await?; + self.send_inter_agent_communication_after_capacity_check( + agent_id, + &state, + communication, + agent_communication_context, + ) + .await + } + + async fn send_inter_agent_communication_after_capacity_check( + &self, + agent_id: ThreadId, + state: &Arc, + communication: InterAgentCommunication, + context: AgentCommunicationContext, + ) -> CodexResult { + self.submit_inter_agent_communication(agent_id, state, communication, context) .await } @@ -193,8 +209,11 @@ impl AgentControl { agent_id: ThreadId, state: &Arc, communication: InterAgentCommunication, + context: AgentCommunicationContext, ) -> CodexResult { let last_task_message = last_task_message_from_communication(&communication); + let communication_for_log = + crate::agent_communication::logging_enabled().then(|| communication.clone()); let result = self .handle_thread_request_result( agent_id, @@ -204,6 +223,16 @@ impl AgentControl { .await, ) .await; + if let (Some(communication), Ok(communication_id)) = + (communication_for_log, result.as_ref()) + { + crate::agent_communication::emit_agent_communication_send( + communication_id, + &context, + &communication, + agent_id, + ); + } if result.is_ok() { match last_task_message { Some(last_task_message) => self @@ -494,8 +523,10 @@ impl AgentControl { message, /*trigger_turn*/ false, ); + let context = + AgentCommunicationContext::new(AgentCommunicationKind::Result, child_thread_id); let _ = control - .send_inter_agent_communication(parent_thread_id, communication) + .send_inter_agent_communication(parent_thread_id, communication, context) .await; return; } @@ -720,27 +751,23 @@ fn agent_matches_prefix(agent_path: Option<&AgentPath>, prefix: &AgentPath) -> b }) } -pub(crate) fn render_input_preview(initial_operation: &Op) -> String { - match initial_operation { - Op::UserInput { items, .. } => items - .iter() - .map(|item| match item { - UserInput::Text { text, .. } => text.clone(), - UserInput::Image { .. } => "[image]".to_string(), - UserInput::LocalImage { path, .. } => { - format!("[local_image:{}]", path.display()) - } - UserInput::Skill { name, path, .. } => { - format!("[skill:${name}]({})", path.display()) - } - UserInput::Mention { name, path, .. } => format!("[mention:${name}]({path})"), - _ => "[input]".to_string(), - }) - .collect::>() - .join("\n"), - Op::InterAgentCommunication { communication } => communication.content.clone(), - _ => String::new(), - } +pub(crate) fn render_input_preview(input: &[UserInput]) -> String { + input + .iter() + .map(|item| match item { + UserInput::Text { text, .. } => text.clone(), + UserInput::Image { .. } => "[image]".to_string(), + UserInput::LocalImage { path, .. } => { + format!("[local_image:{}]", path.display()) + } + UserInput::Skill { name, path, .. } => { + format!("[skill:${name}]({})", path.display()) + } + UserInput::Mention { name, path, .. } => format!("[mention:${name}]({path})"), + _ => "[input]".to_string(), + }) + .collect::>() + .join("\n") } fn last_task_message_from_communication(communication: &InterAgentCommunication) -> Option { diff --git a/codex-rs/core/src/agent/control/execution.rs b/codex-rs/core/src/agent/control/execution.rs index 42aaa50e02..d8559c1499 100644 --- a/codex-rs/core/src/agent/control/execution.rs +++ b/codex-rs/core/src/agent/control/execution.rs @@ -32,7 +32,16 @@ impl AgentControl { thread_id: ThreadId, op: &Op, ) -> CodexResult<()> { - if !op_starts_turn(op) { + self.ensure_execution_capacity_for_turn_start(thread_id, op_starts_turn(op)) + .await + } + + pub(super) async fn ensure_execution_capacity_for_turn_start( + &self, + thread_id: ThreadId, + starts_turn: bool, + ) -> CodexResult<()> { + if !starts_turn { return Ok(()); } let state = self.upgrade()?; diff --git a/codex-rs/core/src/agent/control/spawn.rs b/codex-rs/core/src/agent/control/spawn.rs index c4c40577ca..9078434f1f 100644 --- a/codex-rs/core/src/agent/control/spawn.rs +++ b/codex-rs/core/src/agent/control/spawn.rs @@ -9,6 +9,17 @@ struct SpawnAgentThreadInheritance { exec_policy: Option>, } +/// Initial input delivered after a spawned agent acquires execution capacity. +/// +/// V2 communication spawns keep the communication and its context paired so centralized +/// submission and lifecycle logging cannot receive one without the other. Other spawn sources +/// provide user input directly, making an uncontextualized inter-agent communication +/// unrepresentable. +enum SpawnInitialInput { + UserInput(Vec), + InterAgentCommunication(InterAgentCommunication, AgentCommunicationContext), +} + fn default_agent_nickname_list() -> Vec<&'static str> { AGENT_NAMES .lines() @@ -89,12 +100,12 @@ impl AgentControl { pub(crate) async fn spawn_agent( &self, config: Config, - initial_operation: Op, + initial_input: Vec, session_source: Option, ) -> CodexResult { let spawned_agent = Box::pin(self.spawn_agent_internal( config, - initial_operation, + SpawnInitialInput::UserInput(initial_input), session_source, SpawnAgentOptions::default(), )) @@ -106,12 +117,34 @@ impl AgentControl { pub(crate) async fn spawn_agent_with_metadata( &self, config: Config, - initial_operation: Op, + initial_input: Vec, session_source: Option, options: SpawnAgentOptions, // TODO(jif) drop with new fork. ) -> CodexResult { - Box::pin(self.spawn_agent_internal(config, initial_operation, session_source, options)) - .await + Box::pin(self.spawn_agent_internal( + config, + SpawnInitialInput::UserInput(initial_input), + session_source, + options, + )) + .await + } + + pub(crate) async fn spawn_agent_with_communication( + &self, + config: Config, + communication: InterAgentCommunication, + context: AgentCommunicationContext, + session_source: Option, + options: SpawnAgentOptions, + ) -> CodexResult { + Box::pin(self.spawn_agent_internal( + config, + SpawnInitialInput::InterAgentCommunication(communication, context), + session_source, + options, + )) + .await } pub(crate) async fn ensure_v2_agent_loaded( @@ -197,7 +230,7 @@ impl AgentControl { async fn spawn_agent_internal( &self, config: Config, - initial_operation: Op, + initial_input: SpawnInitialInput, session_source: Option, options: SpawnAgentOptions, ) -> CodexResult { @@ -356,8 +389,21 @@ impl AgentControl { ) .await; - self.send_input_after_capacity_check(new_thread.thread_id, &state, initial_operation) - .await?; + match initial_input { + SpawnInitialInput::UserInput(input) => { + self.send_input_after_capacity_check(new_thread.thread_id, &state, input) + .await?; + } + SpawnInitialInput::InterAgentCommunication(communication, context) => { + self.send_inter_agent_communication_after_capacity_check( + new_thread.thread_id, + &state, + communication, + context, + ) + .await?; + } + } if multi_agent_version != MultiAgentVersion::V2 { let child_reference = agent_metadata .agent_path diff --git a/codex-rs/core/src/agent/control_tests.rs b/codex-rs/core/src/agent/control_tests.rs index 9863bfb7e2..4741c3c101 100644 --- a/codex-rs/core/src/agent/control_tests.rs +++ b/codex-rs/core/src/agent/control_tests.rs @@ -3,6 +3,8 @@ use crate::CodexThread; use crate::StateDbHandle; use crate::ThreadManager; use crate::agent::agent_status_from_event; +use crate::agent_communication::AgentCommunicationContext; +use crate::agent_communication::AgentCommunicationKind; use crate::config::AgentRoleConfig; use crate::config::Config; use crate::config::ConfigBuilder; @@ -65,12 +67,11 @@ async fn test_config() -> (TempDir, Config) { test_config_with_cli_overrides(Vec::new()).await } -fn text_input(text: &str) -> Op { +fn text_input(text: &str) -> Vec { vec![UserInput::Text { text: text.to_string(), text_elements: Vec::new(), }] - .into() } fn assistant_message(text: &str, phase: Option) -> ResponseItem { @@ -311,8 +312,7 @@ async fn send_input_errors_when_manager_dropped() { vec![UserInput::Text { text: "hello".to_string(), text_elements: Vec::new(), - }] - .into(), + }], ) .await .expect_err("send_input should fail without a manager"); @@ -423,8 +423,7 @@ async fn send_input_errors_when_thread_missing() { vec![UserInput::Text { text: "hello".to_string(), text_elements: Vec::new(), - }] - .into(), + }], ) .await .expect_err("send_input should fail for missing thread"); @@ -490,8 +489,7 @@ async fn send_input_submits_user_message() { vec![UserInput::Text { text: "hello from tests".to_string(), text_elements: Vec::new(), - }] - .into(), + }], ) .await .expect("send_input should succeed"); @@ -531,7 +529,11 @@ async fn send_inter_agent_communication_without_turn_queues_message_without_trig let submission_id = harness .control - .send_inter_agent_communication(thread_id, communication.clone()) + .send_inter_agent_communication( + thread_id, + communication.clone(), + AgentCommunicationContext::new(AgentCommunicationKind::Message, ThreadId::new()), + ) .await .expect("send_inter_agent_communication should succeed"); assert!(!submission_id.is_empty()); @@ -656,7 +658,11 @@ async fn ensure_v2_agent_loaded_reloads_registered_unloaded_agent() { ); harness .control - .send_inter_agent_communication(spawned_agent.thread_id, communication.clone()) + .send_inter_agent_communication( + spawned_agent.thread_id, + communication.clone(), + AgentCommunicationContext::new(AgentCommunicationKind::Message, ThreadId::new()), + ) .await .expect("send_inter_agent_communication should succeed after reload"); let expected = ( @@ -803,7 +809,11 @@ async fn encrypted_inter_agent_communication_clears_existing_last_task_message() ); harness .control - .send_inter_agent_communication(spawned_agent.thread_id, communication) + .send_inter_agent_communication( + spawned_agent.thread_id, + communication, + AgentCommunicationContext::new(AgentCommunicationKind::Followup, ThreadId::new()), + ) .await .expect("send_inter_agent_communication should succeed"); diff --git a/codex-rs/core/src/agent_communication.rs b/codex-rs/core/src/agent_communication.rs new file mode 100644 index 0000000000..d4c8ca22d2 --- /dev/null +++ b/codex-rs/core/src/agent_communication.rs @@ -0,0 +1,79 @@ +use codex_protocol::ThreadId; +use codex_protocol::protocol::InterAgentCommunication; + +const AGENT_COMMUNICATION_TARGET: &str = "codex_otel.agent_communication"; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum AgentCommunicationKind { + Spawn, + Message, + Followup, + Result, +} + +impl AgentCommunicationKind { + fn as_str(self) -> &'static str { + match self { + Self::Spawn => "spawn", + Self::Message => "message", + Self::Followup => "followup", + Self::Result => "result", + } + } +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub(crate) struct AgentCommunicationContext { + kind: AgentCommunicationKind, + pub(crate) sender_thread_id: ThreadId, +} + +impl AgentCommunicationContext { + pub(crate) fn new(kind: AgentCommunicationKind, sender_thread_id: ThreadId) -> Self { + Self { + kind, + sender_thread_id, + } + } +} + +pub(crate) fn logging_enabled() -> bool { + tracing::enabled!(target: AGENT_COMMUNICATION_TARGET, tracing::Level::INFO) +} + +pub(crate) fn emit_agent_communication_send( + communication_id: &str, + context: &AgentCommunicationContext, + communication: &InterAgentCommunication, + receiver_thread_id: ThreadId, +) { + tracing::info!( + target: AGENT_COMMUNICATION_TARGET, + { + event.name = "codex.agent_communication", + communication_id, + kind = context.kind.as_str(), + state = "send", + sender_thread_id = %context.sender_thread_id, + receiver_thread_id = %receiver_thread_id, + content = if communication.content.is_empty() { + communication.encrypted_content.as_deref().unwrap_or_default() + } else { + communication.content.as_str() + }, + }, + "agent communication" + ); +} + +pub(crate) fn emit_agent_communication_receive(communication_id: &str) { + tracing::info!( + target: AGENT_COMMUNICATION_TARGET, + { + event.name = "codex.agent_communication", + communication_id, + state = "receive", + }, + "agent communication" + ); +} diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index e9b749fea9..df8753a386 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -31,6 +31,7 @@ pub use codex_thread::TryStartTurnIfIdleError; pub use codex_thread::TryStartTurnIfIdleRejectionReason; pub use session::turn_context::TurnContext; mod agent; +mod agent_communication; mod attestation; mod codex_delegate; mod command_canonicalization; diff --git a/codex-rs/core/src/session/handlers.rs b/codex-rs/core/src/session/handlers.rs index a483ce5e0f..9733a9f66a 100644 --- a/codex-rs/core/src/session/handlers.rs +++ b/codex-rs/core/src/session/handlers.rs @@ -285,6 +285,7 @@ pub async fn inter_agent_communication( sess.input_queue .enqueue_mailbox_communication(communication) .await; + crate::agent_communication::emit_agent_communication_receive(&sub_id); if trigger_turn { sess.maybe_start_turn_for_pending_work_with_sub_id(sub_id) .await; diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index d43724e5db..809fa92015 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -14,6 +14,8 @@ use crate::agent::AgentControl; use crate::agent::AgentStatus; use crate::agent::agent_status_from_event; use crate::agent::status::is_final; +use crate::agent_communication::AgentCommunicationContext; +use crate::agent_communication::AgentCommunicationKind; use crate::attestation::AttestationProvider; use crate::build_available_skills; use crate::compact; @@ -1867,10 +1869,12 @@ impl Session { message, /*trigger_turn*/ false, ); + let context = + AgentCommunicationContext::new(AgentCommunicationKind::Result, self.thread_id); if let Err(err) = self .services .agent_control - .send_inter_agent_communication(parent_thread_id, communication) + .send_inter_agent_communication(parent_thread_id, communication, context) .await { debug!("failed to notify parent thread {parent_thread_id}: {err}"); diff --git a/codex-rs/core/src/tools/handlers/agent_jobs.rs b/codex-rs/core/src/tools/handlers/agent_jobs.rs index 360744973c..6405414731 100644 --- a/codex-rs/core/src/tools/handlers/agent_jobs.rs +++ b/codex-rs/core/src/tools/handlers/agent_jobs.rs @@ -203,7 +203,7 @@ async fn run_agent_job_loop( .agent_control .spawn_agent_with_metadata( options.spawn_config.clone(), - items.into(), + items, Some(SessionSource::SubAgent(SubAgentSource::Other(format!( "agent_job:{job_id}" )))), diff --git a/codex-rs/core/src/tools/handlers/multi_agents_common.rs b/codex-rs/core/src/tools/handlers/multi_agents_common.rs index 8f91ce5b8e..12ec89d651 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_common.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_common.rs @@ -18,7 +18,6 @@ use codex_protocol::openai_models::ReasoningEffort; use codex_protocol::openai_models::ReasoningEffortPreset; use codex_protocol::protocol::CollabAgentRef; use codex_protocol::protocol::CollabAgentStatusEntry; -use codex_protocol::protocol::Op; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::SubAgentSource; use codex_protocol::user_input::UserInput; @@ -163,7 +162,7 @@ pub(crate) fn thread_spawn_source( pub(crate) fn parse_collab_input( message: Option, items: Option>, -) -> Result { +) -> Result, FunctionCallError> { match (message, items) { (Some(_), Some(_)) => Err(FunctionCallError::RespondToModel( "Provide either message or items, but not both".to_string(), @@ -180,8 +179,7 @@ pub(crate) fn parse_collab_input( Ok(vec![UserInput::Text { text: message, text_elements: Vec::new(), - }] - .into()) + }]) } (None, Some(items)) => { if items.is_empty() { @@ -189,7 +187,7 @@ pub(crate) fn parse_collab_input( "Items can't be empty".to_string(), )); } - Ok(items.into()) + Ok(items) } } } diff --git a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs b/codex-rs/core/src/tools/handlers/multi_agents_tests.rs index 8acb0a8c8d..f313fcf0d9 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_tests.rs @@ -1377,8 +1377,7 @@ async fn multi_agent_v2_send_message_accepts_root_target_from_child() { vec![UserInput::Text { text: "inspect this repo".to_string(), text_elements: Vec::new(), - }] - .into(), + }], Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { parent_thread_id: root.thread_id, depth: 1, @@ -1454,8 +1453,7 @@ async fn multi_agent_v2_followup_task_rejects_root_target_from_child() { vec![UserInput::Text { text: "inspect this repo".to_string(), text_elements: Vec::new(), - }] - .into(), + }], Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { parent_thread_id: root.thread_id, depth: 1, @@ -1626,8 +1624,7 @@ async fn multi_agent_v2_list_agents_filters_by_relative_path_prefix() { vec![UserInput::Text { text: "research".to_string(), text_elements: Vec::new(), - }] - .into(), + }], Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { parent_thread_id: root.thread_id, depth: 1, @@ -1647,8 +1644,7 @@ async fn multi_agent_v2_list_agents_filters_by_relative_path_prefix() { vec![UserInput::Text { text: "build".to_string(), text_elements: Vec::new(), - }] - .into(), + }], Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { parent_thread_id: root.thread_id, depth: 2, @@ -4097,8 +4093,7 @@ async fn multi_agent_v2_interrupt_agent_rejects_self_target_by_id() { vec![UserInput::Text { text: "inspect this repo".to_string(), text_elements: Vec::new(), - }] - .into(), + }], Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { parent_thread_id: root.thread_id, depth: 1, @@ -4165,8 +4160,7 @@ async fn multi_agent_v2_interrupt_agent_rejects_self_target_by_task_name() { vec![UserInput::Text { text: "inspect this repo".to_string(), text_elements: Vec::new(), - }] - .into(), + }], Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { parent_thread_id: root.thread_id, depth: 1, diff --git a/codex-rs/core/src/tools/handlers/multi_agents_v2.rs b/codex-rs/core/src/tools/handlers/multi_agents_v2.rs index cae3497a13..8e7f73a762 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_v2.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_v2.rs @@ -19,7 +19,6 @@ use codex_protocol::protocol::CollabWaitingEndEvent; use codex_protocol::protocol::InterAgentCommunication; use codex_protocol::protocol::SubAgentActivityEvent; use codex_protocol::protocol::SubAgentActivityKind; -use codex_protocol::user_input::UserInput; use codex_tools::ToolName; use serde::Deserialize; use serde::Serialize; diff --git a/codex-rs/core/src/tools/handlers/multi_agents_v2/message_tool.rs b/codex-rs/core/src/tools/handlers/multi_agents_v2/message_tool.rs index 7fdb80e8bc..9501b4355a 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_v2/message_tool.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_v2/message_tool.rs @@ -4,6 +4,8 @@ //! resulting `InterAgentCommunication` should wake the target immediately. use super::*; +use crate::agent_communication::AgentCommunicationContext; +use crate::agent_communication::AgentCommunicationKind; use crate::tools::context::FunctionToolOutput; use crate::turn_timing::now_unix_timestamp_ms; use codex_protocol::protocol::InterAgentCommunication; @@ -46,7 +48,7 @@ pub(crate) struct FollowupTaskArgs { pub(crate) message: String, } -fn message_content(message: String) -> Result { +pub(super) fn message_content(message: String) -> Result { if message.trim().is_empty() { return Err(FunctionCallError::RespondToModel( "Empty message can't be sent to an agent".to_string(), @@ -101,10 +103,15 @@ pub(crate) async fn handle_message_string_tool( .unwrap_or_else(AgentPath::root); let communication = communication_from_tool_message(author, receiver_agent_path.clone(), message); + let kind = match mode { + MessageDeliveryMode::QueueOnly => AgentCommunicationKind::Message, + MessageDeliveryMode::TriggerTurn => AgentCommunicationKind::Followup, + }; + let context = AgentCommunicationContext::new(kind, session.thread_id); let result = session .services .agent_control - .send_inter_agent_communication(receiver_thread_id, mode.apply(communication)) + .send_inter_agent_communication(receiver_thread_id, mode.apply(communication), context) .await .map_err(|err| collab_agent_error(receiver_thread_id, err)); result?; diff --git a/codex-rs/core/src/tools/handlers/multi_agents_v2/spawn.rs b/codex-rs/core/src/tools/handlers/multi_agents_v2/spawn.rs index 2186a4a6d6..8eeb6ddd1b 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_v2/spawn.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_v2/spawn.rs @@ -4,11 +4,13 @@ use crate::agent::control::SpawnAgentOptions; use crate::agent::next_thread_spawn_depth; use crate::agent::role::DEFAULT_ROLE_NAME; use crate::agent::role::apply_role_to_config; +use crate::agent_communication::AgentCommunicationContext; +use crate::agent_communication::AgentCommunicationKind; use crate::tools::handlers::multi_agents_spec::SpawnAgentToolOptions; use crate::tools::handlers::multi_agents_spec::create_spawn_agent_tool_v2; +use crate::tools::handlers::multi_agents_v2::message_tool::message_content; use crate::turn_timing::now_unix_timestamp_ms; use codex_protocol::AgentPath; -use codex_protocol::protocol::Op; use codex_tools::ToolSpec; #[derive(Default)] @@ -55,8 +57,7 @@ async fn handle_spawn_agent( .map(str::trim) .filter(|role| !role.is_empty()); - let message = args.message.clone(); - let initial_operation = parse_collab_input(Some(args.message), /*items*/ None)?; + let message = message_content(args.message)?; let session_source = turn.session_source.clone(); let child_depth = next_thread_spawn_depth(&session_source); let mut config = @@ -104,33 +105,28 @@ async fn handle_spawn_agent( "spawned agent is missing a canonical task name".to_string(), ) })?; + let author = turn + .session_source + .get_agent_path() + .unwrap_or_else(AgentPath::root); + let communication = communication_from_tool_message(author, new_agent_path.clone(), message); + let context = AgentCommunicationContext::new(AgentCommunicationKind::Spawn, session.thread_id); let spawned_agent = Box::pin( - session.services.agent_control.spawn_agent_with_metadata( - config, - match initial_operation { - Op::UserInput { items, .. } - if items - .iter() - .all(|item| matches!(item, UserInput::Text { .. })) => - { - let author = turn - .session_source - .get_agent_path() - .unwrap_or_else(AgentPath::root); - let communication = - communication_from_tool_message(author, new_agent_path.clone(), message); - Op::InterAgentCommunication { communication } - } - initial_operation => initial_operation, - }, - Some(spawn_source), - SpawnAgentOptions { - fork_parent_spawn_call_id: fork_mode.as_ref().map(|_| call_id.clone()), - fork_mode, - parent_thread_id: Some(session.thread_id), - environments: Some(turn.environments.to_selections()), - }, - ), + session + .services + .agent_control + .spawn_agent_with_communication( + config, + communication, + context, + Some(spawn_source), + SpawnAgentOptions { + fork_parent_spawn_call_id: fork_mode.as_ref().map(|_| call_id.clone()), + fork_mode, + parent_thread_id: Some(session.thread_id), + environments: Some(turn.environments.to_selections()), + }, + ), ) .await .map_err(collab_spawn_error)?; diff --git a/codex-rs/core/tests/suite/subagent_notifications.rs b/codex-rs/core/tests/suite/subagent_notifications.rs index 951fb92f17..56fd0b9581 100644 --- a/codex-rs/core/tests/suite/subagent_notifications.rs +++ b/codex-rs/core/tests/suite/subagent_notifications.rs @@ -39,10 +39,13 @@ use serde_json::Value; use serde_json::json; use std::fs; use std::path::Path; +use std::sync::Mutex; use std::time::Duration; use test_case::test_case; use tokio::time::Instant; use tokio::time::sleep; +use tracing::Level; +use tracing_test::internal::MockWriter; use wiremock::MockServer; const SPAWN_CALL_ID: &str = "spawn-call-1"; @@ -96,6 +99,13 @@ fn decoded_body(req: &wiremock::Request) -> Option> { } } +fn log_field<'a>(line: &'a str, name: &str) -> Option<&'a str> { + let prefix = format!("{name}="); + line.split_ascii_whitespace() + .find_map(|field| field.strip_prefix(&prefix)) + .map(|value| value.trim_matches('"')) +} + fn has_subagent_notification(req: &ResponsesRequest) -> bool { req.message_input_texts("user") .iter() @@ -1042,8 +1052,16 @@ async fn spawned_multi_agent_v2_child_inherits_parent_developer_context() -> Res Ok(()) } -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[tokio::test] async fn encrypted_multi_agent_v2_spawn_sends_agent_message_to_child() -> Result<()> { + let output: &'static Mutex> = Box::leak(Box::new(Mutex::new(Vec::new()))); + let subscriber = tracing_subscriber::fmt() + .with_ansi(false) + .with_max_level(Level::TRACE) + .with_writer(MockWriter::new(output)) + .finish(); + let _guard = tracing::subscriber::set_default(subscriber); + let server = start_mock_server().await; let encrypted_message = "opaque-encrypted-message"; let spawn_args = serde_json::to_string(&json!({ @@ -1098,6 +1116,7 @@ async fn encrypted_multi_agent_v2_spawn_sends_agent_message_to_child() -> Result .expect("test config should allow feature update"); }); let test = builder.build(&server).await?; + let root_thread_id = test.session_configured.thread_id; test.submit_turn(TURN_1_PROMPT).await?; @@ -1124,6 +1143,41 @@ async fn encrypted_multi_agent_v2_spawn_sends_agent_message_to_child() -> Result })]) ); + let child_thread_id = test + .thread_manager + .list_thread_ids() + .await + .into_iter() + .find(|thread_id| *thread_id != root_thread_id) + .expect("child thread ID"); + let logs = tokio::time::timeout(Duration::from_secs(5), async { + loop { + let logs = String::from_utf8(output.lock().expect("buffer lock").clone()) + .expect("logs should be UTF-8"); + if logs.contains("kind=\"spawn\"") && logs.contains("state=\"receive\"") { + break logs; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("spawn communication logs should be emitted"); + let send = logs + .lines() + .find(|line| line.contains("kind=\"spawn\"") && line.contains("state=\"send\"")) + .expect("spawn send event"); + assert!(send.contains(&format!("sender_thread_id={root_thread_id}"))); + assert!(send.contains(&format!("receiver_thread_id={child_thread_id}"))); + assert!(send.contains(&format!("content=\"{encrypted_message}\""))); + + let communication_id = log_field(send, "communication_id").expect("communication ID"); + logs.lines() + .find(|line| { + line.contains("state=\"receive\"") + && log_field(line, "communication_id") == Some(communication_id) + }) + .expect("correlated receive event"); + Ok(()) } diff --git a/codex-rs/feedback/src/lib.rs b/codex-rs/feedback/src/lib.rs index 7c27d2b3b0..82c5851a66 100644 --- a/codex-rs/feedback/src/lib.rs +++ b/codex-rs/feedback/src/lib.rs @@ -18,6 +18,7 @@ use tracing::Event; use tracing::Level; use tracing::field::Visit; use tracing_subscriber::Layer; +use tracing_subscriber::filter::LevelFilter; use tracing_subscriber::filter::Targets; use tracing_subscriber::fmt::writer::MakeWriter; use tracing_subscriber::registry::LookupSpan; @@ -204,7 +205,11 @@ impl CodexFeedback { .with_target(false) // Capture everything, regardless of the caller's `RUST_LOG`, so feedback includes the // full trace when the user uploads a report. - .with_filter(Targets::new().with_default(Level::TRACE)) + .with_filter( + Targets::new() + .with_default(Level::TRACE) + .with_target("codex_otel.agent_communication", LevelFilter::OFF), + ) } /// Returns a [`tracing_subscriber`] layer that collects structured metadata for feedback. @@ -723,6 +728,28 @@ mod tests { pretty_assertions::assert_eq!(snap.tags.get("cached").map(String::as_str), Some("true")); } + #[test] + fn logger_layer_excludes_agent_communication_content() { + let fb = CodexFeedback::new(); + let _guard = tracing_subscriber::registry() + .with(fb.logger_layer()) + .set_default(); + + tracing::info!(target: "codex_core", "retained-log"); + tracing::info!( + target: "codex_otel.agent_communication", + content = "secret", + "dropped-log" + ); + + let logs = std::str::from_utf8(fb.snapshot(/*session_id*/ None).as_bytes()) + .expect("feedback logs should be UTF-8") + .to_string(); + assert!(logs.contains("retained-log")); + assert!(!logs.contains("secret")); + assert!(!logs.contains("dropped-log")); + } + #[test] fn feedback_attachments_gate_connectivity_diagnostics() { let extra_filename = format!("codex-feedback-extra-{}.jsonl", ThreadId::new()); diff --git a/codex-rs/state/src/log_db.rs b/codex-rs/state/src/log_db.rs index 2cf0f5d9af..2a7c205d62 100644 --- a/codex-rs/state/src/log_db.rs +++ b/codex-rs/state/src/log_db.rs @@ -54,6 +54,7 @@ pub fn default_filter() -> Targets { Targets::new() .with_default(LevelFilter::TRACE) .with_target("log", LevelFilter::OFF) + .with_target("codex_otel.agent_communication", LevelFilter::OFF) .with_target("codex_otel.log_only", LevelFilter::OFF) .with_target("codex_otel.trace_safe", LevelFilter::OFF) } diff --git a/codex-rs/state/src/log_db_filter_tests.rs b/codex-rs/state/src/log_db_filter_tests.rs index 36ee27aba8..7251a5f6a9 100644 --- a/codex-rs/state/src/log_db_filter_tests.rs +++ b/codex-rs/state/src/log_db_filter_tests.rs @@ -1,5 +1,4 @@ use pretty_assertions::assert_eq; -use tracing_subscriber::filter::Targets; use tracing_subscriber::layer::SubscriberExt; use tracing_subscriber::util::SubscriberInitExt; use uuid::Uuid; @@ -16,17 +15,18 @@ async fn sqlite_sink_drops_low_level_opentelemetry_sdk_logs() { let layer = start(runtime.clone()); let guard = tracing_subscriber::registry() - .with( - layer - .clone() - .with_filter(Targets::new().with_default(tracing::Level::TRACE)), - ) + .with(layer.clone().with_filter(default_filter())) .set_default(); tracing::trace!(target: "opentelemetry_sdk", "dropped-trace"); tracing::debug!(target: "opentelemetry_sdk", "dropped-debug"); tracing::info!(target: "opentelemetry_sdk", "retained-info"); tracing::trace!(target: "codex_state", "retained-trace"); + tracing::info!( + target: "codex_otel.agent_communication", + content = "secret", + "dropped-agent-communication" + ); layer.flush().await; drop(guard);