From 9f04d886f3663b584fba87d6fd71001fb6f9fcf3 Mon Sep 17 00:00:00 2001 From: Jenia Varavva Date: Mon, 9 Nov 2026 20:09:05 +0000 Subject: [PATCH] Support cancellation and deadlines for realtime attachments (#60443) ## Why A pending realtime connection handshake can hold up subsequent attachment requests. Callers need to cancel obsolete attempts or limit how long they may take to establish a connection. ## What changed - Add `RealtimeAttachmentOptions` with an optional cancellation receiver and `start_before` deadline to `Op::RealtimeConversationAttach`. - Interrupt pending connection work and finish cleanup before acknowledging cancellation. Reject cancelled or expired requests before replacing an existing conversation, and recheck cancellation and the deadline before publishing a connection. - Preserve a successfully published connection when cancellation arrives later. Dropping a cancellation sender whose value is `false` lets the attempt continue. - Prepare the daemon startup test's socket directory with `codex_uds::prepare_private_socket_directory`. ## Testing Add integration coverage for cancelling a held handshake, skipping cancelled or expired queued requests, preserving the next usable call, and releasing a pending handshake when its deadline expires without a cancellation receiver. GitOrigin-RevId: 7144c1237bcbe3b859f10b654d919ae9ff20cbbb --- codex-rs/Cargo.lock | 1 + codex-rs/cli/Cargo.toml | 1 + codex-rs/cli/tests/daemon_startup.rs | 3 +- codex-rs/core/src/realtime_conversation.rs | 132 +++++++++++++- .../realtime_conversation/existing_call.rs | 9 +- codex-rs/core/src/session/handlers.rs | 9 +- .../core/tests/suite/realtime_attachment.rs | 172 +++++++++++++++++- codex-rs/protocol/src/protocol.rs | 14 ++ 8 files changed, 324 insertions(+), 17 deletions(-) diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 0cdfa65c5f..c3f77e4afe 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2607,6 +2607,7 @@ dependencies = [ "codex-terminal-detection", "codex-thread-store", "codex-tui", + "codex-uds", "codex-utils-absolute-path", "codex-utils-cargo-bin", "codex-utils-cli", diff --git a/codex-rs/cli/Cargo.toml b/codex-rs/cli/Cargo.toml index ffff79d331..501db96af1 100644 --- a/codex-rs/cli/Cargo.toml +++ b/codex-rs/cli/Cargo.toml @@ -130,6 +130,7 @@ app_test_support = { workspace = true } assert_cmd = { workspace = true } assert_matches = { workspace = true } codex-app-server-client = { workspace = true } +codex-uds = { workspace = true } codex-utils-cargo-bin = { workspace = true } codex-utils-pty = { workspace = true } core_test_support = { workspace = true } diff --git a/codex-rs/cli/tests/daemon_startup.rs b/codex-rs/cli/tests/daemon_startup.rs index aadb824861..56a7cc76b3 100644 --- a/codex-rs/cli/tests/daemon_startup.rs +++ b/codex-rs/cli/tests/daemon_startup.rs @@ -259,7 +259,8 @@ async fn daemon_startup(command: &str) -> Result<()> { fs::hard_link(&codex, &managed) .or_else(|_| codex_utils_cargo_bin::copy_executable(&codex, &managed))?; if command != "restrictive-job" { - fs::create_dir(home.path().join("app-server-daemon"))?; + codex_uds::prepare_private_socket_directory(&home.path().join("app-server-daemon")) + .await?; fs::write( home.path().join("app-server-daemon/settings.json"), if persisted { diff --git a/codex-rs/core/src/realtime_conversation.rs b/codex-rs/core/src/realtime_conversation.rs index 6545ed97d3..cf4c4f5a13 100644 --- a/codex-rs/core/src/realtime_conversation.rs +++ b/codex-rs/core/src/realtime_conversation.rs @@ -29,6 +29,7 @@ use codex_api::RealtimeWebsocketEvents; use codex_api::RealtimeWebsocketWriter; use codex_api::build_session_headers; use codex_api::map_api_error; +use codex_async_utils::OrCancelExt; use codex_config::config_toml::RealtimeWsMode; use codex_config::config_toml::RealtimeWsVersion; use codex_features::Feature; @@ -55,6 +56,7 @@ use codex_protocol::protocol::ConversationTextRole; use codex_protocol::protocol::ErrorEvent; use codex_protocol::protocol::Event; use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::RealtimeAttachmentOptions; use codex_protocol::protocol::RealtimeConversationClosedEvent; use codex_protocol::protocol::RealtimeConversationRealtimeEvent; use codex_protocol::protocol::RealtimeConversationSdpEvent; @@ -79,8 +81,10 @@ use std::sync::atomic::AtomicBool; use std::sync::atomic::Ordering; use std::time::Duration; use std::time::Instant; +use std::time::SystemTime; use tokio::sync::Mutex; use tokio::sync::Semaphore; +use tokio::sync::watch; use tokio::task::JoinHandle; use tokio_util::sync::CancellationToken; use tracing::debug; @@ -546,7 +550,29 @@ struct ConversationState { stop_token: CancellationToken, } +#[derive(Clone, Default)] +struct RealtimeAttachmentControl { + token: CancellationToken, + options: RealtimeAttachmentOptions, +} + +impl RealtimeAttachmentControl { + fn is_cancelled(&self) -> bool { + self.token.is_cancelled() + || self + .options + .cancel + .as_ref() + .is_some_and(|cancellation| *cancellation.borrow()) + || self + .options + .start_before + .is_some_and(|deadline| SystemTime::now() >= deadline) + } +} + struct RealtimeStart { + attachment_control: RealtimeAttachmentControl, sideband: Option, api_provider: ApiProvider, http_client_factory: HttpClientFactory, @@ -643,6 +669,9 @@ impl RealtimeConversationManager { ) -> CodexResult { let previous_state = { let mut guard = self.state.lock().await; + if start.attachment_control.is_cancelled() { + return Err(CodexErr::TurnAborted); + } guard.conversation.take() }; if let Some(state) = previous_state { @@ -667,6 +696,7 @@ impl RealtimeConversationManager { ) -> CodexResult { let session_id = start.session_config.session_id.clone(); let RealtimeStart { + attachment_control, sideband, auth_context, api_provider, @@ -687,6 +717,9 @@ impl RealtimeConversationManager { sdp, existing_call_id, } = start; + if attachment_control.is_cancelled() { + return Err(CodexErr::TurnAborted); + } let event_parser = session_config.event_parser; let session_kind = match event_parser { RealtimeEventParser::V1 | RealtimeEventParser::FramelessBidi => RealtimeSessionKind::V1, @@ -769,7 +802,8 @@ impl RealtimeConversationManager { &auth_context, http_client_factory, ) - .await?; + .or_cancel(&attachment_control.token) + .await??; let task = spawn_webrtc_sideband_input_task(RealtimeWebrtcSidebandInputTask { client: client .with_http_client_factory(call.sideband_http_client_factory) @@ -790,6 +824,7 @@ impl RealtimeConversationManager { (task, Some(call.sdp)) } else if let Some(call_id) = existing_call_id { let task = existing_call::attach(existing_call::ExistingCallAttachment { + attachment_cancellation: attachment_control.token.clone(), client, model_client, auth_context, @@ -817,7 +852,8 @@ impl RealtimeConversationManager { }; let connection = client .connect(session_config, extra_headers, default_headers()) - .await + .or_cancel(&attachment_control.token) + .await? .map_err(map_api_error)?; let task = spawn_realtime_input_task( RealtimeInputTask { @@ -837,8 +873,7 @@ impl RealtimeConversationManager { (task, None) }; - let mut state = self.state.lock().await; - state.conversation = Some(ConversationState { + let conversation = ConversationState { session_id, sub_id: sub_id.to_owned(), audio_tx, @@ -850,9 +885,39 @@ impl RealtimeConversationManager { realtime_active: Arc::clone(&realtime_active), route_handoffs: Arc::clone(&route_handoffs), stop_token, - }); - state.mode_instructions = Some(mode_instructions); - state.context_active = true; + }; + let uncommitted = { + let mut state = self.state.lock().await; + // Read the original sources at commit, even if their wakeup has not + // been polled. Hold the watch guard through synchronous publication. + let cancellation = attachment_control + .options + .cancel + .as_ref() + .map(watch::Receiver::borrow); + if cancellation.as_ref().is_some_and(|cancelled| **cancelled) + || attachment_control.token.is_cancelled() + || attachment_control + .options + .start_before + .is_some_and(|deadline| SystemTime::now() >= deadline) + { + Some(conversation) + } else { + // Publication commits success; later cancellation must preserve + // residency until the caller projects the Started event. + state.conversation = Some(conversation); + state.mode_instructions = Some(mode_instructions); + state.context_active = true; + None + } + }; + if let Some(conversation) = uncommitted { + drop(events_rx); + drop(transcript_tail_rx); + stop_conversation_state(conversation, RealtimeFanoutTaskStop::Await).await; + return Err(CodexErr::TurnAborted); + } Ok(RealtimeStartOutput { realtime_active, route_handoffs, @@ -1378,6 +1443,7 @@ pub(crate) async fn handle_start( } struct PreparedRealtimeConversationStart { + attachment_control: RealtimeAttachmentControl, sideband: Option, api_provider: ApiProvider, http_client_factory: HttpClientFactory, @@ -1595,6 +1661,7 @@ async fn prepare_realtime_start( extra_headers.insert(X_CODEX_TURN_METADATA_HEADER, metadata); } Ok(PreparedRealtimeConversationStart { + attachment_control: RealtimeAttachmentControl::default(), sideband: None, auth_context, api_provider, @@ -1823,6 +1890,7 @@ async fn handle_start_inner( prepared_start: PreparedRealtimeConversationStart, ) -> CodexResult<()> { let PreparedRealtimeConversationStart { + attachment_control, sideband, auth_context, api_provider, @@ -1857,6 +1925,7 @@ async fn handle_start_inner( }; let event_parser = session_config.event_parser; let start = RealtimeStart { + attachment_control, sideband, auth_context, api_provider, @@ -2184,14 +2253,57 @@ pub(crate) async fn handle_close(sess: &Arc, sub_id: String) { handle_detach(sess, /*expected*/ None, &sub_id).await; } -/// Uses the ordinary sequential start path and reports its connection result. +/// Completes after the connection is published or the cancelled attempt is cleaned up. pub(crate) async fn handle_attach( sess: &Arc, sub_id: String, params: ConversationStartParams, + options: RealtimeAttachmentOptions, ) -> CodexResult<()> { - let prepared = prepare_realtime_start(sess, params).await?; - handle_start_inner(sess, &sub_id, prepared).await + let attachment_control = RealtimeAttachmentControl { + token: CancellationToken::new(), + options, + }; + if attachment_control.is_cancelled() { + return Err(CodexErr::TurnAborted); + } + let remaining = attachment_control + .options + .start_before + .map(|deadline| deadline.duration_since(SystemTime::now())) + .transpose() + .map_err(|_| CodexErr::TurnAborted)?; + if remaining == Some(Duration::ZERO) { + return Err(CodexErr::TurnAborted); + } + let attach = async { + let mut prepared = prepare_realtime_start(sess, params) + .or_cancel(&attachment_control.token) + .await??; + prepared.attachment_control = attachment_control.clone(); + handle_start_inner(sess, &sub_id, prepared).await + }; + tokio::pin!(attach); + let cancelled = async { + let Some(mut cancellation) = attachment_control.options.cancel.clone() else { + return false; + }; + cancellation.wait_for(|cancelled| *cancelled).await.is_ok() + }; + let expired = async { + match remaining { + Some(remaining) => tokio::time::sleep(remaining).await, + None => std::future::pending::<()>().await, + } + }; + tokio::select! { + biased; + true = cancelled => {} + _ = expired => {} + result = &mut attach => return result, + } + attachment_control.token.cancel(); + attach.await } /// Attaches Core to an existing call using the caller’s channels to receive diff --git a/codex-rs/core/src/realtime_conversation/existing_call.rs b/codex-rs/core/src/realtime_conversation/existing_call.rs index 1f9a6918ef..c8e9b35df6 100644 --- a/codex-rs/core/src/realtime_conversation/existing_call.rs +++ b/codex-rs/core/src/realtime_conversation/existing_call.rs @@ -13,6 +13,7 @@ use codex_api::RealtimeSessionConfig; use codex_api::RealtimeTranscriptState; use codex_api::RealtimeWebsocketClient; use codex_api::map_api_error; +use codex_async_utils::OrCancelExt; use codex_http_client::HttpClientFactory; use codex_login::default_client::default_headers; use codex_protocol::error::Result as CodexResult; @@ -23,6 +24,7 @@ use tokio::task::JoinHandle; use tokio_util::sync::CancellationToken; pub(super) struct ExistingCallAttachment { + pub(super) attachment_cancellation: CancellationToken, pub(super) client: RealtimeWebsocketClient, pub(super) model_client: ModelClient, pub(super) auth_context: codex_model_provider::ProviderAuthContext, @@ -43,6 +45,7 @@ pub(super) struct ExistingCallAttachment { pub(super) async fn attach(attachment: ExistingCallAttachment) -> CodexResult> { let ExistingCallAttachment { + attachment_cancellation, client, model_client, auth_context, @@ -67,7 +70,8 @@ pub(super) async fn attach(attachment: ExistingCallAttachment) -> CodexResult CodexResult { + Op::RealtimeConversationAttach { + params, + options, + reply, + } => { let result = - crate::realtime_conversation::handle_attach(&sess, sub.id, params).await; + crate::realtime_conversation::handle_attach(&sess, sub.id, params, options) + .await; let _ = reply.send(result); false } diff --git a/codex-rs/core/tests/suite/realtime_attachment.rs b/codex-rs/core/tests/suite/realtime_attachment.rs index 38e63ac3fa..444428ea4d 100644 --- a/codex-rs/core/tests/suite/realtime_attachment.rs +++ b/codex-rs/core/tests/suite/realtime_attachment.rs @@ -8,6 +8,7 @@ use codex_protocol::protocol::ConversationTextParams; use codex_protocol::protocol::ConversationTextRole; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::Op; +use codex_protocol::protocol::RealtimeAttachmentOptions; use codex_protocol::protocol::RealtimeConversationVersion; use codex_protocol::protocol::RealtimeOutputModality; use codex_protocol::protocol::RealtimeSidebandTransport; @@ -23,9 +24,11 @@ use pretty_assertions::assert_eq; use serde_json::Value; use serde_json::json; use std::time::Duration; +use std::time::SystemTime; use tokio::net::TcpListener; use tokio::sync::mpsc; use tokio::sync::oneshot; +use tokio::sync::watch; use tokio::time::timeout; use tokio_tungstenite::tungstenite::Message; @@ -363,7 +366,11 @@ async fn start( let (reply, result) = oneshot::channel(); let id = test .codex - .submit(Op::RealtimeConversationAttach { params, reply }) + .submit(Op::RealtimeConversationAttach { + params, + reply, + options: RealtimeAttachmentOptions::default(), + }) .await?; Ok((id, result)) } @@ -420,6 +427,163 @@ async fn assert_text_routes_to(test: &TestCodex, connection: &mut Connected) -> Ok(()) } +/// Cancellation releases a held handshake and obsolete queued work before selecting the next call. +#[tokio::test] +async fn cancelled_attachments_release_the_queue_and_preserve_the_next_call() -> Result<()> { + let api = start_mock_server().await; + let test = test_codex().build_with_auto_env(&api).await?; + let mut original = Sideband::new().await?; + let mut newer = Sideband::new().await?; + let (cancel_original, cancellation) = watch::channel(false); + let (reply, original_result) = oneshot::channel(); + test.codex + .submit(Op::RealtimeConversationAttach { + params: params("original", &original.url), + reply, + options: RealtimeAttachmentOptions { + cancel: Some(cancellation), + start_before: None, + }, + }) + .await?; + // Keep the first upgrade pending until the replacement is already usable. + let mut original_connection = original.connected().await; + + let (cancel_queued, cancellation) = watch::channel(false); + let (reply, queued_result) = oneshot::channel(); + cancel_queued.send_replace(/*value*/ true); + test.codex + .submit(Op::RealtimeConversationAttach { + params: params("queued", &newer.url), + reply, + options: RealtimeAttachmentOptions { + cancel: Some(cancellation), + start_before: None, + }, + }) + .await?; + let (_cancel_expired, cancellation) = watch::channel(false); + let (reply, expired_result) = oneshot::channel(); + test.codex + .submit(Op::RealtimeConversationAttach { + params: params("expired", &newer.url), + reply, + options: RealtimeAttachmentOptions { + cancel: Some(cancellation), + start_before: Some(SystemTime::now() - Duration::from_secs(1)), + }, + }) + .await?; + let (cancel_newer, cancellation) = watch::channel(false); + let (reply, newer_result) = oneshot::channel(); + test.codex + .submit(Op::RealtimeConversationAttach { + params: params("newer", &newer.url), + reply, + options: RealtimeAttachmentOptions { + cancel: Some(cancellation), + start_before: None, + }, + }) + .await?; + // Losing the caller's cancellation sender does not cancel admitted work. + drop(cancel_newer); + cancel_original.send_replace(/*value*/ true); + assert!( + timeout(Duration::from_secs(10), original_result) + .await?? + .is_err() + ); + assert!( + timeout(Duration::from_secs(10), queued_result) + .await?? + .is_err() + ); + assert!( + timeout(Duration::from_secs(10), expired_result) + .await?? + .is_err() + ); + + let connection = newer.connected().await; + connection + .accept + .send(/*value*/ true) + .expect("handshake pending"); + timeout(Duration::from_secs(10), newer_result).await???; + let started = wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::RealtimeConversationStarted(_)) + }) + .await; + let EventMsg::RealtimeConversationStarted(started) = started else { + unreachable!("wait_for_event selected the Started event"); + }; + assert_eq!(started.realtime_session_id.as_deref(), Some("newer")); + assert_eq!(started.version, RealtimeConversationVersion::V3); + assert!(matches!( + newer.connections.try_recv(), + Err(mpsc::error::TryRecvError::Empty) + )); + original_connection + .accept + .send(/*value*/ true) + .expect("original upgrade remains held until cancellation completes"); + assert_eq!( + timeout(Duration::from_secs(10), original_connection.messages.recv()).await?, + None + ); + detach(&test, "original").await?; + let mut current = Connected { + messages: connection.messages, + events: connection.events, + }; + assert_text_routes_to(&test, &mut current).await?; + detach(&test, "newer").await?; + assert_eq!( + timeout(Duration::from_secs(10), current.messages.recv()).await?, + None + ); + test.codex.shutdown_and_wait().await?; + Ok(()) +} + +/// A deadline without a cancellation receiver releases a held handshake so the next call can connect. +#[tokio::test] +async fn attachment_deadline_releases_a_pending_handshake() -> Result<()> { + let api = start_mock_server().await; + let test = test_codex().build_with_auto_env(&api).await?; + let mut sideband = Sideband::new().await?; + let (reply, result) = oneshot::channel(); + test.codex + .submit(Op::RealtimeConversationAttach { + params: params("expired", &sideband.url), + reply, + options: RealtimeAttachmentOptions { + cancel: None, + start_before: Some(SystemTime::now() + Duration::from_secs(60)), + }, + }) + .await?; + let mut connection = sideband.connected().await; + tokio::time::pause(); + tokio::time::advance(Duration::from_secs(61)).await; + let expired = timeout(Duration::from_secs(10), result).await; + tokio::time::resume(); + assert!(expired??.is_err()); + connection + .accept + .send(/*value*/ true) + .expect("handshake still held"); + assert_eq!( + timeout(Duration::from_secs(10), connection.messages.recv()).await?, + None + ); + let (_, mut current) = attached(&test, &mut sideband, "current").await?; + assert_text_routes_to(&test, &mut current).await?; + test.codex.shutdown_and_wait().await?; + Ok(()) +} + /// Starts are sequential; completion follows the handshake and a delayed stop A leaves B usable. pub(super) async fn attachment_replacement_scenario() -> Result> { let api = start_mock_server().await; @@ -541,7 +705,11 @@ async fn completed_turns_reach_history_before_the_next_handoff(flush_tail: bool) params.flush_transcript_tail_on_session_end = flush_tail; let (reply, attached) = oneshot::channel(); test.codex - .submit(Op::RealtimeConversationAttach { params, reply }) + .submit(Op::RealtimeConversationAttach { + params, + reply, + options: RealtimeAttachmentOptions::default(), + }) .await?; let connection = sideband.connected().await; connection diff --git a/codex-rs/protocol/src/protocol.rs b/codex-rs/protocol/src/protocol.rs index e11fb4bdfb..877e37a315 100644 --- a/codex-rs/protocol/src/protocol.rs +++ b/codex-rs/protocol/src/protocol.rs @@ -11,6 +11,7 @@ use std::path::Path; use std::path::PathBuf; use std::str::FromStr; use std::time::Duration; +use std::time::SystemTime; use strum_macros::EnumIter; @@ -82,6 +83,7 @@ use serde_json::Value; use serde_with::serde_as; use strum_macros::Display; use tokio::sync::oneshot; +use tokio::sync::watch; use tracing::error; use ts_rs::TS; @@ -222,6 +224,16 @@ pub struct ConversationStartParams { pub voice: Option, } +/// Controls when an in-process realtime attachment may establish a connection. +#[derive(Debug, Clone, Default)] +pub struct RealtimeAttachmentOptions { + /// Set the watch to `true` to cancel a pending connection. + /// Dropping the sender while its value is `false` lets the attempt continue. + pub cancel: Option>, + /// Optional caller deadline for publishing the connection. Cleanup can finish later. + pub start_before: Option, +} + /// A caller-owned realtime sideband. Dropping either direction seals input and ends /// the attachment after accepted messages drain. The caller supplies bounded channels /// and one complete JSON message per item. @@ -724,8 +736,10 @@ pub enum Op { RealtimeConversationStart(ConversationStartParams), /// Start a realtime conversation and report completion of the connection attempt. + /// If cancellation wins, completion acknowledges connection cleanup. RealtimeConversationAttach { params: ConversationStartParams, + options: RealtimeAttachmentOptions, reply: oneshot::Sender>, },