diff --git a/codex-rs/core/src/agent/control.rs b/codex-rs/core/src/agent/control.rs index 696982b0a13..e6079d630f0 100644 --- a/codex-rs/core/src/agent/control.rs +++ b/codex-rs/core/src/agent/control.rs @@ -352,15 +352,6 @@ impl AgentControl { .ok_or(CodexErr::ThreadNotFound(agent_id)) } - pub(crate) async fn list_live_agent_subtree_thread_ids( - &self, - agent_id: ThreadId, - ) -> CodexResult> { - let mut thread_ids = vec![agent_id]; - thread_ids.extend(self.live_thread_spawn_descendants(agent_id).await?); - Ok(thread_ids) - } - pub(crate) async fn get_agent_config_snapshot( &self, agent_id: ThreadId, diff --git a/codex-rs/core/src/agent/control_tests.rs b/codex-rs/core/src/agent/control_tests.rs index 4057e02c909..afc6d5bcee7 100644 --- a/codex-rs/core/src/agent/control_tests.rs +++ b/codex-rs/core/src/agent/control_tests.rs @@ -2768,6 +2768,7 @@ async fn list_agent_subtree_thread_ids_finds_live_descendants_of_unloaded_root() .expect("grandchild spawn should succeed"); manager.remove_thread(&parent_thread_id).await; + manager.remove_thread(&child_thread_id).await; let mut subtree_thread_ids = manager .list_agent_subtree_thread_ids(parent_thread_id) diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index fc7f74f78cc..357c9db8d73 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/codex-rs/core/src/codex_thread.rs @@ -9,6 +9,7 @@ use codex_otel::SessionTelemetry; use codex_protocol::ThreadId; use codex_protocol::config_types::ApprovalsReviewer; use codex_protocol::config_types::CollaborationMode; +use codex_protocol::config_types::ModeKind; use codex_protocol::config_types::Personality; use codex_protocol::config_types::ReasoningSummary; use codex_protocol::config_types::WindowsSandboxLevel; @@ -81,6 +82,33 @@ pub struct ThreadConfigSnapshot { pub originator: String, } +/// Runtime configuration captured only while a thread is eligible for idle resume handling. +pub struct IdleResumeSnapshot { + pub thread_config_snapshot: ThreadConfigSnapshot, + pub original_config: Arc, +} + +/// Synchronous decision made while an idle thread is reserved against new work. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum IdleResumeDecision { + /// Keep using the existing thread. + Reuse, + /// Shut down the existing thread before replacing it. + Shutdown, +} + +/// Result of atomically resolving whether an existing thread can be reused. +pub enum IdleResumeOutcome { + /// Work is active or queued, so the caller should rejoin or retry later. + Active, + /// The reserved thread can be reused with this captured configuration. + Reused(Box), + /// This call committed a shutdown submission. + ShutdownCommitted, + /// Another call already committed shutdown for the thread. + ShutdownPending, +} + /// Explains why `CodexThread::try_start_turn_if_idle` rejected an automatic /// idle turn. #[derive(Clone, Copy, Debug, PartialEq, Eq)] @@ -130,6 +158,20 @@ impl ThreadConfigSnapshot { &self.environments.environments } + /// Checks the default-mode fields that must match before reusing an idle thread. + pub fn matches_default_resume_mode_and_environments( + &self, + desired_environments: &[TurnEnvironmentSelection], + ) -> bool { + self.collaboration_mode.mode == ModeKind::Default + && self + .collaboration_mode + .settings + .developer_instructions + .is_none() + && self.environments.environments == desired_environments + } + pub fn sandbox_policy(&self) -> SandboxPolicy { codex_sandboxing::compatibility_sandbox_policy_for_permission_profile( &self.permission_profile, @@ -215,6 +257,22 @@ impl CodexThread { self.codex.try_shutdown_if_idle().await } + /// Atomically decides whether an idle thread should be reused or shut down. + /// + /// `decide` runs synchronously while submission delivery is blocked. It must be fast and must + /// not re-enter this thread. + pub async fn resolve_idle_resume(&self, decide: F) -> CodexResult + where + F: FnOnce(&IdleResumeSnapshot) -> IdleResumeDecision + Send, + { + self.codex.resolve_idle_resume(decide).await + } + + /// Returns the immutable source assigned when the thread was created. + pub fn session_source(&self) -> &SessionSource { + &self.session_source + } + /// Wait until the underlying session loop has terminated. pub async fn wait_until_terminated(&self) { self.codex.session_loop_termination.clone().await; diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index cd27242609b..49046e3befe 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -26,6 +26,9 @@ mod config_lock; pub use codex_thread::BackgroundTerminalInfo; pub use codex_thread::CodexThread; pub use codex_thread::CodexThreadSettingsOverrides; +pub use codex_thread::IdleResumeDecision; +pub use codex_thread::IdleResumeOutcome; +pub use codex_thread::IdleResumeSnapshot; pub use codex_thread::ThreadConfigSnapshot; pub use codex_thread::TryStartTurnIfIdleError; pub use codex_thread::TryStartTurnIfIdleRejectionReason; diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 8b3f3a899d1..6ad86938d1d 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -175,6 +175,9 @@ use tracing::warn; use uuid::Uuid; use crate::client::ModelClient; +use crate::codex_thread::IdleResumeDecision; +use crate::codex_thread::IdleResumeOutcome; +use crate::codex_thread::IdleResumeSnapshot; use crate::codex_thread::ThreadConfigSnapshot; #[cfg(test)] use crate::compact::collect_user_messages; @@ -814,11 +817,25 @@ impl Codex { /// 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. + pub async fn try_shutdown_if_idle(&self) -> CodexResult { + match self + .resolve_idle_resume(|_| IdleResumeDecision::Shutdown) + .await? + { + IdleResumeOutcome::Active => Ok(false), + IdleResumeOutcome::ShutdownCommitted | IdleResumeOutcome::ShutdownPending => Ok(true), + IdleResumeOutcome::Reused(_) => unreachable!("shutdown decision cannot reuse thread"), + } + } + #[expect( clippy::await_holding_invalid_type, - reason = "the send lock serializes the idle claim with submission delivery" + reason = "the send lock and idle reservation make the reuse decision atomic with submission delivery" )] - pub async fn try_shutdown_if_idle(&self) -> CodexResult { + pub(crate) async fn resolve_idle_resume(&self, decide: F) -> CodexResult + where + F: FnOnce(&IdleResumeSnapshot) -> IdleResumeDecision + Send, + { if self.submission_transport == SubmissionTransport::ForwardingProxy { return Err(CodexErr::InvalidRequest( "idle shutdown is unavailable through a forwarding proxy".to_string(), @@ -826,26 +843,43 @@ impl Codex { } let _send_guard = self.session.submission_send_lock.lock().await; if self.session.submission_lifecycle.is_closing() { - return Ok(true); + return Ok(IdleResumeOutcome::ShutdownPending); } let Some(mut reservation) = self.session.try_claim_idle_shutdown().await else { - return Ok(false); + return Ok(IdleResumeOutcome::Active); }; - if !reservation.prepare_idle_shutdown_delivery() { - return Ok(false); - } - let sub = Submission { - id: new_submission_id(), - op: Op::Shutdown, - client_user_message_id: None, - trace: current_span_w3c_trace_context(), + let snapshot = { + let state = self.session.state.lock().await; + IdleResumeSnapshot { + thread_config_snapshot: state.session_configuration.thread_config_snapshot(), + original_config: Arc::clone( + &state.session_configuration.original_config_do_not_use, + ), + } }; - if self.tx_sub.send(sub).await.is_err() { - reservation.release_after_failed_delivery(); - return Err(CodexErr::InternalAgentDied); + match decide(&snapshot) { + IdleResumeDecision::Reuse => { + drop(reservation); + Ok(IdleResumeOutcome::Reused(Box::new(snapshot))) + } + IdleResumeDecision::Shutdown => { + if !reservation.prepare_idle_shutdown_delivery() { + return Ok(IdleResumeOutcome::Active); + } + 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(IdleResumeOutcome::ShutdownCommitted) + } } - reservation.commit(); - Ok(true) } /// Persist a thread-level memory mode update for the active session. @@ -1668,12 +1702,20 @@ impl Session { state.session_configuration.provider.clone() } + #[expect( + clippy::await_holding_invalid_type, + reason = "the send lock serializes the config swap and the activity reservation covers the full refresh" + )] pub(crate) async fn refresh_runtime_config(&self, next_config: Config) { // Refresh only the user layer from the incoming snapshot. Preserve thread-local // layers such as request/session overrides that were present when this session // was created. let notify_config_contributors = !self.services.extensions.config_contributors().is_empty(); - let (previous_config, new_config, config) = { + let (reservation, previous_config, new_config, config) = { + let _send_guard = self.submission_send_lock.lock().await; + let Some(reservation) = self.try_reserve_activity() else { + return; + }; let mut state = self.state.lock().await; let previous_config = notify_config_contributors .then(|| Self::build_effective_session_config(&state.session_configuration)); @@ -1687,7 +1729,7 @@ impl Session { state.session_configuration.original_config_do_not_use = Arc::clone(&config); let new_config = notify_config_contributors .then(|| Self::build_effective_session_config(&state.session_configuration)); - (previous_config, new_config, config) + (reservation, previous_config, new_config, config) }; self.emit_config_changed_contributors(previous_config.as_ref(), new_config.as_ref()); self.services.skills_service.clear_cache(); @@ -1709,6 +1751,8 @@ impl Session { ) { self.services.hooks.store(Arc::new(hooks)); } + drop(state); + drop(reservation); } fn emit_config_changed_contributors( diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 8175778b4f7..5e378dbdfe9 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -1611,6 +1611,45 @@ disabled_tools = [ ); } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[expect( + clippy::await_holding_invalid_type, + reason = "the test holds session state to pause refresh after it reserves activity" +)] +async fn refresh_runtime_config_blocks_idle_resume_until_the_config_swap_finishes() { + let (session, _turn_context) = make_session_and_context().await; + let session = Arc::new(session); + let original_config = session.get_config().await; + let next_config = (*original_config).clone(); + let state = session.state.lock().await; + let refresh = tokio::spawn({ + let session = Arc::clone(&session); + async move { session.refresh_runtime_config(next_config).await } + }); + + timeout(StdDuration::from_secs(2), async { + while let Ok(guard) = session.submission_send_lock.try_lock() { + drop(guard); + tokio::task::yield_now().await; + } + }) + .await + .expect("refresh should acquire the submission barrier"); + assert!(session.submission_send_lock.try_lock().is_err()); + assert!(session.try_claim_idle_shutdown().await.is_none()); + + drop(state); + refresh.await.expect("config refresh should not panic"); + + let refreshed_config = session.get_config().await; + assert!(!Arc::ptr_eq(&refreshed_config, &original_config)); + let idle = session + .try_claim_idle_shutdown() + .await + .expect("completed refresh should leave the session idle"); + drop(idle); +} + #[test] fn collect_explicit_app_ids_from_skill_items_includes_linked_mentions() { let connectors = vec![make_connector("calendar", "Calendar")]; @@ -5655,6 +5694,7 @@ async fn make_session_with_config_and_rx( Some(config.multi_agent_version_from_features()), ) .await?; + session.thread_activity.mark_initialized(); Ok((session, rx_event)) } @@ -10306,6 +10346,104 @@ async fn idle_shutdown_claim_waits_for_submissions_and_blocks_automatic_turns() drop(shutdown); } +#[tokio::test] +async fn idle_resume_resolution_is_atomic_with_pending_activity() { + 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); + let mut codex = Codex { + tx_sub, + rx_event, + agent_status, + session: Arc::clone(&session), + session_loop_termination: completed_session_loop_termination(), + submission_transport: SubmissionTransport::Canonical, + }; + + let pending = session + .try_reserve_activity() + .expect("open session should reserve activity"); + let outcome = codex + .resolve_idle_resume(|_| panic!("active resolution must not call the decider")) + .await + .expect("active resolution should succeed"); + assert!(matches!(outcome, IdleResumeOutcome::Active)); + drop(pending); + + codex.submission_transport = SubmissionTransport::ForwardingProxy; + let proxy_result = codex + .resolve_idle_resume(|_| panic!("proxy resolution must not call the decider")) + .await; + let Err(proxy_error) = proxy_result else { + panic!("forwarding proxies should reject idle resolution"); + }; + assert!(matches!(proxy_error, CodexErr::InvalidRequest(_))); + codex.submission_transport = SubmissionTransport::Canonical; + + let expected_config = session.get_config().await; + let outcome = codex + .resolve_idle_resume(|snapshot| { + assert!(session.submission_send_lock.try_lock().is_err()); + let expected_environments = snapshot + .thread_config_snapshot + .environment_selections() + .to_vec(); + assert!( + snapshot + .thread_config_snapshot + .matches_default_resume_mode_and_environments(&expected_environments) + ); + let mut non_default = snapshot.thread_config_snapshot.clone(); + non_default.collaboration_mode.mode = ModeKind::Plan; + assert!( + !non_default.matches_default_resume_mode_and_environments(&expected_environments) + ); + let mut developer_override = snapshot.thread_config_snapshot.clone(); + developer_override + .collaboration_mode + .settings + .developer_instructions = Some("override".to_string()); + assert!( + !developer_override + .matches_default_resume_mode_and_environments(&expected_environments) + ); + let unexpected_environment = TurnEnvironmentSelection { + environment_id: "unexpected".to_string(), + cwd: PathUri::from_abs_path(snapshot.thread_config_snapshot.cwd()), + }; + assert!( + !snapshot + .thread_config_snapshot + .matches_default_resume_mode_and_environments(&[unexpected_environment]) + ); + IdleResumeDecision::Reuse + }) + .await + .expect("idle reuse resolution should succeed"); + let IdleResumeOutcome::Reused(snapshot) = outcome else { + panic!("idle thread should be reusable"); + }; + assert!(Arc::ptr_eq(&snapshot.original_config, &expected_config)); + assert!(session.try_reserve_activity().is_some()); + + let outcome = codex + .resolve_idle_resume(|_| IdleResumeDecision::Shutdown) + .await + .expect("idle shutdown resolution should succeed"); + assert!(matches!(outcome, IdleResumeOutcome::ShutdownCommitted)); + assert_eq!( + rx_sub.recv().await.expect("shutdown submission").op, + Op::Shutdown + ); + let pending = codex + .resolve_idle_resume(|_| panic!("pending shutdown must not call the decider")) + .await + .expect("pending shutdown resolution should succeed"); + assert!(matches!(pending, IdleResumeOutcome::ShutdownPending)); +} + #[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; diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index 26cf0ff9411..07a4ca5d5b9 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -624,11 +624,7 @@ impl ThreadManager { } } - for descendant_id in self - .agent_control() - .list_live_agent_subtree_thread_ids(thread_id) - .await? - { + for descendant_id in self.in_memory_agent_subtree_thread_ids(thread_id) { if seen_thread_ids.insert(descendant_id) { subtree_thread_ids.push(descendant_id); } @@ -637,6 +633,17 @@ impl ThreadManager { Ok(subtree_thread_ids) } + /// List `thread_id` plus descendants registered in the in-memory activity graph. + pub fn in_memory_agent_subtree_thread_ids(&self, thread_id: ThreadId) -> Vec { + let mut thread_ids = vec![thread_id]; + thread_ids.extend( + self.state + .thread_activity_gate + .descendant_thread_ids(thread_id), + ); + thread_ids + } + pub async fn start_thread(&self, config: Config) -> CodexResult { // Box delegated thread-spawn futures so these convenience wrappers do // not inline the full spawn path into every caller's async state. @@ -1065,6 +1072,7 @@ impl ThreadManager { .await } + #[cfg(test)] pub(crate) fn agent_control(&self) -> AgentControl { AgentControl::new(Arc::downgrade(&self.state), /*rollout_budget*/ None) } diff --git a/codex-rs/core/tests/suite/pending_input.rs b/codex-rs/core/tests/suite/pending_input.rs index 9c9852a2764..264ca982f2c 100644 --- a/codex-rs/core/tests/suite/pending_input.rs +++ b/codex-rs/core/tests/suite/pending_input.rs @@ -192,7 +192,7 @@ async fn steer_user_input(codex: &CodexThread, text: &str) { .expect("steer user input"); } -async fn submit_queue_only_agent_mail(codex: &CodexThread, text: &str) { +async fn submit_agent_mail(codex: &CodexThread, text: &str, trigger_turn: bool) { codex .submit(Op::InterAgentCommunication { communication: InterAgentCommunication::new( @@ -200,7 +200,7 @@ async fn submit_queue_only_agent_mail(codex: &CodexThread, text: &str) { AgentPath::root(), Vec::new(), text.to_string(), - /*trigger_turn*/ false, + trigger_turn, ), }) .await @@ -215,6 +215,10 @@ async fn submit_queue_only_agent_mail(codex: &CodexThread, text: &str) { .await; } +async fn submit_queue_only_agent_mail(codex: &CodexThread, text: &str) { + submit_agent_mail(codex, text, /*trigger_turn*/ false).await; +} + async fn wait_for_reasoning_item_started(codex: &CodexThread) { wait_for_event(codex, |event| { matches!( @@ -596,6 +600,64 @@ async fn injected_user_input_triggers_follow_up_request_with_deltas() { server.shutdown().await; } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn trigger_turn_mail_starts_after_the_current_turn_completes() { + let (gate_first_completed_tx, gate_first_completed_rx) = oneshot::channel(); + let first_chunks = vec![ + chunk(ev_response_created("resp-1")), + chunk(ev_message_item_added("msg-1", "")), + chunk(ev_output_text_delta("first answer")), + chunk(ev_message_item_done("msg-1", "first answer")), + gated_chunk(gate_first_completed_rx, vec![ev_completed("resp-1")]), + ]; + let (server, _completions) = + start_streaming_sse_server(vec![first_chunks, response_completed_chunks("resp-2")]).await; + let codex = build_codex(&server).await; + + submit_user_input(&codex, "first prompt").await; + wait_for_agent_message(&codex, "first answer").await; + submit_agent_mail(&codex, "queued trigger", /*trigger_turn*/ true).await; + gate_first_completed_tx + .send(()) + .expect("first response gate should be open"); + + let mut completed_turn_id = None; + let started = wait_for_event(&codex, |event| match event { + EventMsg::TurnComplete(event) => { + completed_turn_id = Some(event.turn_id.clone()); + false + } + EventMsg::TurnStarted(_) => completed_turn_id.is_some(), + _ => false, + }) + .await; + let EventMsg::TurnStarted(started) = started else { + unreachable!("wait predicate only accepts turn/started events"); + }; + let first_turn_id = completed_turn_id.expect("current turn should complete first"); + let second_turn_id = started.turn_id; + assert_ne!(first_turn_id, second_turn_id); + wait_for_turn_complete(&codex).await; + + let requests = server.requests().await; + assert_eq!(requests.len(), 2); + let second: Value = from_slice(&requests[1]).expect("parse follow-up request"); + let queued_message_is_visible = second["input"] + .as_array() + .into_iter() + .flatten() + .any(|item| { + item["type"] == "agent_message" + && item["content"] + .as_array() + .into_iter() + .flatten() + .any(|content| content["text"] == "queued trigger") + }); + assert!(queued_message_is_visible); + server.shutdown().await; +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn queued_inter_agent_mail_triggers_follow_up_after_reasoning_item() { let (gate_reasoning_done_tx, gate_reasoning_done_rx) = oneshot::channel();