Keep interrupts responsive during host service-tier updates (#60644)

## Why

Waiting for host acceptance of a service-tier update in the submission loop can block interrupts while the host receipt is pending.

## What changed

Move the host acceptance wait out of the submission loop while keeping explicit thread settings replies pending until acceptance completes. Report handoff failures to the caller while preserving the locally committed settings, and allow a subsequent explicit request for the same tier to retry a failed handoff.

## Testing

Add a WebSocket regression test covering interrupts with pending host receipts, failure reporting, same-tier retries, and continued thread operation after startup caller cancellation.

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