diff --git a/codex-rs/core/src/agent/child_config.rs b/codex-rs/core/src/agent/child_config.rs index 4aa8a725e8..1e2d4363c6 100644 --- a/codex-rs/core/src/agent/child_config.rs +++ b/codex-rs/core/src/agent/child_config.rs @@ -3,6 +3,7 @@ //! Spawn and reload share live runtime policy; role and model precedence, full-history //! inheritance, and validation messages remain the same for each multi-agent version. +use crate::agent::host::AgentRuntimeServices; use crate::agent::role::DEFAULT_ROLE_NAME; use crate::agent::role::apply_role_to_config; use crate::config::Config; @@ -69,6 +70,14 @@ pub(crate) async fn prepare_agent_spawn_config( options.reasoning_effort, ) .await?; + let shares_service_tier = match &session.agent_runtime { + AgentRuntimeServices::Controller => true, + AgentRuntimeServices::Host(services) => services.service_tier.is_some(), + }; + if !shares_service_tier { + // Inherit the parent's thread preference before applying child role overrides. + config.service_tier = session.thread_settings_snapshot().await.service_tier; + } if !options.full_history_fork || (options.version == SpawnConfigVersion::V2 && options.role_name.is_some()) { @@ -82,7 +91,9 @@ pub(crate) async fn prepare_agent_spawn_config( .clone_from(&turn.developer_instructions); } } - apply_spawn_agent_service_tier(session, &mut config).await?; + if shares_service_tier { + apply_spawn_agent_service_tier(session, &mut config).await?; + } apply_spawn_agent_runtime_overrides(&mut config, turn)?; // Remember an applied configured default so cold reload reapplies its restrictions. @@ -271,7 +282,8 @@ pub(crate) async fn apply_spawn_agent_service_tier( session: &Session, config: &mut Config, ) -> Result<(), String> { - let Some(service_tier) = session.services.agent_control.service_tier() else { + let Some(service_tier) = crate::agent::host::service_tier(session, config.service_tier.clone()) + else { config.service_tier = None; return Ok(()); }; diff --git a/codex-rs/core/src/agent/control/residency.rs b/codex-rs/core/src/agent/control/residency.rs index 3ea01481ab..e16b1ccab4 100644 --- a/codex-rs/core/src/agent/control/residency.rs +++ b/codex-rs/core/src/agent/control/residency.rs @@ -4,6 +4,7 @@ use super::LocalAgentControl; use super::LocalAgentRuntime; use super::runtime::AgentTreeMembership; use crate::agent::AgentStatus; +use crate::agent::host::AgentRuntimeServices; use crate::agent::registry::EvictedEnvironments; use crate::codex_thread::CodexThread; use crate::config::Config; @@ -84,6 +85,10 @@ impl ThreadManagerState { // Raw bindings keep the original pending owner configuration alive. let environments = EvictedEnvironments { config: Arc::new(config), + has_independent_service_tier: matches!( + &thread.session.agent_runtime, + AgentRuntimeServices::Host(services) if services.service_tier.is_none() + ), selections: thread.environment_selections().await, bindings: thread .session diff --git a/codex-rs/core/src/agent/control/spawn.rs b/codex-rs/core/src/agent/control/spawn.rs index e557f814c9..7449da552d 100644 --- a/codex-rs/core/src/agent/control/spawn.rs +++ b/codex-rs/core/src/agent/control/spawn.rs @@ -286,6 +286,10 @@ impl LocalAgentControl { let mut environment_selections = evicted_environments .as_ref() .map(|environments| environments.selections.clone()); + let independent_service_tier = evicted_environments + .as_ref() + .filter(|evicted| evicted.has_independent_service_tier) + .map(|evicted| evicted.config.service_tier.clone()); let evicted_environments = evicted_environments.map(|environments| environments.bindings); let stored_thread = state @@ -384,7 +388,7 @@ impl LocalAgentControl { CodexErr::InvalidRequest(format!("permission_profile is invalid: {err}")) })?; } - config.service_tier = self.root_service_tier(); + config.service_tier = independent_service_tier.unwrap_or_else(|| self.root_service_tier()); if let Some(model) = stored_model { config.model = Some(model); } diff --git a/codex-rs/core/src/agent/host/mod.rs b/codex-rs/core/src/agent/host/mod.rs index 4eccdcaa77..0fe8d8b13d 100644 --- a/codex-rs/core/src/agent/host/mod.rs +++ b/codex-rs/core/src/agent/host/mod.rs @@ -3,6 +3,7 @@ mod completion; mod invocation; mod mcp_metadata; mod ownership; +mod service_tier; mod session_reader; mod spawn; mod types; @@ -20,6 +21,10 @@ pub(crate) use invocation::tool_invocation; pub use mcp_metadata::AgentMcpMetadataRecipient; pub use ownership::AgentRuntimeOwnership; pub(crate) use ownership::AgentRuntimeServices; +pub use service_tier::AgentServiceTier; +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; pub use session_reader::AgentRef; pub use session_reader::AgentSessionReader; pub use types::HostThreadSpawnReservation; diff --git a/codex-rs/core/src/agent/host/ownership.rs b/codex-rs/core/src/agent/host/ownership.rs index 37125d04cb..3cfb43e6f4 100644 --- a/codex-rs/core/src/agent/host/ownership.rs +++ b/codex-rs/core/src/agent/host/ownership.rs @@ -1,5 +1,6 @@ //! Captures coordination providers once, after the thread's startup contributors finish. +use super::AgentServiceTier; use super::AgentSessionReader; use super::AgentTurnAdmission; use super::AgentUsageAccounting; @@ -17,6 +18,7 @@ pub enum AgentRuntimeOwnership { /// Requires an [`AgentTurnAdmission`] provider. Core skips controller rollout-budget /// accounting and budget reminders. Install an [`AgentUsageAccounting`] provider /// to receive usage reports and an [`AgentSessionReader`] to render the V2 subagent roster. + /// An optional [`AgentServiceTier`] shares the root's committed tier with its children. Host, } @@ -29,6 +31,7 @@ pub(crate) struct HostAgentServices { pub(crate) session_reader: Option>, pub(crate) turn_admission: Arc, pub(crate) usage_accounting: Option>, + pub(crate) service_tier: Option>, } impl AgentRuntimeServices { @@ -51,6 +54,7 @@ impl AgentRuntimeServices { .get::>() .as_deref() .cloned(), + service_tier: data.get::>().as_deref().cloned(), })) } } diff --git a/codex-rs/core/src/agent/host/service_tier.rs b/codex-rs/core/src/agent/host/service_tier.rs new file mode 100644 index 0000000000..c0c11acd5e --- /dev/null +++ b/codex-rs/core/src/agent/host/service_tier.rs @@ -0,0 +1,59 @@ +//! Select host-owned shared service-tier state at existing runtime configuration boundaries. + +use crate::agent::api::AgentConfigUpdate; +use crate::agent::host::AgentRuntimeServices; +use crate::session::session::Session; +use codex_protocol::SessionId; + +/// The host's local 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; +/// 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); +} + +pub(crate) fn host_child_uses_shared_service_tier(session: &Session) -> bool { + matches!(&session.agent_runtime, AgentRuntimeServices::Host(services) if services.service_tier.is_some()) + && session.session_id() != SessionId::from(session.thread_id()) +} + +pub(crate) 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()), + } +} + +pub(crate) fn propagate_config_update(session: &Session, update: AgentConfigUpdate) { + match &session.agent_runtime { + AgentRuntimeServices::Controller => session + .services + .agent_control + .propagate_config_update(update), + 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); + } + } + } +} diff --git a/codex-rs/core/src/agent/registry.rs b/codex-rs/core/src/agent/registry.rs index 1b2855fcd8..eeed8f450d 100644 --- a/codex-rs/core/src/agent/registry.rs +++ b/codex-rs/core/src/agent/registry.rs @@ -43,6 +43,7 @@ struct ActiveAgents { #[derive(Clone)] pub(crate) struct EvictedEnvironments { pub(crate) config: Arc, + pub(crate) has_independent_service_tier: bool, pub(crate) selections: Vec, pub(crate) bindings: TurnEnvironmentSnapshot, } diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index 0df7afaefc..b0afbbbebe 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -162,6 +162,7 @@ mod thread_store_builder; pub use agent::host::AgentMcpMetadataRecipient; pub use agent::host::AgentRef; pub use agent::host::AgentRuntimeOwnership; +pub use agent::host::AgentServiceTier; pub use agent::host::AgentSessionReader; pub use agent::host::AgentTurnAdmission; pub use agent::host::AgentTurnAdmissionRequest; diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 91b361126b..13ba7bdd54 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -204,6 +204,7 @@ use futures::prelude::*; use rmcp::model::RequestId; use serde_json::Value; use tokio::sync::Mutex; +use tokio::sync::MutexGuard; use tokio::sync::RwLock; use tokio::sync::oneshot; use tokio::sync::watch; @@ -2206,8 +2207,9 @@ impl Session { .update_configured_selections(&updated.environments); } state.session_configuration = updated; - if root_service_tier_changed { - self.services.agent_control.propagate_config_update( + if root_service_tier_changed && !state.shared_service_tier_pending { + crate::agent::host::propagate_config_update( + self, AgentConfigUpdate::ServiceTier( state .session_configuration @@ -2312,6 +2314,25 @@ impl Session { .thread_settings_snapshot(&state.session_configuration.environments) } + /// Keep this guard through registration and publication to order the initial tier + /// with settings commits from the already-running submission loop. + pub(crate) async fn lock_pending_service_tier( + &self, + ) -> Option<(MutexGuard<'_, SessionState>, AgentConfigUpdate)> { + let state = self.state.lock().await; + if !state.shared_service_tier_pending { + return None; + } + let update = AgentConfigUpdate::ServiceTier( + state + .session_configuration + .step_settings + .service_tier + .clone(), + ); + Some((state, update)) + } + pub(crate) async fn restorable_thread_settings(&self) -> CodexThreadSettingsOverrides { let state = self.state.lock().await; state @@ -4248,6 +4269,18 @@ impl Session { .await } + fn inherits_root_service_tier(&self, session_source: &SessionSource) -> bool { + match &self.agent_runtime { + crate::agent::host::AgentRuntimeServices::Controller => matches!( + session_source, + SessionSource::SubAgent(SubAgentSource::ThreadSpawn { .. }) + ), + crate::agent::host::AgentRuntimeServices::Host(_) => { + crate::agent::host::host_child_uses_shared_service_tier(self) + } + } + } + #[expect( clippy::await_holding_invalid_type, reason = "step settings and both environment snapshots must be captured under active_turn then state" @@ -4278,11 +4311,9 @@ impl Session { .then(|| self.services.turn_environments.configured_snapshot_now()), ) }; - if matches!( - turn_context.session_source, - SessionSource::SubAgent(SubAgentSource::ThreadSpawn { .. }) - ) { - let root_service_tier = self.services.agent_control.service_tier(); + 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()); 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/session.rs b/codex-rs/core/src/session/session.rs index be85f1fb3c..660428d728 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -1092,6 +1092,8 @@ impl Session { } } } + let has_provided_agent_control = + matches!(&agent_control, AgentControlInit::Provided { .. }); let (agent_control, local_agent_runtime): (Arc, _) = match agent_control { AgentControlInit::Local(control) => { let control = control.with_session_id( @@ -1878,6 +1880,13 @@ impl Session { // This placeholder is used only for teardown, before the session is published. Err(err) => (AgentRuntimeServices::Controller, Some(err)), }; + state.shared_service_tier_pending = parent_thread_id.is_none() + && match &agent_runtime { + AgentRuntimeServices::Controller => has_provided_agent_control, + AgentRuntimeServices::Host(services) => { + session_id == SessionId::from(thread_id) && services.service_tier.is_some() + } + }; let executed_tool_calls = crate::state::ExecutedToolCalls::new(&config.features, &initial_history); let codex_responses_headers = diff --git a/codex-rs/core/src/session/startup_prewarm.rs b/codex-rs/core/src/session/startup_prewarm.rs index 5d3fa4d681..0207e90609 100644 --- a/codex-rs/core/src/session/startup_prewarm.rs +++ b/codex-rs/core/src/session/startup_prewarm.rs @@ -29,8 +29,6 @@ use codex_otel::STARTUP_PREWARM_DURATION_METRIC; use codex_otel::SessionTelemetry; use codex_protocol::error::Result as CodexResult; use codex_protocol::models::BaseInstructions; -use codex_protocol::protocol::SessionSource; -use codex_protocol::protocol::SubAgentSource; #[derive(Clone, Copy, Debug)] pub(crate) enum PrewarmInput { @@ -342,12 +340,14 @@ async fn schedule_startup_prewarm_inner( let preconnect_model_info = Arc::clone(startup_turn_context.model_info()); // Spawned subagents inherit the root's selection, with the same feature and model filtering // that capture applies to the actual request. - let preconnect_service_tier = if matches!( - startup_turn_context.session_source, - SessionSource::SubAgent(SubAgentSource::ThreadSpawn { .. }) - ) { + let inherits_root_service_tier = + session.inherits_root_service_tier(&startup_turn_context.session_source); + let preconnect_service_tier = if inherits_root_service_tier { crate::session::get_service_tier( - session.services.agent_control.service_tier(), + crate::agent::host::service_tier( + &session, + startup_turn_context.initial_settings.service_tier.clone(), + ), &session.features(), &preconnect_model_info, ) diff --git a/codex-rs/core/src/session/turn_input.rs b/codex-rs/core/src/session/turn_input.rs index ff6abd45b7..5e32f8d984 100644 --- a/codex-rs/core/src/session/turn_input.rs +++ b/codex-rs/core/src/session/turn_input.rs @@ -245,6 +245,14 @@ pub(super) async fn handle( submission_id: String, ) -> CodexResult { let request = request.into(); + if request.request.start.service_tier.is_some() + && crate::agent::host::host_child_uses_shared_service_tier(session) + { + return Err(CodexErr::InvalidRequest( + "per-turn service tier overrides are unavailable for a child using the shared service tier" + .into(), + )); + } if let SubmittedTurnInput::ResponseItem(item) | SubmittedTurnInput::AnnotatedResponseItem(AnnotatedResponseItem { item, .. }) = &request.request.input diff --git a/codex-rs/core/src/state/session.rs b/codex-rs/core/src/state/session.rs index acd1077fe4..d702c09da3 100644 --- a/codex-rs/core/src/state/session.rs +++ b/codex-rs/core/src/state/session.rs @@ -72,6 +72,8 @@ impl ReasoningEffortPin { /// Persistent, session-scoped state previously stored directly on `Session`. 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, /// 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. @@ -141,6 +143,7 @@ impl SessionState { active_disabled_plugin_ids: Vec::new(), persisted_dynamic_tools: Arc::clone(&session_configuration.dynamic_tools), session_configuration, + shared_service_tier_pending: false, 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 19d93d5aca..6556313eb5 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -13,7 +13,6 @@ use crate::CodexAppsToolsCache; use crate::CodexThreadSettingsOverrides; use crate::agent::AgentStatus; use crate::agent::LocalAgentControl; -use crate::agent::api::AgentConfigUpdate; use crate::agent::api::AgentControl; use crate::agent::control::AgentControlInit; use crate::agent::host::InheritedSessionId; @@ -2513,11 +2512,6 @@ impl ThreadManagerState { }) => Some(*parent_thread_id), _ => None, }); - // Host controllers can already be shared with live threads. Publish root settings - // only after registration succeeds, so a failed start cannot update their tree. - let initial_host_config = (self.agent_control_factory.is_some() - && parent_thread_id.is_none()) - .then(|| AgentConfigUpdate::ServiceTier(config.service_tier.clone())); // Both resume entry points must restore identities before children can be loaded lazily. if let InitialHistory::Resumed(resumed) = &initial_history && initial_history.get_multi_agent_version() == Some(MultiAgentVersion::V2) @@ -2798,7 +2792,7 @@ impl ThreadManagerState { } } let finalize_result = self - .finalize_thread_spawn(session, io, tracked_session_source, initial_host_config) + .finalize_thread_spawn(session, io, tracked_session_source) .await; let new_thread = match finalize_result { Ok(new_thread) => new_thread, @@ -2830,7 +2824,6 @@ impl ThreadManagerState { session: Arc, io: SessionIo, session_source: SessionSource, - initial_host_config: Option, ) -> CodexResult { let thread_id = session.thread_id(); let event = io.next_event().await?; @@ -2851,23 +2844,31 @@ impl ThreadManagerState { .mcp_ema_cancellation() .clone(); { + let pending_service_tier = if self.agent_control_factory.is_some() + || matches!( + &session.agent_runtime, + crate::agent::host::AgentRuntimeServices::Host(services) + if services.service_tier.is_some() + ) { + session.lock_pending_service_tier().await + } else { + None + }; let mut threads = self.threads.write().await; if let std::collections::hash_map::Entry::Vacant(e) = threads.entry(thread_id) { - if let Some(update) = initial_host_config { - session - .services - .agent_control - .propagate_config_update(update); - } ema_cancellation.publish(); let thread = Arc::new(CodexThread::new( - session, + Arc::clone(&session), io, ThreadStartupMetadata::from(&session_configured), session_configured.rollout_path.clone(), session_source, )); e.insert(thread.clone()); + if let Some((mut state, update)) = pending_service_tier { + crate::agent::host::propagate_config_update(&session, update); + state.shared_service_tier_pending = false; + } return Ok(NewThread { thread_id, thread, diff --git a/codex-rs/core/src/thread_manager_tests.rs b/codex-rs/core/src/thread_manager_tests.rs index e0a741a263..db8709b181 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/codex-rs/core/src/thread_manager_tests.rs @@ -33,6 +33,7 @@ use codex_protocol::capabilities::CapabilityRootLocation; use codex_protocol::capabilities::SelectedCapabilityRoot; use codex_protocol::config_types::WindowsSandboxLevel; use codex_protocol::error::CodexErrKind; +use codex_protocol::error::CodexErrorDetails; use codex_protocol::mcp::ClientMcpExtensions; use codex_protocol::mcp::MCP_APP_UI_EXTENSION_ID; use codex_protocol::mcp::OPENAI_FORM_EXTENSION_ID; @@ -180,6 +181,94 @@ async fn startup_fixture() -> (tempfile::TempDir, Config, ThreadManager) { (home, config, manager) } +#[tokio::test] +#[expect( + clippy::await_holding_invalid_type, + reason = "hold startup before registration while the submission loop applies settings" +)] +async fn shared_service_tier_is_published_by_registered_startup() { + let (_home, mut config, manager) = startup_fixture().await; + config.ephemeral = true; + let thread_id = ThreadId::new(); + let controller = Arc::new(LocalAgentControl::default().with_session_id( + codex_protocol::SessionId::from(thread_id), + /*max_threads*/ usize::MAX, + )); + let manager = manager.with_agent_control_factory({ + let controller: Arc = controller.clone(); + move |_| { + let controller = Arc::clone(&controller); + async move { Ok(controller) } + } + }); + + for (initial_tier, updated_tier, previous_shared_tier) in [ + ("flex", None, None), + ("default", Some("flex".into()), Some("default")), + ] { + config.service_tier = Some(initial_tier.into()); + let startup = Arc::new(SessionStartup::default()); + // An ephemeral startup skips persistence until the submission loop is live. + let persistence = startup.persistence.lock().await; + let mut start = std::pin::pin!(manager.start_thread_inner( + StartThreadOptions { + reserved_thread_id: Some(thread_id), + ..StartThreadOptions::new(config.clone()) + }, + /*forked_from_thread_id*/ None, + Some(Arc::clone(&startup)), + )); + let io = std::future::poll_fn(|cx| { + assert!(start.as_mut().poll(cx).is_pending()); + startup + .io + .get() + .map_or(std::task::Poll::Pending, std::task::Poll::Ready) + }) + .await; + let (reply, applied) = tokio::sync::oneshot::channel(); + io.submit(Op::ThreadSettings { + thread_settings: codex_protocol::protocol::ThreadSettingsOverrides { + service_tier: Some(updated_tier), + ..Default::default() + }, + reply: Some(reply), + }) + .await + .expect("submit settings before registration"); + applied + .await + .expect("receive settings reply") + .expect("apply settings"); + assert_eq!( + controller.root_service_tier().as_deref(), + previous_shared_tier + ); + + drop(persistence); + let started = start.await; + if previous_shared_tier.is_some() { + let error = started.err().expect("reject the duplicate runtime"); + assert!( + matches!(error.details(), CodexErrorDetails::InvalidRequest(message) + if *message == format!("thread {thread_id} is already running")) + ); + startup.cleanup().await; + } else { + started.expect("register the first runtime"); + } + assert_eq!(controller.root_service_tier().as_deref(), Some("default")); + } + + manager + .get_thread(thread_id) + .await + .expect("find registered thread") + .shutdown_and_wait() + .await + .expect("stop registered thread"); +} + fn managed_stack(requirements: Option) -> ConfigLayerStack { let config = codex_config::ConfigRequirements { exec_policy: requirements.map(|policy| { diff --git a/codex-rs/core/tests/suite/agent_websocket.rs b/codex-rs/core/tests/suite/agent_websocket.rs index 6304ffecdd..77b513acee 100644 --- a/codex-rs/core/tests/suite/agent_websocket.rs +++ b/codex-rs/core/tests/suite/agent_websocket.rs @@ -1,28 +1,45 @@ +use super::agent_control::HostRuntime; use anyhow::Result; +use codex_core::AgentRuntimeOwnership; +use codex_core::AgentServiceTier; +use codex_core::AgentTurnAdmission; +use codex_core::InheritedSessionId; 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::config_types::ServiceTier; +use codex_protocol::dynamic_tools::DynamicToolFunctionSpec; +use codex_protocol::dynamic_tools::DynamicToolResponse; +use codex_protocol::dynamic_tools::DynamicToolSpec; +use codex_protocol::error::CodexErrorDetails; use codex_protocol::openai_models::ModelServiceTier; use codex_protocol::openai_models::ToolMode; use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::Op; +use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::ThreadSettingsOverrides; use codex_protocol::user_input::UserInput; use core_test_support::responses::WebSocketConnectionConfig; use core_test_support::responses::ev_assistant_message; use core_test_support::responses::ev_completed; use core_test_support::responses::ev_exec_command_call; +use core_test_support::responses::ev_function_call; use core_test_support::responses::ev_function_call_with_namespace; use core_test_support::responses::ev_response_created; use core_test_support::responses::start_websocket_server; use core_test_support::responses::start_websocket_server_with_headers; use core_test_support::skip_if_no_network; +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 pretty_assertions::assert_eq; use serde_json::Value; use serde_json::json; +use std::sync::Arc; +use std::sync::Mutex; use std::time::Duration; const WS_V2_BETA_HEADER_VALUE: &str = "responses_websockets=2026-02-06"; @@ -161,6 +178,186 @@ async fn spawned_agent_prewarm_handshake_inherits_effective_root_service_tier( Ok(()) } +#[derive(Default)] +struct HostTier(Mutex>); + +impl AgentServiceTier for HostTier { + fn current(&self) -> Option { + self.0.lock().expect("host tier").clone() + } + + fn update(&self, value: Option) { + *self.0.lock().expect("host tier") = value; + } +} + +#[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(())); + + let server = start_websocket_server(vec![ + vec![vec![ + ev_response_created("fixture-warm"), + ev_completed("fixture-warm"), + ]], + vec![ + vec![ + ev_response_created("child-warm"), + ev_completed("child-warm"), + ], + vec![ + ev_response_created("child-tool"), + ev_function_call("hold", "hold_turn", "{}"), + ev_completed("child-tool"), + ], + vec![ + ev_response_created("child-done"), + ev_assistant_message("msg-done", "done"), + ev_completed("child-done"), + ], + ], + ]) + .await; + let test = test_codex() + .with_model("gpt-5.4") + .with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing()) + .with_model_info_override("gpt-5.4", |model| { + model.service_tiers.push(ModelServiceTier { + id: "priority".into(), + name: "Fast".into(), + description: "Priority processing".into(), + }); + }) + .with_config(|config| { + 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 tier = Arc::new(HostTier::default()); + let mut options = test.start_thread_options().await; + 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; + options + .thread_extension_init + .insert(AgentRuntimeOwnership::Host); + options + .thread_extension_init + .insert(Arc::new(HostRuntime) as Arc); + options + .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 mut options = test.start_thread_options().await; + options.session_source = Some(SessionSource::Custom("host".into())); + options.config.service_tier = None; + options.dynamic_tools = vec![DynamicToolSpec::Function(DynamicToolFunctionSpec { + name: "hold_turn".into(), + description: "Wait for the client to release this turn.".into(), + input_schema: json!({"type": "object", "properties": {}}), + defer_loading: false, + })]; + options.thread_extension_init = thread_extension_init; + options + .thread_extension_init + .insert(InheritedSessionId(root.session_configured.session_id)); + let child = test.thread_manager.start_thread(options).await?; + + let warmup = tokio::time::timeout( + Duration::from_secs(5), + server.wait_for_request(/*connection_index*/ 1, /*request_index*/ 0), + ) + .await? + .body_json(); + assert_eq!( + server.handshakes()[1].header(X_CODEX_ROUTING_HINT_HEADER), + Some("model=gpt-5.4;tier=priority".into()) + ); + assert_eq!(warmup["generate"], false); + assert_eq!(warmup["service_tier"], "priority"); + + submit_thread_settings( + &child.thread, + ThreadSettingsOverrides { + service_tier: Some(None), + ..Default::default() + }, + ) + .await?; + assert_eq!(tier.current().as_deref(), Some("priority")); + + let request = TurnInputRequest::user_input(vec![UserInput::Text { + text: "Wait for the tier change.".into(), + text_elements: Vec::new(), + }]); + let error = child + .thread + .start_or_steer_turn(request.clone().on_start(TurnStartOptions { + service_tier: Some("default".into()), + ..Default::default() + })) + .await + .expect_err("child uses the shared tier"); + assert!( + matches!(error.details(), CodexErrorDetails::InvalidRequest(message) if message.contains("shared service tier")), + "{error:?}" + ); + + child.thread.start_or_steer_turn(request).await?; + wait_for_event( + &child.thread, + |event| matches!(event, EventMsg::DynamicToolCallRequest(call) if call.call_id == "hold"), + ) + .await; + submit_thread_settings( + &root.thread, + ThreadSettingsOverrides { + service_tier: Some(Some("flex".into())), + ..Default::default() + }, + ) + .await?; + child + .thread + .submit(Op::DynamicToolResponse { + id: "hold".into(), + response: DynamicToolResponse { + content_items: vec![], + success: true, + }, + }) + .await?; + wait_for_successful_turn(&child.thread).await; + + let connections = server.connections(); + let [_, first, second] = connections[1].as_slice() else { + panic!("expected child warmup and two inference requests"); + }; + assert_eq!( + [ + first.body_json()["service_tier"].clone(), + second.body_json()["service_tier"].clone(), + ], + [json!("priority"), json!("flex")] + ); + child.thread.shutdown_and_wait().await?; + root.thread.shutdown_and_wait().await?; + test.codex.shutdown_and_wait().await?; + server.shutdown().await; + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn websocket_model_switch_to_responses_lite_omits_top_level_tools() -> Result<()> { skip_if_no_network!(Ok(())); diff --git a/codex-rs/core/tests/suite/scenarios_reserved_spawn.rs b/codex-rs/core/tests/suite/scenarios_reserved_spawn.rs index 488312e58d..4060f09d64 100644 --- a/codex-rs/core/tests/suite/scenarios_reserved_spawn.rs +++ b/codex-rs/core/tests/suite/scenarios_reserved_spawn.rs @@ -13,6 +13,7 @@ use codex_core::SpawnAgentForkMode; use codex_core::StartIfIdleSubmission; use codex_core::ThreadCreateRequest; use codex_core::ThreadSpawnError; +use codex_core::TurnStartOptions; use codex_extension_api::ExtensionFuture; use codex_extension_api::Instructions; use codex_extension_api::ThreadInstructionsProvider; @@ -109,7 +110,7 @@ async fn host_spawn_invocation( server: &MockServer, parent: &Arc, call_id: &str, - message: &str, + input: TurnInputRequest, ) -> Result { responses::mount_sse_once( server, @@ -119,9 +120,7 @@ async fn host_spawn_invocation( ]), ) .await; - parent - .start_or_steer_turn(TurnInputRequest::user_input(vec![text(message)])) - .await?; + parent.start_or_steer_turn(input).await?; Ok(wait_for_event_match(parent, |event| match event { EventMsg::DynamicToolCallRequest(call) => Some(call.clone()), _ => None, @@ -190,7 +189,14 @@ async fn same_manager_creation_preserves_live_bindings_and_attachments( )); Arc::new(extensions.build()) }) - .with_config(|config| config.allow_persistent_dynamic_tool_updates = true), + .with_config(|config| { + config.allow_persistent_dynamic_tool_updates = true; + config.service_tier = Some("flex".into()); + config + .features + .enable(Feature::FastMode) + .expect("test config should allow fast mode"); + }), network_policy.clone(), Some(provider.clone()), ) @@ -198,8 +204,16 @@ async fn same_manager_creation_preserves_live_bindings_and_attachments( let root_id = test.session_configured.thread_id; let captured = parent.config_snapshot().await; let parent_message = "Start the worker."; - let invocation = - host_spawn_invocation(&server, &parent, "spawn-worker", parent_message).await?; + let invocation = host_spawn_invocation( + &server, + &parent, + "spawn-worker", + TurnInputRequest::user_input(vec![text(parent_message)]).on_start(TurnStartOptions { + service_tier: Some("priority".into()), + ..Default::default() + }), + ) + .await?; let context = parent.tool_invocation(&invocation.turn_id, &invocation.call_id)?; test.thread_store .add_thread_attachment(AddThreadAttachmentParams { @@ -253,6 +267,7 @@ async fn same_manager_creation_preserves_live_bindings_and_attachments( let spawned = parent .create_child(&context, request(), reservation()) .await?; + assert_eq!(spawned.config.service_tier.as_deref(), Some("flex")); assert_eq!(spawned.config.approval_policy, captured.approval_policy); assert_eq!(spawned.config.environments, captured.environments); assert_eq!( @@ -299,6 +314,7 @@ async fn same_manager_creation_preserves_live_bindings_and_attachments( && report.usage.total_tokens == 42 )); let child_request = child_request.single_request(); + assert_eq!(child_request.body_json()["service_tier"], "flex"); assert!(child_request.body_contains_text(UPDATED_RULE)); assert!(!child_request.body_contains_text("only as direct tool calls")); assert_eq!( @@ -362,8 +378,13 @@ async fn cancelled_creation_cleans_registered_child() -> Result<()> { /*thread_instructions_provider*/ None, ) .await?; - let invocation = - host_spawn_invocation(&server, &parent, "cancel-spawn", "Start a worker.").await?; + let invocation = host_spawn_invocation( + &server, + &parent, + "cancel-spawn", + TurnInputRequest::user_input(vec![text("Start a worker.")]), + ) + .await?; let context = parent.tool_invocation(&invocation.turn_id, &invocation.call_id)?; let spawn = tokio::spawn({ let parent = Arc::clone(&parent); diff --git a/codex-rs/core/tests/suite/subagent_service_tier.rs b/codex-rs/core/tests/suite/subagent_service_tier.rs index 74d5592540..3f3aaaec1b 100644 --- a/codex-rs/core/tests/suite/subagent_service_tier.rs +++ b/codex-rs/core/tests/suite/subagent_service_tier.rs @@ -1,9 +1,12 @@ //! Verifies that root service-tier changes reach existing child turns and future work. +use super::agent_control::HostRuntime; use anyhow::Result; +use codex_core::AgentRuntimeOwnership; use codex_core::TurnInputRequest; use codex_core::config::AgentRoleConfig; use codex_core::config::Config; +use codex_extension_api::ExtensionRegistryBuilder; use codex_features::Feature; use codex_protocol::config_types::SERVICE_TIER_DEFAULT_REQUEST_VALUE; use codex_protocol::protocol::EventMsg; @@ -24,6 +27,7 @@ use core_test_support::wait_for_event; use core_test_support::wait_for_event_match; use pretty_assertions::assert_eq; use serde_json::json; +use std::sync::Arc; use test_case::test_case; const ROOT_PROMPT: &str = "spawn the service-tier worker"; @@ -304,13 +308,26 @@ async fn root_service_tier_change_updates_existing_subagent( Ok(()) } +#[test_case(AgentRuntimeOwnership::Controller, None; "controller")] +#[test_case(AgentRuntimeOwnership::Host, Some("priority"); "host without a tier provider")] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn evicted_role_subagent_uses_root_service_tier_after_reload() -> Result<()> { +async fn evicted_role_subagent_uses_owned_service_tier_after_reload( + owner: AgentRuntimeOwnership, + expected: Option<&str>, +) -> Result<()> { let server = start_mock_server().await; + let mut extensions = ExtensionRegistryBuilder::new(); + if owner == AgentRuntimeOwnership::Host { + extensions.thread_lifecycle_contributor(Arc::new(HostRuntime)); + } let mut builder = test_codex() .with_model("gpt-5.6-sol") - .with_config(|config| { - config.service_tier = Some("priority".to_string()); + .with_extensions(Arc::new(extensions.build())) + .with_config(move |config| { + config.service_tier = Some(match owner { + AgentRuntimeOwnership::Controller => "priority".to_string(), + AgentRuntimeOwnership::Host => SERVICE_TIER_DEFAULT_REQUEST_VALUE.to_string(), + }); config.multi_agent_v2.max_concurrent_threads_per_session = 2; configure_priority_role(config); }); @@ -375,10 +392,13 @@ async fn evicted_role_subagent_uses_root_service_tier_after_reload() -> Result<( .ensure_multi_agent_v2_child_loaded(original_thread_id) .await?; let reloaded_thread = test.thread_manager.get_thread(original_thread_id).await?; + let expected_configured_tier = match owner { + AgentRuntimeOwnership::Controller => test.codex.config_snapshot().await.service_tier, + AgentRuntimeOwnership::Host => expected.map(str::to_string), + }; assert_eq!( reloaded_thread.config_snapshot().await.service_tier, - test.codex.config_snapshot().await.service_tier, - "reload ignores the role tier and preserves the root-owned preference" + expected_configured_tier, ); let reloaded_request = mount_completed_child(&server, FOLLOWUP_PROMPT, ROOT_PROMPT).await; @@ -389,7 +409,7 @@ async fn evicted_role_subagent_uses_root_service_tier_after_reload() -> Result<( }])) .await?; wait_for_turn_complete(&reloaded_thread).await; - assert_request_service_tier(&reloaded_request, /*expected*/ None); + assert_request_service_tier(&reloaded_request, expected); reloaded_thread.shutdown_and_wait().await?; test.codex.shutdown_and_wait().await?;