openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-11-06 20:20:34 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add shared service-tier support for host-owned agents (#59770)

Browse Source 
## What changed

- Expose an optional `AgentServiceTier` provider so host-owned children, including detached children, inherit the root's committed tier during spawn preparation, startup prewarm, and each inference step.
- Reject per-turn tier overrides for children using the shared provider. Root updates apply to subsequent requests without changing requests already in flight.
- Preserve independent configured tiers and role overrides when no provider is installed, including across eviction and reload. Spawned children inherit the parent's thread preference rather than its per-turn override.
- Publish the latest committed root tier only after successful runtime registration, preventing failed duplicate starts from changing shared settings.

## Testing

Add coverage for detached-child WebSocket prewarm and between-step tier changes, shared-child override rejection, independent-tier inheritance and reload, and settings updates before registration.

GitOrigin-RevId: f0f3676813a1c59bd77baf95121e2d13e0163a8e
This commit is contained in:
rhan-oai
2026-11-06 00:24:27 +00:00
committed by copyberry
parent 74417e9c6b
commit 94ca55a93c
18 changed files with 517 additions and 47 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/agent/child_config.rs
+14 -2
View File
@@ -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(());
};
codex-rs/core/src/agent/control/residency.rs
+5
View File
@@ -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
codex-rs/core/src/agent/control/spawn.rs
+5 -1
View File
@@ -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);
}
codex-rs/core/src/agent/host/mod.rs
+5
View File
@@ -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;
codex-rs/core/src/agent/host/ownership.rs
+4
View File
@@ -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<Arc<dyn AgentSessionReader>>,
pub(crate) turn_admission: Arc<dyn AgentTurnAdmission>,
pub(crate) usage_accounting: Option<Arc<dyn AgentUsageAccounting>>,
pub(crate) service_tier: Option<Arc<dyn AgentServiceTier>>,
}
impl AgentRuntimeServices {
@@ -51,6 +54,7 @@ impl AgentRuntimeServices {
.get::<Arc<dyn AgentUsageAccounting>>()
.as_deref()
.cloned(),
service_tier: data.get::<Arc<dyn AgentServiceTier>>().as_deref().cloned(),
}))
}
}
codex-rs/core/src/agent/host/service_tier.rs
+59
View File
@@ -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<dyn AgentServiceTier>` 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<String>;
fn update(&self, service_tier: Option<String>);
}
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<String>,
) -> Option<String> {
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);
}
}
}
}
codex-rs/core/src/agent/registry.rs
+1
View File
@@ -43,6 +43,7 @@ struct ActiveAgents {
#[derive(Clone)]
pub(crate) struct EvictedEnvironments {
pub(crate) config: Arc<Config>,
pub(crate) has_independent_service_tier: bool,
pub(crate) selections: Vec<TurnEnvironmentSelection>,
pub(crate) bindings: TurnEnvironmentSnapshot,
}
codex-rs/core/src/lib.rs
+1
View File
@@ -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;
codex-rs/core/src/session/mod.rs
+38 -7
View File
@@ -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;
codex-rs/core/src/session/session.rs
+9
View File
@@ -1092,6 +1092,8 @@ impl Session {
}
}
}
let has_provided_agent_control =
matches!(&agent_control, AgentControlInit::Provided { .. });
let (agent_control, local_agent_runtime): (Arc<dyn AgentControl>, _) = 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 =
codex-rs/core/src/session/startup_prewarm.rs
+7 -7
View File
@@ -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,
)
codex-rs/core/src/session/turn_input.rs
+8
View File
@@ -245,6 +245,14 @@ pub(super) async fn handle(
submission_id: String,
) -> CodexResult<TurnInputSubmission> {
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
codex-rs/core/src/state/session.rs
+3
View File
@@ -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(),
codex-rs/core/src/thread_manager.rs
+16 -15
View File
@@ -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<Session>,
io: SessionIo,
session_source: SessionSource,
initial_host_config: Option<AgentConfigUpdate>,
) -> CodexResult<NewThread> {
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,
codex-rs/core/src/thread_manager_tests.rs
+89
View File
@@ -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<dyn AgentControl> = 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<RequirementsExecPolicy>) -> ConfigLayerStack {
let config = codex_config::ConfigRequirements {
exec_policy: requirements.map(|policy| {
codex-rs/core/tests/suite/agent_websocket.rs
+197
View File
@@ -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<Option<String>>);
impl AgentServiceTier for HostTier {
fn current(&self) -> Option<String> {
self.0.lock().expect("host tier").clone()
}
fn update(&self, value: Option<String>) {
*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<dyn AgentTurnAdmission>);
options
.thread_extension_init
.insert(tier.clone() as Arc<dyn AgentServiceTier>);
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(()));
codex-rs/core/tests/suite/scenarios_reserved_spawn.rs
+30 -9
View File
@@ -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<CodexThread>,
call_id: &str,
message: &str,
input: TurnInputRequest,
) -> Result<DynamicToolCallRequest> {
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);
codex-rs/core/tests/suite/subagent_service_tier.rs
+26 -6
View File
@@ -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?;