core: make local idle shutdown atomic

This commit is contained in:
Winston Howes
2026-07-06 17:42:27 -07:00
parent 7226904ed1
commit bc530e25f6
10 changed files with 480 additions and 17 deletions
+4 -9
View File
@@ -49,6 +49,7 @@ use crate::session::Codex;
use crate::session::CodexSpawnArgs;
use crate::session::CodexSpawnOk;
use crate::session::SUBMISSION_CHANNEL_CAPACITY;
use crate::session::SubmissionTransport;
use crate::session::emit_subagent_session_started;
use crate::session::session::Session;
use crate::session::turn_context::TurnContext;
@@ -186,6 +187,7 @@ pub(crate) async fn run_codex_thread_interactive(
agent_status: codex.agent_status.clone(),
session: Arc::clone(&codex.session),
session_loop_termination: codex.session_loop_termination.clone(),
submission_transport: SubmissionTransport::ForwardingProxy,
})
}
@@ -232,7 +234,6 @@ pub(crate) async fn run_codex_thread_one_shot(
// Bridge events so we can observe completion and shut down automatically.
let (tx_bridge, rx_bridge) = async_channel::bounded(SUBMISSION_CHANNEL_CAPACITY);
let ops_tx = io.tx_sub.clone();
let agent_status = io.agent_status.clone();
let session = Arc::clone(&io.session);
let session_loop_termination = io.session_loop_termination.clone();
@@ -245,14 +246,7 @@ pub(crate) async fn run_codex_thread_one_shot(
);
let _ = tx_bridge.send(event).await;
if should_shutdown {
let _ = ops_tx
.send(Submission {
id: "shutdown".to_string(),
op: Op::Shutdown {},
client_user_message_id: None,
trace: None,
})
.await;
let _ = io_for_bridge.submit(Op::Shutdown {}).await;
child_cancel.cancel();
break;
}
@@ -271,6 +265,7 @@ pub(crate) async fn run_codex_thread_one_shot(
agent_status,
session,
session_loop_termination,
submission_transport: SubmissionTransport::ForwardingProxy,
})
}
+33 -4
View File
@@ -48,6 +48,7 @@ async fn forward_events_filters_private_events_before_blocked_send_is_cancelled(
agent_status,
session: Arc::clone(&session),
session_loop_termination: completed_session_loop_termination(),
submission_transport: SubmissionTransport::Canonical,
});
let (tx_out, rx_out) = bounded(1);
@@ -136,7 +137,7 @@ async fn forward_events_filters_private_events_before_blocked_send_is_cancelled(
}
#[tokio::test]
async fn forward_ops_preserves_submission_trace_context() {
async fn forwarding_proxy_reserves_once_and_delivers_shutdown() {
let (tx_sub, rx_sub) = bounded(SUBMISSION_CHANNEL_CAPACITY);
let (_tx_events, rx_events) = bounded(SUBMISSION_CHANNEL_CAPACITY);
let (_agent_status_tx, agent_status) = watch::channel(AgentStatus::PendingInit);
@@ -145,10 +146,20 @@ async fn forward_ops_preserves_submission_trace_context() {
tx_sub,
rx_event: rx_events,
agent_status,
session,
session: Arc::clone(&session),
session_loop_termination: completed_session_loop_termination(),
submission_transport: SubmissionTransport::Canonical,
});
let (tx_ops, rx_ops) = bounded(1);
let (_tx_proxy_events, rx_proxy_events) = bounded(SUBMISSION_CHANNEL_CAPACITY);
let proxy = Codex {
tx_sub: tx_ops,
rx_event: rx_proxy_events,
agent_status: codex.agent_status.clone(),
session: Arc::clone(&session),
session_loop_termination: completed_session_loop_termination(),
submission_transport: SubmissionTransport::ForwardingProxy,
};
let cancel = CancellationToken::new();
let forward = tokio::spawn(forward_ops(Arc::clone(&codex), rx_ops, cancel));
@@ -163,8 +174,7 @@ async fn forward_ops_preserves_submission_trace_context() {
tracestate: Some("vendor=state".to_string()),
}),
};
tx_ops.send(submission.clone()).await.unwrap();
drop(tx_ops);
proxy.submit_with_id(submission.clone()).await.unwrap();
let forwarded = timeout(Duration::from_secs(1), rx_sub.recv())
.await
@@ -173,6 +183,23 @@ async fn forward_ops_preserves_submission_trace_context() {
assert_eq!(submission.id, forwarded.id);
assert_eq!(submission.op, forwarded.op);
assert_eq!(submission.trace, forwarded.trace);
session.finish_submission();
let idle_reservation = session
.try_claim_idle_shutdown()
.await
.expect("the forwarded submission should own exactly one reservation");
drop(idle_reservation);
proxy.submit(Op::Shutdown).await.unwrap();
let forwarded_shutdown = timeout(Duration::from_secs(1), rx_sub.recv())
.await
.expect("forward_ops hung while delivering shutdown")
.expect("forwarded shutdown missing");
assert!(matches!(forwarded_shutdown.op, Op::Shutdown));
session.finish_submission();
drop(proxy);
timeout(Duration::from_secs(1), forward)
.await
@@ -223,6 +250,7 @@ async fn handle_request_permissions_uses_tool_call_id_for_round_trip() {
agent_status,
session: Arc::clone(&parent_session),
session_loop_termination: completed_session_loop_termination(),
submission_transport: SubmissionTransport::Canonical,
});
let call_id = "tool-call-1".to_string();
@@ -328,6 +356,7 @@ async fn handle_exec_approval_uses_call_id_for_guardian_review_and_approval_id_f
agent_status,
session: Arc::clone(&parent_session),
session_loop_termination: completed_session_loop_termination(),
submission_transport: SubmissionTransport::Canonical,
});
let cancel_token = CancellationToken::new();
+32
View File
@@ -211,6 +211,10 @@ impl CodexThread {
self.codex.shutdown_and_wait().await
}
pub async fn try_shutdown_if_idle(&self) -> CodexResult<bool> {
self.codex.try_shutdown_if_idle().await
}
/// Wait until the underlying session loop has terminated.
pub async fn wait_until_terminated(&self) {
self.codex.session_loop_termination.clone().await;
@@ -452,6 +456,9 @@ impl CodexThread {
/// Records a user-role session-prefix message without creating a new user turn boundary.
pub(crate) async fn inject_user_message_without_turn(&self, message: String) {
let Some(_reservation) = self.codex.session.try_reserve_activity() else {
return;
};
let item = ResponseItem::Message {
id: None,
role: "user".to_string(),
@@ -472,6 +479,11 @@ impl CodexThread {
"items must not be empty".to_string(),
));
}
let Some(_reservation) = self.codex.session.try_reserve_activity() else {
return Err(CodexErr::InvalidRequest(
"thread is shutting down".to_string(),
));
};
let turn_context = self.codex.session.new_default_turn().await;
if self.codex.session.reference_context_item().await.is_none() {
@@ -634,6 +646,11 @@ impl CodexThread {
server: &str,
uri: &str,
) -> anyhow::Result<serde_json::Value> {
let _reservation = self
.codex
.session
.try_reserve_activity()
.ok_or_else(|| anyhow::anyhow!("thread is shutting down"))?;
let result = self
.current_mcp_runtime()
.await
@@ -651,6 +668,11 @@ impl CodexThread {
arguments: Option<serde_json::Value>,
meta: Option<serde_json::Value>,
) -> anyhow::Result<CallToolResult> {
let _reservation = self
.codex
.session
.try_reserve_activity()
.ok_or_else(|| anyhow::anyhow!("thread is shutting down"))?;
self.current_mcp_runtime()
.await
.manager_arc()
@@ -663,6 +685,11 @@ impl CodexThread {
}
pub async fn increment_out_of_band_elicitation_count(&self) -> CodexResult<i64> {
let Some(_reservation) = self.codex.session.try_reserve_activity() else {
return Err(CodexErr::InvalidRequest(
"thread is shutting down".to_string(),
));
};
let mut elicitations = self.out_of_band_elicitations.lock().await;
let incremented = elicitations.count.checked_add(1).ok_or_else(|| {
CodexErr::Fatal("out-of-band elicitation count overflowed".to_string())
@@ -675,6 +702,11 @@ impl CodexThread {
}
pub async fn decrement_out_of_band_elicitation_count(&self) -> CodexResult<i64> {
let Some(_reservation) = self.codex.session.try_reserve_activity() else {
return Err(CodexErr::InvalidRequest(
"thread is shutting down".to_string(),
));
};
let mut elicitations = self.out_of_band_elicitations.lock().await;
if elicitations.count == 0 {
return Err(CodexErr::InvalidRequest(
@@ -1154,6 +1154,7 @@ mod tests {
agent_status,
session,
session_loop_termination: crate::session::completed_session_loop_termination(),
submission_transport: crate::session::SubmissionTransport::Canonical,
},
cancel_token: CancellationToken::new(),
reuse_key,
+1
View File
@@ -841,6 +841,7 @@ pub(super) async fn submission_loop(
}
.instrument(dispatch_span)
.await;
sess.finish_submission();
if should_exit {
shutdown_received = true;
break;
+7
View File
@@ -62,6 +62,12 @@ impl Session {
));
}
let Some(reservation) = self.try_reserve_activity() else {
return Err(TryStartTurnIfIdleError::new(
TryStartTurnIfIdleRejectionReason::Busy,
input,
));
};
let turn_state = {
let mut active_turn = self.active_turn.lock().await;
if active_turn.is_some() {
@@ -73,6 +79,7 @@ impl Session {
let active_turn = active_turn.get_or_insert_with(ActiveTurn::default);
Arc::clone(&active_turn.turn_state)
};
drop(reservation);
if self.input_queue.has_trigger_turn_mailbox_items().await {
self.clear_reserved_idle_turn(&turn_state).await;
+93 -4
View File
@@ -393,10 +393,17 @@ pub struct Codex {
// Shared future for the background submission loop completion so multiple
// callers can wait for shutdown.
pub(crate) session_loop_termination: SessionLoopTermination,
pub(crate) submission_transport: SubmissionTransport,
}
pub(crate) type SessionLoopTermination = Shared<BoxFuture<'static, ()>>;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) enum SubmissionTransport {
Canonical,
ForwardingProxy,
}
/// Wrapper returned by [`Codex::spawn`] containing the spawned [`Codex`] and
/// the unique session id.
pub struct CodexSpawnOk {
@@ -721,6 +728,7 @@ impl Codex {
agent_status: agent_status_rx,
session,
session_loop_termination: session_loop_termination_from_handle(session_loop_handle),
submission_transport: SubmissionTransport::Canonical,
};
Ok(CodexSpawnOk { codex, thread_id })
@@ -767,17 +775,73 @@ impl Codex {
/// Use sparingly: prefer `submit()` so Codex is responsible for generating
/// unique IDs for each submission.
#[expect(
clippy::await_holding_invalid_type,
reason = "submission ordering must remain serialized while a bounded-channel send waits"
)]
pub async fn submit_with_id(&self, mut sub: Submission) -> CodexResult<()> {
if sub.trace.is_none() {
sub.trace = current_span_w3c_trace_context();
}
self.tx_sub
.send(sub)
.await
.map_err(|_| CodexErr::InternalAgentDied)?;
if self.submission_transport == SubmissionTransport::ForwardingProxy {
self.tx_sub
.send(sub)
.await
.map_err(|_| CodexErr::InternalAgentDied)?;
return Ok(());
}
let is_shutdown = matches!(&sub.op, Op::Shutdown);
let _send_guard = self.session.submission_send_lock.lock().await;
let Some(reservation) = self.session.submission_lifecycle.try_reserve(is_shutdown) else {
if is_shutdown {
return Ok(());
}
return Err(CodexErr::InvalidRequest(
"thread is shutting down".to_string(),
));
};
if self.tx_sub.send(sub).await.is_err() {
reservation.release_after_failed_delivery();
return Err(CodexErr::InternalAgentDied);
}
reservation.commit();
Ok(())
}
/// Atomically submits shutdown only when no work is active, queued, or
/// being dispatched. Once claimed, new submissions and automatic idle work
/// are rejected so shutdown cannot race a newly starting turn.
#[expect(
clippy::await_holding_invalid_type,
reason = "the send lock serializes the idle claim with submission delivery"
)]
pub async fn try_shutdown_if_idle(&self) -> CodexResult<bool> {
if self.submission_transport == SubmissionTransport::ForwardingProxy {
return Err(CodexErr::InvalidRequest(
"idle shutdown is unavailable through a forwarding proxy".to_string(),
));
}
let _send_guard = self.session.submission_send_lock.lock().await;
if self.session.submission_lifecycle.is_closing() {
return Ok(true);
}
let Some(reservation) = self.session.try_claim_idle_shutdown().await else {
return Ok(false);
};
let sub = Submission {
id: new_submission_id(),
op: Op::Shutdown,
client_user_message_id: None,
trace: current_span_w3c_trace_context(),
};
if self.tx_sub.send(sub).await.is_err() {
reservation.release_after_failed_delivery();
return Err(CodexErr::InternalAgentDied);
}
reservation.commit();
Ok(true)
}
/// Persist a thread-level memory mode update for the active session.
///
/// This is a local-only operation that updates rollout metadata directly
@@ -978,6 +1042,28 @@ fn push_prompt_fragment(
}
impl Session {
async fn has_idle_shutdown_blocker(&self) -> bool {
matches!(self.agent_status.borrow().clone(), AgentStatus::Running)
|| self.conversation.running_state().await.is_some()
|| *self.services.elicitations.subscribe().borrow()
|| self.input_queue.has_pending_mailbox_items().await
|| self.active_turn.lock().await.is_some()
|| !self.list_background_terminals().await.is_empty()
}
pub(crate) async fn try_claim_idle_shutdown(
&self,
) -> Option<session::SubmissionReservation<'_>> {
if self.has_idle_shutdown_blocker().await {
return None;
}
let reservation = self.submission_lifecycle.try_reserve_idle_shutdown()?;
if self.has_idle_shutdown_blocker().await {
return None;
}
Some(reservation)
}
pub(crate) async fn app_server_client_metadata(&self) -> AppServerClientMetadata {
let state = self.state.lock().await;
AppServerClientMetadata {
@@ -1194,6 +1280,9 @@ impl Session {
}
pub(crate) async fn route_realtime_text_input(self: &Arc<Self>, text: String) {
let Some(_reservation) = self.try_reserve_activity() else {
return;
};
handlers::user_input_or_turn_inner(
self,
Uuid::now_v7().to_string(),
+115
View File
@@ -20,11 +20,114 @@ use codex_protocol::protocol::ThreadHistoryMode;
use codex_protocol::protocol::ThreadSource;
use codex_protocol::protocol::TurnEnvironmentSelections;
use std::sync::OnceLock;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering;
use tokio::sync::Semaphore;
/// Context for an initialized model agent
///
/// A session has at most 1 running task at a time, and can be interrupted by user input.
pub(crate) struct SubmissionLifecycle {
state: AtomicUsize,
}
const SUBMISSION_CLOSING: usize = 1 << (usize::BITS - 1);
const SUBMISSION_COUNT_MASK: usize = SUBMISSION_CLOSING - 1;
impl Default for SubmissionLifecycle {
fn default() -> Self {
Self {
state: AtomicUsize::new(0),
}
}
}
impl SubmissionLifecycle {
pub(crate) fn is_closing(&self) -> bool {
self.state.load(Ordering::SeqCst) & SUBMISSION_CLOSING != 0
}
pub(crate) fn try_reserve(&self, close: bool) -> Option<SubmissionReservation<'_>> {
self.state
.fetch_update(Ordering::SeqCst, Ordering::SeqCst, |state| {
if state & SUBMISSION_CLOSING != 0 {
return None;
}
let count = state & SUBMISSION_COUNT_MASK;
assert!(count < SUBMISSION_COUNT_MASK, "submission count overflow");
Some((state + 1) | if close { SUBMISSION_CLOSING } else { 0 })
})
.ok()?;
Some(SubmissionReservation {
lifecycle: self,
clear_closing_on_drop: close,
active: true,
})
}
pub(crate) fn try_reserve_idle_shutdown(&self) -> Option<SubmissionReservation<'_>> {
self.state
.compare_exchange(
0,
SUBMISSION_CLOSING | 1,
Ordering::SeqCst,
Ordering::SeqCst,
)
.ok()
.map(|_| SubmissionReservation {
lifecycle: self,
clear_closing_on_drop: true,
active: true,
})
}
pub(crate) fn finish_submission(&self) {
self.release(/*clear_closing*/ false);
}
fn release(&self, clear_closing: bool) {
let _ = self
.state
.fetch_update(Ordering::SeqCst, Ordering::SeqCst, |state| {
if state & SUBMISSION_COUNT_MASK == 0 {
return None;
}
let mut next = state - 1;
if clear_closing {
next &= !SUBMISSION_CLOSING;
}
Some(next)
});
}
}
pub(crate) struct SubmissionReservation<'a> {
lifecycle: &'a SubmissionLifecycle,
clear_closing_on_drop: bool,
active: bool,
}
impl SubmissionReservation<'_> {
pub(crate) fn commit(mut self) {
self.active = false;
}
pub(crate) fn release_after_failed_delivery(mut self) {
if self.active {
self.lifecycle.release(/*clear_closing*/ false);
self.active = false;
}
}
}
impl Drop for SubmissionReservation<'_> {
fn drop(&mut self) {
if self.active {
self.lifecycle.release(self.clear_closing_on_drop);
}
}
}
pub(crate) struct Session {
pub(crate) thread_id: ThreadId,
pub(crate) installation_id: String,
@@ -41,6 +144,8 @@ pub(crate) struct Session {
pub(super) pending_mcp_server_refresh_config: Mutex<Option<McpServerRefreshConfig>>,
pub(crate) conversation: Arc<RealtimeConversationManager>,
pub(crate) active_turn: Mutex<Option<ActiveTurn>>,
pub(crate) submission_lifecycle: SubmissionLifecycle,
pub(crate) submission_send_lock: Mutex<()>,
pub(crate) input_queue: InputQueue,
pub(crate) guardian_review_session: GuardianReviewSessionManager,
pub(crate) services: SessionServices,
@@ -465,6 +570,14 @@ impl Session {
self.thread_id
}
pub(crate) fn try_reserve_activity(&self) -> Option<SubmissionReservation<'_>> {
self.submission_lifecycle.try_reserve(/*close*/ false)
}
pub(crate) fn finish_submission(&self) {
self.submission_lifecycle.finish_submission();
}
/// Returns the identity shared by the root thread and all descendant threads.
pub(crate) fn session_id(&self) -> SessionId {
self.services.agent_control.session_id()
@@ -1137,6 +1250,8 @@ impl Session {
pending_mcp_server_refresh_config: Mutex::new(None),
conversation: Arc::new(RealtimeConversationManager::new()),
active_turn: Mutex::new(None),
submission_lifecycle: SubmissionLifecycle::default(),
submission_send_lock: Mutex::new(()),
input_queue: InputQueue::new(),
guardian_review_session: GuardianReviewSessionManager::default(),
services,
+190
View File
@@ -1,3 +1,4 @@
use super::session::SubmissionLifecycle;
use super::turn_context::TurnEnvironment;
use super::*;
use crate::agents_md_manager::AgentsMdManager;
@@ -5518,6 +5519,8 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) {
pending_mcp_server_refresh_config: Mutex::new(None),
conversation: Arc::new(RealtimeConversationManager::new()),
active_turn: Mutex::new(None),
submission_lifecycle: SubmissionLifecycle::default(),
submission_send_lock: Mutex::new(()),
input_queue: super::input_queue::InputQueue::new(),
guardian_review_session: crate::guardian::GuardianReviewSessionManager::default(),
services,
@@ -6448,6 +6451,7 @@ async fn submit_with_id_captures_current_span_trace_context() {
agent_status,
session: Arc::new(session),
session_loop_termination: completed_session_loop_termination(),
submission_transport: SubmissionTransport::Canonical,
};
let _trace_test_context = install_test_tracing("codex-core-tests");
@@ -6483,6 +6487,146 @@ async fn submit_with_id_captures_current_span_trace_context() {
assert_eq!(submitted.trace, Some(expected_trace));
}
async fn codex_with_disconnected_submission_channel() -> (Codex, Arc<Session>) {
let (session, _turn_context) = make_session_and_context().await;
let session = Arc::new(session);
let (tx_sub, rx_sub) = async_channel::bounded(1);
drop(rx_sub);
let (_tx_event, rx_event) = async_channel::unbounded();
let (_agent_status_tx, agent_status) = watch::channel(AgentStatus::PendingInit);
let codex = Codex {
tx_sub,
rx_event,
agent_status,
session: Arc::clone(&session),
session_loop_termination: completed_session_loop_termination(),
submission_transport: SubmissionTransport::Canonical,
};
(codex, session)
}
fn assert_failed_shutdown_stays_closed(session: &Session) {
assert!(session.submission_lifecycle.is_closing());
}
#[tokio::test]
async fn failed_idle_shutdown_delivery_stays_closed() {
let (codex, session) = codex_with_disconnected_submission_channel().await;
assert!(matches!(
codex.try_shutdown_if_idle().await,
Err(CodexErr::InternalAgentDied)
));
assert_failed_shutdown_stays_closed(&session);
}
#[tokio::test]
async fn failed_explicit_shutdown_delivery_stays_closed() {
let (codex, session) = codex_with_disconnected_submission_channel().await;
assert!(matches!(
codex
.submit_with_id(Submission {
id: "shutdown".to_string(),
op: Op::Shutdown,
client_user_message_id: None,
trace: None,
})
.await,
Err(CodexErr::InternalAgentDied)
));
assert_failed_shutdown_stays_closed(&session);
}
#[tokio::test]
async fn cancelled_submission_and_idle_shutdown_release_lifecycle_reservations() {
let (session, _turn_context) = make_session_and_context().await;
let session = Arc::new(session);
let (tx_sub, rx_sub) = async_channel::bounded(1);
let (_tx_event, rx_event) = async_channel::unbounded();
let (_agent_status_tx, agent_status) = watch::channel(AgentStatus::PendingInit);
tx_sub
.send(Submission {
id: "filler".into(),
op: Op::Interrupt,
client_user_message_id: None,
trace: None,
})
.await
.expect("fill submission channel");
let codex = Arc::new(Codex {
tx_sub,
rx_event,
agent_status,
session: Arc::clone(&session),
session_loop_termination: completed_session_loop_termination(),
submission_transport: SubmissionTransport::Canonical,
});
let blocked_submit = {
let codex = Arc::clone(&codex);
tokio::spawn(async move { codex.submit(Op::Interrupt).await })
};
tokio::time::timeout(StdDuration::from_secs(1), async {
while session.submission_send_lock.try_lock().is_ok() {
tokio::task::yield_now().await;
}
})
.await
.expect("submission should hold the send lock while blocking");
blocked_submit.abort();
assert!(
blocked_submit
.await
.expect_err("blocked submission should be cancelled")
.is_cancelled()
);
let released_submission = session
.try_claim_idle_shutdown()
.await
.expect("cancelled submission should release its lifecycle reservation");
drop(released_submission);
assert!(!session.submission_lifecycle.is_closing());
let blocked_shutdown = {
let codex = Arc::clone(&codex);
tokio::spawn(async move { codex.try_shutdown_if_idle().await })
};
tokio::time::timeout(StdDuration::from_secs(1), async {
while !session.submission_lifecycle.is_closing() {
tokio::task::yield_now().await;
}
})
.await
.expect("idle shutdown should claim lifecycle state before blocking");
blocked_shutdown.abort();
assert!(
blocked_shutdown
.await
.expect_err("blocked idle shutdown should be cancelled")
.is_cancelled()
);
let released_shutdown = session
.try_claim_idle_shutdown()
.await
.expect("cancelled idle shutdown should release its lifecycle reservation");
drop(released_shutdown);
assert!(!session.submission_lifecycle.is_closing());
let filler = rx_sub.recv().await.expect("filler submission");
assert_eq!(filler.id, "filler");
assert!(
codex
.try_shutdown_if_idle()
.await
.expect("idle shutdown should submit after cancellation")
);
let shutdown = rx_sub.recv().await.expect("shutdown submission");
assert_eq!(shutdown.op, Op::Shutdown);
}
#[tokio::test]
async fn new_default_turn_captures_current_span_trace_id() {
let (session, _turn_context) = make_session_and_context().await;
@@ -7180,6 +7324,7 @@ async fn shutdown_and_wait_allows_multiple_waiters() {
agent_status,
session: Arc::new(session),
session_loop_termination: session_loop_termination_from_handle(session_loop_handle),
submission_transport: SubmissionTransport::Canonical,
});
let waiter_1 = {
@@ -7218,6 +7363,7 @@ async fn shutdown_and_wait_waits_when_shutdown_is_already_in_progress() {
agent_status,
session: Arc::new(session),
session_loop_termination: session_loop_termination_from_handle(session_loop_handle),
submission_transport: SubmissionTransport::Canonical,
});
let waiter = {
@@ -7256,6 +7402,7 @@ async fn shutdown_and_wait_shuts_down_cached_guardian_subagent() {
agent_status: parent_agent_status,
session: Arc::clone(&parent_session),
session_loop_termination: session_loop_termination_from_handle(parent_session_loop_handle),
submission_transport: SubmissionTransport::Canonical,
};
let (child_session, _child_turn_context) = make_session_and_context().await;
@@ -7279,6 +7426,7 @@ async fn shutdown_and_wait_shuts_down_cached_guardian_subagent() {
agent_status: child_agent_status,
session: Arc::new(child_session),
session_loop_termination: session_loop_termination_from_handle(child_session_loop_handle),
submission_transport: SubmissionTransport::Canonical,
};
parent_session
.guardian_review_session
@@ -7312,6 +7460,7 @@ async fn cached_guardian_subagent_exposes_its_rollout_path() {
agent_status: child_agent_status,
session: Arc::new(child_session),
session_loop_termination: session_loop_termination_from_handle(child_session_loop_handle),
submission_transport: SubmissionTransport::Canonical,
};
parent_session
.guardian_review_session
@@ -7345,6 +7494,7 @@ async fn shutdown_and_wait_shuts_down_tracked_ephemeral_guardian_review() {
agent_status: parent_agent_status,
session: Arc::clone(&parent_session),
session_loop_termination: session_loop_termination_from_handle(parent_session_loop_handle),
submission_transport: SubmissionTransport::Canonical,
};
let (child_session, _child_turn_context) = make_session_and_context().await;
@@ -7368,6 +7518,7 @@ async fn shutdown_and_wait_shuts_down_tracked_ephemeral_guardian_review() {
agent_status: child_agent_status,
session: Arc::new(child_session),
session_loop_termination: session_loop_termination_from_handle(child_session_loop_handle),
submission_transport: SubmissionTransport::Canonical,
};
parent_session
.guardian_review_session
@@ -7644,6 +7795,8 @@ where
pending_mcp_server_refresh_config: Mutex::new(None),
conversation: Arc::new(RealtimeConversationManager::new()),
active_turn: Mutex::new(None),
submission_lifecycle: SubmissionLifecycle::default(),
submission_send_lock: Mutex::new(()),
input_queue: super::input_queue::InputQueue::new(),
guardian_review_session: crate::guardian::GuardianReviewSessionManager::default(),
services,
@@ -9951,6 +10104,43 @@ async fn try_start_turn_if_idle_rejects_active_turn_without_injecting() {
sess.abort_all_tasks(TurnAbortReason::Interrupted).await;
}
#[tokio::test]
async fn idle_shutdown_claim_waits_for_submissions_and_blocks_automatic_turns() {
let (sess, _tc, _rx) = make_session_and_context_with_rx().await;
let pending = sess
.try_reserve_activity()
.expect("open session should reserve a submission");
assert!(sess.try_claim_idle_shutdown().await.is_none());
drop(pending);
sess.input_queue
.enqueue_mailbox_communication(InterAgentCommunication::new(
AgentPath::root(),
AgentPath::root(),
Vec::new(),
"pending context".to_string(),
/*trigger_turn*/ false,
))
.await;
assert!(sess.try_claim_idle_shutdown().await.is_none());
sess.input_queue.drain_mailbox_input_items().await;
let shutdown = sess
.try_claim_idle_shutdown()
.await
.expect("idle session should claim shutdown");
let item = user_message("synthetic idle input");
let err = sess
.try_start_turn_if_idle(vec![item.clone()])
.await
.expect_err("claimed shutdown should reject automatic idle work");
assert_eq!(TryStartTurnIfIdleRejectionReason::Busy, err.reason());
assert_eq!(vec![item], err.into_input());
assert!(sess.active_turn.lock().await.is_none());
drop(shutdown);
}
#[tokio::test]
async fn try_start_turn_if_idle_rejects_plan_mode_without_injecting() {
let (sess, _tc, _rx) = make_session_and_context_with_rx().await;
+4
View File
@@ -474,6 +474,9 @@ impl Session {
return;
}
let Some(reservation) = self.try_reserve_activity() else {
return;
};
{
let mut active_turn = self.active_turn.lock().await;
if active_turn.is_some() {
@@ -481,6 +484,7 @@ impl Session {
}
*active_turn = Some(ActiveTurn::default());
}
drop(reservation);
let turn_context = self.new_default_turn_with_sub_id(sub_id).await;
self.maybe_emit_model_warnings_for_turn(turn_context.as_ref())