diff --git a/codex-rs/core/src/agent/host/service_tier.rs b/codex-rs/core/src/agent/host/service_tier.rs index 8f15f89030..185099c228 100644 --- a/codex-rs/core/src/agent/host/service_tier.rs +++ b/codex-rs/core/src/agent/host/service_tier.rs @@ -8,7 +8,7 @@ use futures::FutureExt; use futures::future::BoxFuture; use futures::future::Shared; -pub(crate) type ServiceTierAcceptance = Shared>; +pub(crate) type ServiceTierAcceptance = Shared>>; /// The host's view of the service tier shared within an agent session. /// @@ -27,9 +27,12 @@ pub trait AgentServiceTier: Send + Sync { /// Hand off the committed tier in call order and return a receipt for local acceptance. /// Core calls this method under its commit locks: the handoff must be fast and nonblocking. /// Drive the update independently of the receipt: dropping the receipt only cancels the wait. - /// Core awaits it outside the state and registry locks. Explicit settings updates that - /// request the same committed tier share this receipt until the tier changes. - fn update(&self, service_tier: Option) -> BoxFuture<'static, ()>; + /// Core waits outside the state and registry locks. `Ok(())` confirms local acceptance. + /// `Err` leaves acceptance unconfirmed and carries a caller-safe message. Core surfaces it + /// in an explicit thread settings reply. A later settings request with a reply can hand off + /// the same tier again after Core observes a failed receipt. The host must preserve commit + /// order and tolerate repeated handoffs. + fn update(&self, service_tier: Option) -> BoxFuture<'static, Result<(), String>>; } pub(crate) fn host_child_uses_shared_service_tier(session: &Session) -> bool { diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index af15c8d5a0..e97fe781ab 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/codex-rs/core/src/codex_thread.rs @@ -826,7 +826,9 @@ impl CodexThread { } /// Queues settings for future turns and waits for Core to accept or reject them. - /// Rejections are returned to the caller instead of emitted as thread errors. + /// Rejections are returned to the caller instead of emitted as thread errors. A host-root + /// service-tier update also waits for local host acceptance outside Core's submission loop. + /// A handoff error leaves the settings committed locally. Callers can retry the same tier. pub async fn update_thread_settings( &self, thread_settings: impl Into>, @@ -857,7 +859,11 @@ impl CodexThread { settings: CodexThreadSettingsOverrides, ) -> ConstraintResult<()> { let updates = self.thread_settings_update(settings); - self.session.update_settings(updates).await.map(|_| ()) + let commit = self.session.update_settings(updates).await?; + if let Some(accepted) = commit.host_acceptance { + let _ = accepted.await; + } + Ok(()) } /// Replaces tools for subsequent model steps and future turns, preserving in-flight calls. diff --git a/codex-rs/core/src/session/handlers.rs b/codex-rs/core/src/session/handlers.rs index f055cdba19..746ff4f220 100644 --- a/codex-rs/core/src/session/handlers.rs +++ b/codex-rs/core/src/session/handlers.rs @@ -683,13 +683,33 @@ pub(super) async fn submission_loop( request: thread_settings, turn_extension_init: sub.turn_extension_init, }; - match thread_settings::update(&sess, thread_settings).await { - Ok(snapshot) => { - // Reply first: the caller may hold a lock its event consumer needs. - if let Some(reply) = reply { - let _ = reply.send(Ok(())); + match thread_settings::update(&sess, thread_settings, reply.is_some()).await { + Ok(commit) => { + match (reply, commit.host_acceptance) { + (Some(mut reply), Some(acceptance)) => { + // Leave the submission loop free to process interrupts. + tokio::spawn(async move { + tokio::select! { + _ = reply.closed() => {} + result = acceptance => { + let result = result.map_err(|error| { + CodexErr::from(std::io::Error::other(format!( + "service tier committed locally; host acceptance unconfirmed: {error}" + ))) + }); + let _ = reply.send(result); + } + } + }); + } + (Some(reply), None) => { + // The caller may hold a lock its event consumer needs. + let _ = reply.send(Ok(())); + } + (None, _) => {} } - thread_settings::emit_applied(&sess, sub.id.clone(), snapshot).await; + thread_settings::emit_applied(&sess, sub.id.clone(), commit.snapshot) + .await; } Err(error) => { let message = format!("invalid thread settings override: {error}"); diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 1809f27bba..72c69d087a 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -2152,6 +2152,8 @@ impl Session { /// /// The predicate must be fast and side-effect-free. It must not block, acquire /// other locks, or call back into `Session`. + /// + /// Returns any host receipt for the caller to await when it needs acknowledgment. async fn update_settings_if( &self, updates: SessionSettingsUpdate, @@ -2159,7 +2161,6 @@ impl Session { ) -> ConstraintResult> { let notify_config_contributors = !self.services.extensions.config_contributors().is_empty(); let windows_sandbox_changed; - let mut tier_acceptance = None; let ( commit, previous_config, @@ -2223,7 +2224,17 @@ impl Session { .update_configured_selections(&updated.environments); } state.session_configuration = updated; - if root_service_tier_changed && !state.shared_service_tier_pending { + let explicit_service_tier = updates.step_settings.service_tier.is_some(); + let retry_failed_handoff = updates.host_acceptance_in_reply + && explicit_service_tier + && state + .service_tier_acceptance + .as_ref() + .and_then(futures::future::Shared::peek) + .is_some_and(Result::is_err); + if (root_service_tier_changed || retry_failed_handoff) + && !state.shared_service_tier_pending + { state.service_tier_acceptance = crate::agent::host::propagate_config_update( self, AgentConfigUpdate::ServiceTier( @@ -2235,9 +2246,11 @@ impl Session { ), ); } - if root_service_tier_changed || updates.step_settings.service_tier.is_some() { - tier_acceptance = state.service_tier_acceptance.clone(); - } + let host_acceptance = if root_service_tier_changed || explicit_service_tier { + state.service_tier_acceptance.clone() + } else { + None + }; let new_config = notify_config_contributors .then(|| self.build_effective_session_config(&state.session_configuration)); let commit = SessionSettingsCommit { @@ -2245,6 +2258,7 @@ impl Session { snapshot: state .session_configuration .thread_settings_snapshot(&state.session_configuration.environments), + host_acceptance, }; ( commit, @@ -2280,9 +2294,6 @@ impl Session { if mcp_inputs_changed { self.schedule_mcp_prewarm(); } - if let Some(accepted) = tier_acceptance { - accepted.await; - } Ok(Some(commit)) } diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 54374144fd..04b6314a28 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -666,6 +666,8 @@ impl SessionConfiguration { pub(crate) struct SessionSettingsCommit { pub(crate) configuration: SessionConfiguration, pub(crate) snapshot: ThreadSettingsSnapshot, + /// Receipt for the host's local acceptance of this commit, when applicable. + pub(crate) host_acceptance: Option>>>, } #[derive(Default, Clone)] @@ -682,6 +684,8 @@ pub(crate) struct SessionSettingsUpdate { pub(crate) active_permission_profile: Option, pub(crate) windows_sandbox_level: Option, pub(crate) service_tier_for_turn: Option, + /// An explicit settings reply waits for host acceptance and can retry an observed failure. + pub(crate) host_acceptance_in_reply: bool, pub(crate) app_server_client_name: Option, pub(crate) app_server_client_version: Option, pub(crate) disabled_plugin_ids: Option>, diff --git a/codex-rs/core/src/session/thread_settings.rs b/codex-rs/core/src/session/thread_settings.rs index c3a82eff0f..facc5973bb 100644 --- a/codex-rs/core/src/session/thread_settings.rs +++ b/codex-rs/core/src/session/thread_settings.rs @@ -2,6 +2,7 @@ //! with checkpoints written directly to storage. use super::session::Session; +use super::session::SessionSettingsCommit; use super::session::SessionSettingsUpdate; use super::step_settings::StepSettingsUpdate; use crate::WithTurnExtensionData; @@ -77,8 +78,10 @@ impl Session { pub(super) async fn update( session: &Session, overrides: impl Into>, -) -> ConstraintResult { - let updates = prepare_update(overrides, &session.services.selected_capability_roots); + reply_requested: bool, +) -> ConstraintResult { + let mut updates = prepare_update(overrides, &session.services.selected_capability_roots); + updates.host_acceptance_in_reply = reply_requested; let commit = session.update_settings(updates).await?; let active_turn = session.active_turn.lock().await; // A registered running task retains its lineage so it can be recovered after a daemon @@ -89,7 +92,7 @@ pub(super) async fn update( { session.state.lock().await.last_started_turn_id = None; } - Ok(commit.snapshot) + Ok(commit) } /// Converts protocol overrides into the internal settings update shape. diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index 7757f386f9..b7d841c399 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -2980,7 +2980,7 @@ impl ThreadManagerState { } // Cancelling this wait must leave the registered runtime and the host handoff intact. if let Some(accepted) = tier_acceptance { - accepted.await; + let _ = accepted.await; } Ok(new_thread) } diff --git a/codex-rs/core/tests/suite/agent_websocket.rs b/codex-rs/core/tests/suite/agent_websocket.rs index 57ec34557c..6f8aaf4d6f 100644 --- a/codex-rs/core/tests/suite/agent_websocket.rs +++ b/codex-rs/core/tests/suite/agent_websocket.rs @@ -1,12 +1,15 @@ use super::agent_control::HostRuntime; +use anyhow::Context; use anyhow::Result; use codex_core::AgentRuntimeOwnership; use codex_core::AgentServiceTier; use codex_core::AgentTurnAdmission; use codex_core::CodexThreadSettingsOverrides; use codex_core::InheritedSessionId; +use codex_core::SteerSubmission; use codex_core::ThreadManager; use codex_core::TurnInputRequest; +use codex_core::TurnInputSubmission; use codex_core::TurnStartOptions; use codex_core::X_CODEX_ROUTING_HINT_HEADER; use codex_features::Feature; @@ -200,19 +203,17 @@ impl AgentServiceTier for HostTier { Box::pin(async { self.value.borrow().clone() }) } - fn update(&self, value: Option) -> BoxFuture<'static, ()> { + fn update(&self, value: Option) -> BoxFuture<'static, Result<(), String>> { let (polling, waiting) = oneshot::channel(); let (accept, accepted) = oneshot::channel(); - self.updates - .send(TierHandoff { - value, - waiting, - accept, - }) - .expect("tier handoff"); + let _ = self.updates.send(TierHandoff { + value, + waiting, + accept, + }); Box::pin(async move { let _ = polling.send(()); - accepted.await.expect("local tier acceptance"); + accepted.await.map_err(|_| "tier worker stopped".into()) }) } } @@ -256,6 +257,227 @@ async fn accept_committed_tier( }).await? } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn pending_host_tier_settings_reply_does_not_block_interrupt() -> Result<()> { + skip_if_no_network!(Ok(())); + + let completed = ev_completed("root-turn"); + let server = start_websocket_server(vec![ + vec![vec![ + ev_response_created("fixture-warm"), + ev_completed("fixture-warm"), + ]], + vec![ + vec![ev_response_created("root-warm"), ev_completed("root-warm")], + vec![ + ev_response_created("root-after-cancel"), + ev_assistant_message("still-running", "Ready."), + ev_completed("root-after-cancel"), + ], + vec![ev_response_created("root-turn"), completed.clone()], + ], + ]) + .await; + let release_completion = server.gate_response_event(&completed); + let test = test_codex() + .with_model("gpt-5.4") + .with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing()) + .with_config(|config| { + // Keep sampling behind the gate until the explicit interrupt. + config.features.disable(Feature::InstantInterrupt).unwrap(); + config + .features + .enable(Feature::FastMode) + .expect("fast mode"); + }) + .build_with_websocket_server(&server) + .await?; + tokio::time::timeout( + Duration::from_secs(5), + server.wait_for_request(/*connection_index*/ 0, /*request_index*/ 0), + ) + .await?; + + let (updates, mut handoffs) = mpsc::unbounded_channel(); + let tier = Arc::new(HostTier { + value: watch::Sender::new(None), + updates, + }); + let mut options = test.start_thread_options().await; + let root_id = test.thread_manager.reserve_thread_id(); + options.reserved_thread_id = Some(root_id); + options.config.service_tier = Some("default".into()); + options + .thread_extension_init + .insert(AgentRuntimeOwnership::Host); + options + .thread_extension_init + .insert(Arc::new(HostRuntime) as Arc); + options + .thread_extension_init + .insert(tier as Arc); + let manager = Arc::clone(&test.thread_manager); + let start = tokio::spawn(async move { manager.start_thread(options).await }); + let TierHandoff { + value, + waiting, + accept, + } = tokio::time::timeout(Duration::from_secs(10), handoffs.recv()) + .await? + .expect("startup tier handoff"); + assert_eq!(value.as_deref(), Some("default")); + tokio::time::timeout(Duration::from_secs(10), waiting).await??; + tokio::time::timeout( + Duration::from_secs(5), + server.wait_for_request(/*connection_index*/ 1, /*request_index*/ 0), + ) + .await?; + start.abort(); + let Err(error) = start.await else { + anyhow::bail!("startup caller returned while the receipt was pending"); + }; + assert!(error.is_cancelled()); + let root = test.thread_manager.get_thread(root_id).await?; + accept + .send(()) + .expect("registered root retains the handoff"); + + let request = TurnInputRequest::user_input(vec![UserInput::Text { + text: "Confirm the root still runs after startup cancellation.".into(), + text_elements: Vec::new(), + }]); + tokio::time::timeout(Duration::from_secs(10), root.start_or_steer_turn(request)) + .await + .context("root did not accept a turn after startup cancellation")??; + wait_for_successful_turn(&root).await; + + let request = TurnInputRequest::user_input(vec![UserInput::Text { + text: "Wait for an interrupt.".into(), + text_elements: Vec::new(), + }]) + .with_thread_settings(ThreadSettingsOverrides { + service_tier: Some(Some("flex".into())), + ..Default::default() + }); + let TurnInputSubmission::Started { turn_id, .. } = + tokio::time::timeout(Duration::from_secs(10), root.start_or_steer_turn(request)).await?? + else { + anyhow::bail!("expected a new root turn"); + }; + let request = tokio::time::timeout( + Duration::from_secs(10), + server.wait_for_request(/*connection_index*/ 1, /*request_index*/ 2), + ) + .await?; + let handoff = handoffs.try_recv().expect("turn tier handoff"); + assert_eq!(handoff.value.as_deref(), Some("flex")); + assert_eq!(request.body_json()["service_tier"], "flex"); + handoff.accept.send(()).expect("accept the turn tier"); + + // Keep both host acknowledgements pending while the same turn is running. + submit_thread_settings( + &root, + ThreadSettingsOverrides { + service_tier: Some(Some("default".into())), + ..Default::default() + }, + ) + .await?; + let no_reply = handoffs.try_recv().expect("no-reply tier handoff"); + assert_eq!(no_reply.value.as_deref(), Some("default")); + + let steer = TurnInputRequest::user_input(vec![UserInput::Text { + text: "Stay on this task.".into(), + text_elements: Vec::new(), + }]) + .with_thread_settings(ThreadSettingsOverrides { + service_tier: Some(Some("flex".into())), + ..Default::default() + }); + assert_eq!( + tokio::time::timeout( + Duration::from_secs(10), + root.steer_turn(steer, turn_id.clone()) + ) + .await??, + SteerSubmission::Steered { + turn_id: turn_id.clone() + } + ); + let steered = handoffs.try_recv().expect("steered tier handoff"); + assert_eq!(steered.value.as_deref(), Some("flex")); + + let thread = Arc::clone(&root); + let settings = tokio::spawn(async move { + thread + .update_thread_settings(ThreadSettingsOverrides { + service_tier: Some(Some("default".into())), + ..Default::default() + }) + .await + }); + let TierHandoff { + value, + waiting, + accept, + } = tokio::time::timeout(Duration::from_secs(10), handoffs.recv()) + .await? + .expect("tier handoff"); + assert_eq!(value.as_deref(), Some("default")); + tokio::time::timeout(Duration::from_secs(10), waiting).await??; + + root.submit(Op::Interrupt).await?; + wait_for_event(&root, |event| { + matches!(event, EventMsg::TurnAborted(aborted) if aborted.turn_id.as_deref() == Some(turn_id.as_str())) + }) + .await; + drop((no_reply, steered)); + drop(release_completion); + assert!( + !settings.is_finished(), + "interrupt does not accept the tier" + ); + drop(accept); + let error = tokio::time::timeout(Duration::from_secs(10), settings) + .await?? + .expect_err("the host dropped its receipt"); + assert_eq!( + error.to_string(), + "service tier committed locally; host acceptance unconfirmed: tier worker stopped" + ); + assert_eq!(root.thread_settings_snapshot().await.service_tier, value); + + tokio::time::timeout( + Duration::from_secs(10), + root.restore_thread_settings(CodexThreadSettingsOverrides { + service_tier: Some(value.clone()), + ..Default::default() + }), + ) + .await??; + assert!( + handoffs.is_empty(), + "restore does not retry a failed handoff" + ); + + accept_committed_tier( + root.update_thread_settings(ThreadSettingsOverrides { + service_tier: Some(value), + ..Default::default() + }), + &mut handoffs, + &test.thread_manager, + root_id, + "default", + ) + .await?; + + root.shutdown_and_wait().await?; + test.codex.shutdown_and_wait().await?; + server.shutdown().await; + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn detached_host_child_inherits_root_service_tier_at_prewarm_and_each_step() -> Result<()> { skip_if_no_network!(Ok(()));