mirror of
https://github.com/openai/codex.git
synced 2026-11-08 20:40:33 +00:00
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
This commit is contained in:
@@ -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(());
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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<BoxFuture<'static, ()>>;
|
||||
|
||||
/// The host's 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;
|
||||
/// 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<String>;
|
||||
fn update(&self, service_tier: Option<String>);
|
||||
fn current(&self) -> BoxFuture<'_, Option<String>>;
|
||||
|
||||
/// 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<String>) -> 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<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()),
|
||||
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<ServiceTierAcceptance> {
|
||||
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())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -2160,6 +2160,7 @@ impl Session {
|
||||
) -> ConstraintResult<Option<SessionSettingsCommit>> {
|
||||
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;
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
@@ -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<ServiceTierAcceptance>,
|
||||
/// 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(),
|
||||
|
||||
@@ -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<Session>,
|
||||
io: SessionIo,
|
||||
session_source: SessionSource,
|
||||
) -> CodexResult<NewThread> {
|
||||
) -> CodexResult<(NewThread, Option<ServiceTierAcceptance>)> {
|
||||
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));
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user