From f5ac7a757e8f6f4fe922f3daab0fcaba02f1fa72 Mon Sep 17 00:00:00 2001 From: rhan-oai Date: Sat, 7 Nov 2026 06:41:12 +0000 Subject: [PATCH] Support asynchronous host service-tier reads and acceptance (#60075) ## What changed Make `AgentServiceTier` reads asynchronous and have updates return a receipt for local acceptance. Await acceptance during root startup and settings updates outside state and registry locks, sharing the receipt with explicit same-tier retries. Await shared-tier reads during child preparation, startup prewarm, and step capture. Allow cancellation while capturing a step, and keep host propagation independent of the acceptance wait. ## Testing Extend the host-child WebSocket test to cover delayed acceptance, same-tier retries without duplicate updates, accessible settings while acceptance is pending, and tier inheritance during prewarm and subsequent steps. GitOrigin-RevId: 0e2e1899834c37684e904c4a36e66344a0c6c7f3 --- codex-rs/core/src/agent/child_config.rs | 3 +- codex-rs/core/src/agent/host/mod.rs | 1 + codex-rs/core/src/agent/host/service_tier.rs | 60 ++++++---- codex-rs/core/src/session/mod.rs | 16 ++- codex-rs/core/src/session/startup_prewarm.rs | 3 +- codex-rs/core/src/state/session.rs | 4 + codex-rs/core/src/thread_manager.rs | 24 ++-- codex-rs/core/tests/suite/agent_websocket.rs | 114 ++++++++++++++++--- 8 files changed, 176 insertions(+), 49 deletions(-) diff --git a/codex-rs/core/src/agent/child_config.rs b/codex-rs/core/src/agent/child_config.rs index 1e2d4363c6..22f741c608 100644 --- a/codex-rs/core/src/agent/child_config.rs +++ b/codex-rs/core/src/agent/child_config.rs @@ -282,7 +282,8 @@ pub(crate) async fn apply_spawn_agent_service_tier( session: &Session, config: &mut Config, ) -> Result<(), String> { - let Some(service_tier) = crate::agent::host::service_tier(session, config.service_tier.clone()) + let Some(service_tier) = + crate::agent::host::service_tier(session, config.service_tier.clone()).await else { config.service_tier = None; return Ok(()); diff --git a/codex-rs/core/src/agent/host/mod.rs b/codex-rs/core/src/agent/host/mod.rs index 0fe8d8b13d..06a15b39c0 100644 --- a/codex-rs/core/src/agent/host/mod.rs +++ b/codex-rs/core/src/agent/host/mod.rs @@ -22,6 +22,7 @@ pub use mcp_metadata::AgentMcpMetadataRecipient; pub use ownership::AgentRuntimeOwnership; pub(crate) use ownership::AgentRuntimeServices; pub use service_tier::AgentServiceTier; +pub(crate) use service_tier::ServiceTierAcceptance; pub(crate) use service_tier::host_child_uses_shared_service_tier; pub(crate) use service_tier::propagate_config_update; pub(crate) use service_tier::service_tier; diff --git a/codex-rs/core/src/agent/host/service_tier.rs b/codex-rs/core/src/agent/host/service_tier.rs index c0c11acd5e..8f15f89030 100644 --- a/codex-rs/core/src/agent/host/service_tier.rs +++ b/codex-rs/core/src/agent/host/service_tier.rs @@ -4,21 +4,32 @@ use crate::agent::api::AgentConfigUpdate; use crate::agent::host::AgentRuntimeServices; use crate::session::session::Session; use codex_protocol::SessionId; +use futures::FutureExt; +use futures::future::BoxFuture; +use futures::future::Shared; -/// The host's local view of the service tier shared within an agent session. +pub(crate) type ServiceTierAcceptance = Shared>; + +/// The host's view of the service tier shared within an agent session. /// /// Install an `Arc` in thread extension data before startup. Host-owned /// children with a provider reject explicit per-turn tier overrides. Host-owned runtimes without /// a provider retain their configured tier and do not propagate shared tier updates. /// Controller-owned runtimes use their controller. -/// Both methods must be cheap, synchronous and non-reentrant: Core may call them while -/// holding configuration or registration locks. They must perform no network I/O. The host publishes -/// committed root settings into its cache and arranges propagation to other runtimes. -/// Children read the latest cached value at the existing preparation and step boundaries; +/// Core awaits reads during preparation and step capture. A host may answer from a cache. +/// Hosts own cross-runtime propagation and recovery. Until an update reaches a runtime, +/// its `current` may return the prior tier. +/// Children read the latest value at the existing preparation and step boundaries; /// updates do not change requests already in flight. `None` restores the default tier. pub trait AgentServiceTier: Send + Sync { - fn current(&self) -> Option; - fn update(&self, service_tier: Option); + fn current(&self) -> BoxFuture<'_, Option>; + + /// 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, ()>; } pub(crate) fn host_child_uses_shared_service_tier(session: &Session) -> bool { @@ -26,34 +37,39 @@ pub(crate) fn host_child_uses_shared_service_tier(session: &Session) -> bool { && session.session_id() != SessionId::from(session.thread_id()) } -pub(crate) fn service_tier( +pub(crate) async fn service_tier( session: &Session, configured_service_tier: Option, ) -> Option { match &session.agent_runtime { AgentRuntimeServices::Controller => session.services.agent_control.service_tier(), - AgentRuntimeServices::Host(services) => services - .service_tier - .as_ref() - .map_or(configured_service_tier, |provider| provider.current()), + AgentRuntimeServices::Host(services) => match &services.service_tier { + Some(provider) => provider.current().await, + None => configured_service_tier, + }, } } -pub(crate) fn propagate_config_update(session: &Session, update: AgentConfigUpdate) { +pub(crate) fn propagate_config_update( + session: &Session, + update: AgentConfigUpdate, +) -> Option { match &session.agent_runtime { - AgentRuntimeServices::Controller => session - .services - .agent_control - .propagate_config_update(update), + AgentRuntimeServices::Controller => { + session + .services + .agent_control + .propagate_config_update(update); + None + } AgentRuntimeServices::Host(services) => { // Only the root can publish shared settings. Compare session and thread IDs // because a detached child may have no structural parent ID. - if session.session_id() == SessionId::from(session.thread_id()) - && let Some(provider) = &services.service_tier - { - let AgentConfigUpdate::ServiceTier(tier) = update; - provider.update(tier); + if session.session_id() != SessionId::from(session.thread_id()) { + return None; } + let AgentConfigUpdate::ServiceTier(tier) = update; + Some(services.service_tier.as_ref()?.update(tier).shared()) } } } diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index c52dd4e9c7..814799a59f 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -2160,6 +2160,7 @@ 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, @@ -2224,7 +2225,7 @@ impl Session { } state.session_configuration = updated; if root_service_tier_changed && !state.shared_service_tier_pending { - crate::agent::host::propagate_config_update( + state.service_tier_acceptance = crate::agent::host::propagate_config_update( self, AgentConfigUpdate::ServiceTier( state @@ -2235,6 +2236,9 @@ impl Session { ), ); } + if root_service_tier_changed || updates.step_settings.service_tier.is_some() { + tier_acceptance = state.service_tier_acceptance.clone(); + } let new_config = notify_config_contributors .then(|| self.build_effective_session_config(&state.session_configuration)); let commit = SessionSettingsCommit { @@ -2277,6 +2281,9 @@ impl Session { if mcp_inputs_changed { self.schedule_mcp_prewarm(); } + if let Some(accepted) = tier_acceptance { + accepted.await; + } Ok(Some(commit)) } @@ -4313,8 +4320,11 @@ impl Session { ) }; if self.inherits_root_service_tier(&turn_context.session_source) { - let root_service_tier = - crate::agent::host::service_tier(self, settings.selected().service_tier.clone()); + let root_service_tier = tokio::select! { + biased; + tier = crate::agent::host::service_tier(self, settings.selected().service_tier.clone()) => tier, + _ = cancellation_token.cancelled() => return Err(CodexErr::TurnAborted), + }; if settings.selected().service_tier != root_service_tier { let mut selected = settings.selected().clone(); selected.service_tier = root_service_tier; diff --git a/codex-rs/core/src/session/startup_prewarm.rs b/codex-rs/core/src/session/startup_prewarm.rs index 0207e90609..3aa83e6db1 100644 --- a/codex-rs/core/src/session/startup_prewarm.rs +++ b/codex-rs/core/src/session/startup_prewarm.rs @@ -347,7 +347,8 @@ async fn schedule_startup_prewarm_inner( crate::agent::host::service_tier( &session, startup_turn_context.initial_settings.service_tier.clone(), - ), + ) + .await, &session.features(), &preconnect_model_info, ) diff --git a/codex-rs/core/src/state/session.rs b/codex-rs/core/src/state/session.rs index d702c09da3..e90bd010c6 100644 --- a/codex-rs/core/src/state/session.rs +++ b/codex-rs/core/src/state/session.rs @@ -16,6 +16,7 @@ use super::auto_compact_window::AutoCompactWindowIds; use super::auto_compact_window::AutoCompactWindowSnapshot; use super::auto_compact_window::ContextWindowIdGenerator; use crate::TurnStartOptions; +use crate::agent::host::ServiceTierAcceptance; use crate::context_manager::ContextManager; use crate::context_manager::HistoryReplacement; use crate::session::PreviousTurnSettings; @@ -74,6 +75,8 @@ pub(crate) struct SessionState { pub(crate) session_configuration: SessionConfiguration, /// Shared root tier publication starts when the runtime wins registration. pub(crate) shared_service_tier_pending: bool, + /// Receipt for the current host root tier, retained for callers retrying that commit. + pub(crate) service_tier_acceptance: Option, /// Catalog to carry across compaction, independent of an explicit resume override. pub(crate) persisted_dynamic_tools: Arc<[DynamicToolSpec]>, /// Plugin selection of the last admitted task; settings updates take effect on the next task. @@ -144,6 +147,7 @@ impl SessionState { persisted_dynamic_tools: Arc::clone(&session_configuration.dynamic_tools), session_configuration, shared_service_tier_pending: false, + service_tier_acceptance: None, base_instructions_provenance: None, history, pending_audio_duration_ids: BTreeSet::new(), diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index 8cbc7bc9b2..af6c240112 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -16,6 +16,7 @@ use crate::agent::LocalAgentControl; use crate::agent::api::AgentControl; use crate::agent::control::AgentControlInit; use crate::agent::host::InheritedSessionId; +use crate::agent::host::ServiceTierAcceptance; use crate::agents_md_manager::SessionInstructions; use crate::attestation::AttestationProvider; use crate::codex_thread::CodexThread; @@ -2894,8 +2895,8 @@ impl ThreadManagerState { let finalize_result = self .finalize_thread_spawn(session, io, tracked_session_source) .await; - let new_thread = match finalize_result { - Ok(new_thread) => new_thread, + let (new_thread, tier_acceptance) = match finalize_result { + Ok(registered) => registered, Err(error) => { if let Some(startup_guard) = startup_guard { startup_guard.cleanup().await; @@ -2916,6 +2917,10 @@ impl ThreadManagerState { if let Some(startup_guard) = startup_guard { startup_guard.disarm(); } + // Cancelling this wait must leave the registered runtime and the host handoff intact. + if let Some(accepted) = tier_acceptance { + accepted.await; + } Ok(new_thread) } @@ -2924,7 +2929,7 @@ impl ThreadManagerState { session: Arc, io: SessionIo, session_source: SessionSource, - ) -> CodexResult { + ) -> CodexResult<(NewThread, Option)> { let thread_id = session.thread_id(); let event = io.next_event().await?; let session_configured = match event { @@ -2965,15 +2970,18 @@ impl ThreadManagerState { session_source, )); e.insert(thread.clone()); - if let Some((mut state, update)) = pending_service_tier { - crate::agent::host::propagate_config_update(&session, update); + let tier_acceptance = pending_service_tier.and_then(|(mut state, update)| { + state.service_tier_acceptance = + crate::agent::host::propagate_config_update(&session, update); state.shared_service_tier_pending = false; - } - return Ok(NewThread { + state.service_tier_acceptance.clone() + }); + let new_thread = NewThread { thread_id, thread, session_configured, - }); + }; + return Ok((new_thread, tier_acceptance)); } } diff --git a/codex-rs/core/tests/suite/agent_websocket.rs b/codex-rs/core/tests/suite/agent_websocket.rs index 77b513acee..57ec34557c 100644 --- a/codex-rs/core/tests/suite/agent_websocket.rs +++ b/codex-rs/core/tests/suite/agent_websocket.rs @@ -3,12 +3,15 @@ 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::ThreadManager; use codex_core::TurnInputRequest; use codex_core::TurnStartOptions; use codex_core::X_CODEX_ROUTING_HINT_HEADER; use codex_features::Feature; use codex_login::CodexAuth; +use codex_protocol::ThreadId; use codex_protocol::config_types::ServiceTier; use codex_protocol::dynamic_tools::DynamicToolFunctionSpec; use codex_protocol::dynamic_tools::DynamicToolResponse; @@ -35,12 +38,15 @@ use core_test_support::submit_thread_settings; use core_test_support::test_codex::test_codex; use core_test_support::wait_for_event; use core_test_support::wait_for_successful_turn; +use futures::future::BoxFuture; use pretty_assertions::assert_eq; use serde_json::Value; use serde_json::json; use std::sync::Arc; -use std::sync::Mutex; use std::time::Duration; +use tokio::sync::mpsc; +use tokio::sync::oneshot; +use tokio::sync::watch; const WS_V2_BETA_HEADER_VALUE: &str = "responses_websockets=2026-02-06"; @@ -178,19 +184,78 @@ async fn spawned_agent_prewarm_handshake_inherits_effective_root_service_tier( Ok(()) } -#[derive(Default)] -struct HostTier(Mutex>); +struct HostTier { + value: watch::Sender>, + updates: mpsc::UnboundedSender, +} + +struct TierHandoff { + value: Option, + waiting: oneshot::Receiver<()>, + accept: oneshot::Sender<()>, +} impl AgentServiceTier for HostTier { - fn current(&self) -> Option { - self.0.lock().expect("host tier").clone() + fn current(&self) -> BoxFuture<'_, Option> { + Box::pin(async { self.value.borrow().clone() }) } - fn update(&self, value: Option) { - *self.0.lock().expect("host tier") = value; + fn update(&self, value: Option) -> BoxFuture<'static, ()> { + let (polling, waiting) = oneshot::channel(); + let (accept, accepted) = oneshot::channel(); + self.updates + .send(TierHandoff { + value, + waiting, + accept, + }) + .expect("tier handoff"); + Box::pin(async move { + let _ = polling.send(()); + accepted.await.expect("local tier acceptance"); + }) } } +async fn accept_committed_tier( + operation: impl Future>, + handoffs: &mut mpsc::UnboundedReceiver, + manager: &ThreadManager, + root_id: ThreadId, + expected: &str, +) -> Result<(T, Option)> { + tokio::time::timeout(Duration::from_secs(10), async { + tokio::pin!(operation); + let (root, value, accept) = tokio::select! { + biased; + _ = &mut operation => anyhow::bail!("root returned before local tier acceptance"), + handoff = async { + let TierHandoff { value, waiting, accept } = handoffs.recv().await.expect("tier handoff"); + waiting.await.expect("Core polled the receipt"); + let root = manager.get_thread(root_id).await?; + assert_eq!(value.as_deref(), Some(expected)); + anyhow::Ok((root, value, accept)) + } => handoff?, + }; + // Use the direct settings API so the retry cannot queue behind the first caller. + let retry = root.restore_thread_settings(CodexThreadSettingsOverrides { + service_tier: Some(value.clone()), + ..Default::default() + }); + tokio::pin!(retry); + tokio::select! { + biased; + _ = &mut retry => anyhow::bail!("same-tier retry returned before local acceptance"), + snapshot = root.thread_settings_snapshot() => assert_eq!(snapshot.service_tier, value), + } + accept.send(()).expect("Core is waiting for acceptance"); + retry.await?; + let result = operation.await?; + assert!(handoffs.is_empty(), "same-tier retry published another update"); + Ok((result, value)) + }).await? +} + #[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(())); @@ -242,8 +307,14 @@ async fn detached_host_child_inherits_root_service_tier_at_prewarm_and_each_step ) .await?; - let tier = Arc::new(HostTier::default()); + 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("priority".into()); // The host root only changes settings. Keep WebSocket traffic on the child. options.config.model_provider.supports_websockets = false; @@ -257,7 +328,16 @@ async fn detached_host_child_inherits_root_service_tier_at_prewarm_and_each_step .thread_extension_init .insert(tier.clone() as Arc); let thread_extension_init = options.thread_extension_init.clone(); - let root = test.thread_manager.start_thread(options).await?; + let (root, initial_tier) = accept_committed_tier( + test.thread_manager.start_thread(options), + &mut handoffs, + &test.thread_manager, + root_id, + "priority", + ) + .await?; + assert_eq!(tier.current().await, None); + tier.value.send_replace(initial_tier); let mut options = test.start_thread_options().await; options.session_source = Some(SessionSource::Custom("host".into())); @@ -295,7 +375,7 @@ async fn detached_host_child_inherits_root_service_tier_at_prewarm_and_each_step }, ) .await?; - assert_eq!(tier.current().as_deref(), Some("priority")); + assert_eq!(tier.current().await.as_deref(), Some("priority")); let request = TurnInputRequest::user_input(vec![UserInput::Text { text: "Wait for the tier change.".into(), @@ -320,14 +400,20 @@ async fn detached_host_child_inherits_root_service_tier_at_prewarm_and_each_step |event| matches!(event, EventMsg::DynamicToolCallRequest(call) if call.call_id == "hold"), ) .await; - submit_thread_settings( - &root.thread, - ThreadSettingsOverrides { + let ((), accepted_tier) = accept_committed_tier( + root.thread.update_thread_settings(ThreadSettingsOverrides { service_tier: Some(Some("flex".into())), ..Default::default() - }, + }), + &mut handoffs, + &test.thread_manager, + root_id, + "flex", ) .await?; + // The root has its acknowledgment. The child observes the change after host propagation. + assert_eq!(tier.current().await.as_deref(), Some("priority")); + tier.value.send_replace(accepted_tier); child .thread .submit(Op::DynamicToolResponse {