diff --git a/codex-rs/core/config.schema.json b/codex-rs/core/config.schema.json index 4aba1bab6cd2..ac21ee392fe3 100644 --- a/codex-rs/core/config.schema.json +++ b/codex-rs/core/config.schema.json @@ -2937,6 +2937,12 @@ "MultiAgentV2ConfigToml": { "additionalProperties": false, "properties": { + "completed_agent_unload_delay_secs": { + "description": "Unload idle terminal children after this delay. Omit to disable automatic unload.", + "format": "uint64", + "minimum": 0.0, + "type": "integer" + }, "default_wait_timeout_ms": { "format": "int64", "maximum": 3600000.0, diff --git a/codex-rs/core/src/agent/control.rs b/codex-rs/core/src/agent/control.rs index 895cd26802d7..020978fc0498 100644 --- a/codex-rs/core/src/agent/control.rs +++ b/codex-rs/core/src/agent/control.rs @@ -241,14 +241,44 @@ impl LocalAgentControl { context: AgentCommunicationContext, start_options: TurnStartOptions, ) -> CodexResult { - self.submit_inter_agent_communication( - agent_id, - state, - communication, - context, - start_options, - ) - .await + if communication.trigger_turn { + return self + .submit_inter_agent_communication( + agent_id, + state, + communication, + context, + start_options, + ) + .await; + } + let result = self + .submit_inter_agent_communication( + agent_id, + state, + communication.clone(), + context.clone(), + start_options.clone(), + ) + .await; + // Eviction can remove a loaded recipient before its residency pin succeeds. + // Retry queue-only delivery once through the mailbox path, including results. + // ThreadNotFound precedes submission; accepted requests are never replayed. + if result + .as_ref() + .is_err_and(|err| matches!(err.details(), CodexErrorDetails::ThreadNotFound(_))) + { + self.submit_inter_agent_communication( + agent_id, + state, + communication, + context, + start_options, + ) + .await + } else { + result + } } async fn submit_inter_agent_communication( diff --git a/codex-rs/core/src/agent/control/residency.rs b/codex-rs/core/src/agent/control/residency.rs index 4755c221f8be..61f6a69755eb 100644 --- a/codex-rs/core/src/agent/control/residency.rs +++ b/codex-rs/core/src/agent/control/residency.rs @@ -6,7 +6,10 @@ use super::runtime::AgentTreeMembership; use crate::agent::AgentStatus; use crate::codex_thread::CodexThread; use crate::config::Config; +use crate::session::session::Session; use crate::thread_manager::ThreadManagerState; +use codex_history::InitialHistory; +use codex_history::ResumedHistory; use codex_protocol::ThreadId; use codex_protocol::error::AgentErrorContext; use codex_protocol::error::CodexErr; @@ -15,9 +18,13 @@ use codex_protocol::error::CodexErrorDetails; use codex_protocol::error::Result as CodexResult; use codex_protocol::protocol::MultiAgentVersion; use codex_protocol::protocol::SessionSource; +use codex_thread_store::ReadThreadParams; +use std::collections::HashMap; use std::collections::VecDeque; use std::sync::Arc; use std::sync::Mutex; +use std::time::Duration; +use tokio_util::sync::CancellationToken; use tracing::warn; /// The result of trying to unload one exact MAv2 child runtime. @@ -51,6 +58,17 @@ impl ThreadManagerState { if !is_resident_candidate(&thread) || !is_unloadable(&thread).await { return Ok(ThreadEvictionOutcome::Busy); } + self.start_v2_eviction(thread, membership, residency_guard) + .await? + } + + fn start_v2_eviction( + self: &Arc, + thread: Arc, + membership: &AgentTreeMembership, + residency_guard: tokio::sync::OwnedRwLockWriteGuard<()>, + ) -> tokio::task::JoinHandle> { + let thread_id = thread.session.thread_id; let manager = Arc::clone(self); let teardown = membership .clone() @@ -97,7 +115,6 @@ impl ThreadManagerState { ThreadEvictionOutcome::NotCurrent }) }) - .await? } } @@ -110,6 +127,21 @@ pub(super) struct V2Residency { struct V2ResidencyState { residents: VecDeque, pending_slots: usize, + unload_timers: HashMap>, + #[cfg(test)] + unload_attempts: HashMap, +} + +impl Drop for V2Residency { + fn drop(&mut self) { + let state = self + .state + .get_mut() + .unwrap_or_else(std::sync::PoisonError::into_inner); + for timer in state.unload_timers.values() { + timer.cancel(); + } + } } pub(super) struct V2ResidencySlot { @@ -148,24 +180,144 @@ impl LocalAgentControl { .await } + pub(super) async fn touch_loaded_v2_residency( + &self, + state: &Arc, + thread_id: ThreadId, + ) -> CodexResult>> { + let thread = state.get_thread(thread_id).await?; + let guard = self.runtime.pin_v2_residency(&thread).await?; + self.runtime.cancel_completed_agent_unload(thread_id); + Ok(guard) + } + pub(super) fn forget_v2_residency(&self, thread_id: ThreadId) { self.runtime.residency.remove(thread_id); } } impl LocalAgentRuntime { + #[cfg(test)] + pub(crate) fn has_completed_agent_unload_timer(&self, thread_id: ThreadId) -> bool { + self.residency + .state + .lock() + .unwrap() + .unload_timers + .contains_key(&thread_id) + } + + #[cfg(test)] + pub(crate) fn completed_agent_unload_attempts(&self, thread_id: ThreadId) -> usize { + self.residency + .state + .lock() + .unwrap() + .unload_attempts + .get(&thread_id) + .copied() + .unwrap_or_default() + } + + pub(crate) fn cancel_completed_agent_unload(&self, thread_id: ThreadId) { + self.residency.cancel_unload_timer(thread_id); + } + + pub(crate) fn schedule_completed_agent_unload(&self, session: &Arc, delay_secs: u64) { + let thread_id = session.thread_id; + let Some(deadline) = + tokio::time::Instant::now().checked_add(Duration::from_secs(delay_secs)) + else { + warn!(%thread_id, "completed agent unload delay exceeds the timer range; skip unload"); + return; + }; + let Ok(membership) = self.admit_start() else { + return; + }; + let timer = Arc::new(CancellationToken::new()); + { + let mut state = self + .residency + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + #[cfg(test)] + state.unload_attempts.insert(thread_id, 0); + if let Some(previous) = state.unload_timers.insert(thread_id, Arc::clone(&timer)) { + previous.cancel(); + } + } + let residency = Arc::downgrade(&self.residency); + let manager = self.manager.clone(); + let session = Arc::downgrade(session); + let shutdown = self.shutdown.clone(); + tokio::spawn(async move { + let unload = async { + tokio::time::sleep_until(deadline).await; + for attempt in 0..2 { + let done = match (residency.upgrade(), manager.upgrade(), session.upgrade()) { + (Some(residency), Some(manager), Some(session)) => { + let done = residency + .try_unload_completed(&manager, &membership, &session, &timer) + .await; + #[cfg(test)] + { + let mut state = residency.state.lock().unwrap(); + if state + .unload_timers + .get(&thread_id) + .is_some_and(|current| Arc::ptr_eq(current, &timer)) + { + *state.unload_attempts.entry(thread_id).or_default() += 1; + } + } + done + } + _ => true, + }; + if done { + return; + } + if attempt == 0 { + tokio::time::sleep(Duration::from_secs(5)).await; + } + } + warn!(%thread_id, "completed agent unload is blocked after one retry; keep runtime loaded"); + }; + tokio::select! { + biased; + _ = shutdown.cancelled() => {}, + _ = timer.cancelled() => {}, + _ = unload => {}, + } + if let Some(residency) = residency.upgrade() { + let mut state = residency + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if state + .unload_timers + .get(&thread_id) + .is_some_and(|current| Arc::ptr_eq(current, &timer)) + { + state.unload_timers.remove(&thread_id); + } + } + }); + } + /// Pins and touches the registered runtime without waiting for unrelated eviction. pub(crate) async fn pin_v2_residency( &self, - state: &ThreadManagerState, - thread: &Arc, + thread: &CodexThread, ) -> CodexResult>> { if !is_resident_candidate(thread) { return Ok(None); } let guard = Arc::clone(&thread.residency_gate).read_owned().await; let thread_id = thread.session.thread_id; - if !Arc::ptr_eq(thread, &state.get_thread(thread_id).await?) { + let registered = self.upgrade()?.get_thread(thread_id).await?; + if !std::ptr::eq(thread, registered.as_ref()) { return Err(CodexErr::ThreadNotFound(thread_id)); } self.residency.touch(thread_id); @@ -174,6 +326,129 @@ impl LocalAgentRuntime { } impl V2Residency { + async fn try_unload_completed( + self: &Arc, + manager: &Arc, + membership: &AgentTreeMembership, + session: &Arc, + timer: &CancellationToken, + ) -> bool { + let thread_id = session.thread_id; + let Ok(thread) = manager.get_thread(thread_id).await else { + return true; + }; + if !Arc::ptr_eq(session, &thread.session) || !is_resident_candidate(&thread) { + return true; + } + if !is_completed_unloadable(&thread).await || has_background_work(&thread).await { + return false; + } + let runtime = &session.services.local_agent_runtime; + // Startup publishes metadata after the initial turn is admitted. + if runtime + .registry + .agent_metadata_for_thread(thread_id) + .is_none() + { + return false; + } + let reloadable = async { + let stored = manager + .read_stored_thread(ReadThreadParams { + thread_id, + include_archived: true, + include_history: false, + }) + .await?; + let Some(history) = + super::spawn::load_agent_model_context(manager, thread_id, stored.history_mode) + .await? + else { + return Ok::<_, CodexErr>(false); + }; + let history = InitialHistory::Resumed(ResumedHistory { + history_revision: history.revision, + conversation_id: thread_id, + history: Arc::new(history.items), + rollout_path: stored.rollout_path, + }); + Ok(history.get_multi_agent_version() == Some(MultiAgentVersion::V2)) + } + .await; + match reloadable { + Ok(true) => {} + Ok(false) => { + warn!(%thread_id, "completed agent has no reloadable saved history; keep runtime loaded"); + return true; + } + Err(err) => { + warn!(%thread_id, "failed to check completed agent history: {err}"); + return false; + } + } + // Read saved history without blocking input. Recheck the runtime and its + // activity after acquiring exclusion because a turn can start during the read. + let Ok(guard) = Arc::clone(&thread.residency_gate).try_write_owned() else { + return false; + }; + if timer.is_cancelled() + || !manager + .get_thread(thread_id) + .await + .is_ok_and(|current| Arc::ptr_eq(¤t, &thread)) + || !is_completed_unloadable(&thread).await + || has_background_work(&thread).await + || timer.is_cancelled() + { + return false; + } + let Ok(descendants) = runtime.live_thread_spawn_descendants(thread_id).await else { + return false; + }; + let mut checked_threads = vec![Arc::clone(&thread)]; + for descendant in descendants { + if let Ok(descendant) = manager.get_thread(descendant).await { + if !is_completed_unloadable(&descendant).await + || has_background_work(&descendant).await + { + return false; + } + checked_threads.push(descendant); + } + } + let eviction = { + // No await between the final admission checks and teardown spawn. A turn + // admitted after this point can deliver its result to the unloaded parent. + let mut admission_guards = Vec::with_capacity(checked_threads.len()); + for checked in &checked_threads { + let Ok(active) = checked.session.active_turn.try_lock() else { + return false; + }; + if active.is_some() { + return false; + } + admission_guards.push(active); + } + if timer.is_cancelled() { + return false; + } + manager.start_v2_eviction(thread, membership, guard) + }; + finish_eviction(eviction).await + } + + fn cancel_unload_timer(&self, thread_id: ThreadId) { + if let Some(timer) = self + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .unload_timers + .remove(&thread_id) + { + timer.cancel(); + } + } + async fn reserve_slot( self: Arc, manager: &Arc, @@ -267,10 +542,16 @@ impl V2Residency { .state .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); - touch_resident(&mut state.residents, thread_id); + // Initial input can pin a published runtime before its reserved slot commits. + // V2 spawn and lazy reload reserve a slot, then commit it after admission or + // registration. Only commit_slot inserts residents, so startup counts once. + if state.residents.contains(&thread_id) { + touch_resident(&mut state.residents, thread_id); + } } fn remove(&self, thread_id: ThreadId) { + self.cancel_unload_timer(thread_id); self.state .lock() .unwrap_or_else(std::sync::PoisonError::into_inner) @@ -301,6 +582,23 @@ fn touch_resident(residents: &mut VecDeque, thread_id: ThreadId) { residents.push_back(thread_id); } +async fn finish_eviction( + eviction: tokio::task::JoinHandle>, +) -> bool { + match eviction.await { + Ok(Ok(ThreadEvictionOutcome::Evicted | ThreadEvictionOutcome::NotCurrent)) => true, + Ok(Ok(ThreadEvictionOutcome::Busy)) => false, + Ok(Err(err)) => { + warn!("v2 resident eviction failed: {err}"); + false + } + Err(err) => { + warn!("v2 resident eviction task failed: {err}"); + false + } + } +} + fn is_resident_candidate(thread: &CodexThread) -> bool { thread.multi_agent_version() == Some(MultiAgentVersion::V2) && is_v2_resident_session_source(&thread.session_source) @@ -326,6 +624,20 @@ async fn is_unloadable(thread: &CodexThread) -> bool { .await } +async fn is_completed_unloadable(thread: &CodexThread) -> bool { + matches!( + thread.agent_status().await, + AgentStatus::Completed(_) | AgentStatus::Errored(_) | AgentStatus::Interrupted + ) && is_unloadable(thread).await +} + +async fn has_background_work(thread: &CodexThread) -> bool { + let awaiting_elicitation = *thread.session.subscribe_elicitation_pause_state().borrow(); + awaiting_elicitation + || thread.session.services.code_mode_service.has_active_cells() + || !thread.session.list_background_terminals().await.is_empty() +} + #[cfg(test)] #[path = "residency_tests.rs"] mod tests; diff --git a/codex-rs/core/src/agent/control/residency_tests.rs b/codex-rs/core/src/agent/control/residency_tests.rs index 5bf450e456ba..27132c3057a4 100644 --- a/codex-rs/core/src/agent/control/residency_tests.rs +++ b/codex-rs/core/src/agent/control/residency_tests.rs @@ -18,6 +18,99 @@ use codex_protocol::protocol::TurnAbortedEvent; use codex_protocol::protocol::TurnCompleteEvent; use pretty_assertions::assert_eq; use std::sync::Arc; +use std::time::Duration; + +#[test] +fn completed_agent_unload_pin_during_startup_does_not_count_reservation_twice() { + let residency = Arc::new(super::V2Residency::default()); + let first_id = ThreadId::new(); + assert!(residency.try_reserve_pending_slot(2)); + let first = super::V2ResidencySlot { + residency: Arc::clone(&residency), + active: true, + }; + // Initial input pins the published thread before spawn commits its reserved slot. + residency.touch(first_id); + assert!( + residency.try_reserve_pending_slot(2), + "a pending startup must use exactly one slot" + ); + let second = super::V2ResidencySlot { + residency: Arc::clone(&residency), + active: true, + }; + first.commit(first_id); + assert!(!residency.try_reserve_pending_slot(2)); + drop(second); + assert!(residency.try_reserve_pending_slot(2)); +} + +#[tokio::test] +async fn completed_agent_unload_timer_does_not_retain_session_or_delay_tree_shutdown() { + let (session, turn) = crate::session::tests::make_session_and_context().await; + let session = Arc::new(session); + let weak_session = Arc::downgrade(&session); + let runtime = session.services.local_agent_runtime.clone(); + runtime.schedule_completed_agent_unload(&session, 60); + drop(turn); + drop(session); + assert!( + weak_session.upgrade().is_none(), + "pending timer must not own the session" + ); + tokio::time::timeout(Duration::from_secs(1), runtime.request_shutdown().wait()) + .await + .expect("pending timer must leave the tree promptly") + .expect("tree shutdown"); +} + +#[tokio::test] +async fn completed_agent_unload_replaces_deadline_and_turn_admission_cancels_it() { + let (session, _) = crate::session::tests::make_session_and_context().await; + let session = Arc::new(session); + let runtime = &session.services.local_agent_runtime; + runtime.schedule_completed_agent_unload(&session, 60); + let first = + Arc::clone(&runtime.residency.state.lock().unwrap().unload_timers[&session.thread_id]); + runtime.schedule_completed_agent_unload(&session, 120); + let second = + Arc::clone(&runtime.residency.state.lock().unwrap().unload_timers[&session.thread_id]); + assert!(first.is_cancelled()); + assert!(!second.is_cancelled()); + assert!(!Arc::ptr_eq(&first, &second)); + session.record_started_turn("new-turn", None).await; + assert!(second.is_cancelled()); + assert!( + runtime + .residency + .state + .lock() + .unwrap() + .unload_timers + .is_empty() + ); + tokio::time::timeout(Duration::from_secs(1), runtime.request_shutdown().wait()) + .await + .expect("cancelled deadlines must leave the tree") + .expect("tree shutdown"); +} + +#[tokio::test] +async fn completed_agent_unload_skips_unrepresentable_deadline() { + let (session, _) = crate::session::tests::make_session_and_context().await; + let session = Arc::new(session); + let runtime = &session.services.local_agent_runtime; + runtime.schedule_completed_agent_unload(&session, u64::MAX); + assert!( + runtime + .residency + .state + .lock() + .unwrap() + .unload_timers + .is_empty() + ); +} #[tokio::test] async fn residency_slot_reservation_unloads_oldest_idle_v2_agent() { diff --git a/codex-rs/core/src/agent/control/spawn.rs b/codex-rs/core/src/agent/control/spawn.rs index 70cf8b7ef07e..acaf47994bfb 100644 --- a/codex-rs/core/src/agent/control/spawn.rs +++ b/codex-rs/core/src/agent/control/spawn.rs @@ -157,7 +157,7 @@ fn retain_forked_developer_message(item: &mut ResponseItem, usage_hint_texts: &[ !content.is_empty() && set_annotated_content(item, content).is_some() } -async fn load_agent_model_context( +pub(super) async fn load_agent_model_context( state: &ThreadManagerState, thread_id: ThreadId, history_mode: ThreadHistoryMode, @@ -342,9 +342,9 @@ impl LocalAgentControl { } } if owner_thread_id.is_none() - && let Ok(thread) = state.get_thread(thread_id).await + && let Ok(pin) = self.touch_loaded_v2_residency(&state, thread_id).await { - return self.runtime.pin_v2_residency(&state, &thread).await; + return Ok(pin); } if self .runtime @@ -397,7 +397,9 @@ impl LocalAgentControl { } if let Ok(thread) = state.get_thread(thread_id).await { self.validate_loaded_v2_child(&thread, parent_thread_id)?; - return self.runtime.pin_v2_residency(&state, &thread).await; + if let Ok(pin) = self.touch_loaded_v2_residency(&state, thread_id).await { + return Ok(pin); + } } } let parent = if let Some(parent) = parent { @@ -631,9 +633,7 @@ impl LocalAgentControl { // Finish registration even if the sender is cancelled while acquiring its pin. tokio::spawn(async move { let _membership = membership; - let residency_pin = runtime - .pin_v2_residency(&state, &reloaded_thread.thread) - .await?; + let residency_pin = runtime.pin_v2_residency(&reloaded_thread.thread).await?; runtime.registry.clear_evicted_environments(thread_id); residency_slot.commit(reloaded_thread.thread_id); // Register before listeners can forward events from the resumed thread. @@ -666,7 +666,7 @@ impl LocalAgentControl { } self.runtime.registry.clear_evicted_environments(thread_id); drop(residency_slot); - return self.runtime.pin_v2_residency(&state, &thread).await; + return self.touch_loaded_v2_residency(&state, thread_id).await; } Err(err) } diff --git a/codex-rs/core/src/agent/control_tests.rs b/codex-rs/core/src/agent/control_tests.rs index 6d97eceb2ed4..9f80bcba4065 100644 --- a/codex-rs/core/src/agent/control_tests.rs +++ b/codex-rs/core/src/agent/control_tests.rs @@ -104,6 +104,7 @@ use codex_thread_store::LocalThreadStoreConfig; use codex_thread_store::PersistContext; use codex_thread_store::ThreadStore; use codex_utils_path_uri::PathUri; +use core_test_support::responses::strip_metadata_from_items; use core_test_support::responses::strip_response_item_ids; use pretty_assertions::assert_eq; use std::future::Future; @@ -1081,6 +1082,793 @@ async fn ensure_v2_agent_loaded_reloads_registered_unloaded_agent() { check_v2_agent_reload(V2ReloadRoute::Sender).await; } +#[derive(Clone, Copy)] +enum CompletedAgentUnloadTask { + Completed, + Errored, + Interrupted, + Pending, +} + +impl crate::tasks::SessionTask for CompletedAgentUnloadTask { + fn kind(&self) -> crate::state::TaskKind { + crate::state::TaskKind::Regular + } + + fn span_name(&self) -> &'static str { + "completed_agent_unload_test" + } + + async fn run( + self: Arc, + _session: Arc, + _ctx: Arc, + _input: Vec, + _cancellation_token: CancellationToken, + ) -> crate::tasks::SessionTaskResult { + match self.as_ref() { + Self::Completed => Ok(Some("completed child result".to_string())), + Self::Errored => Err(CodexErr::Fatal("child failure".to_string())), + Self::Interrupted => Err(CodexErr::TurnAborted), + Self::Pending => std::future::pending().await, + } + } +} + +#[derive(Clone, Copy)] +enum CompletedAgentUnloadBlocker { + Disabled, + Root, + Ephemeral, + Sleep, + Elicitation, + CodeModeCell, + BackgroundTerminal, + TriggerMail, + ResumedTurn, + ActiveDescendant, + TreeShutdown, + ExtremeDelay, +} + +async fn persist_completed_agent_unload_fixture( + thread: &Arc, + item: &ResponseItem, +) -> Option> { + let guard = thread + .session + .services + .local_agent_runtime + .pin_v2_residency(thread) + .await + .expect("pin fixture setup"); + // Injection into an active turn is queued. Replacement discards that queue. + thread + .session + .abort_all_tasks(TurnAbortReason::Replaced) + .await; + thread + .inject_response_items(vec![item.clone()]) + .await + .expect("record fixture history while idle"); + thread.flush_rollout().await.expect("flush fixture history"); + let history = thread.session.clone_history().await; + let actual = strip_response_item_ids(&strip_metadata_from_items( + &history.raw_items().cloned().collect::>(), + )); + let expected = strip_response_item_ids(&strip_metadata_from_items(std::slice::from_ref(item))); + assert!( + actual.contains(&expected[0]), + "fixture history must exist before terminal completion: actual={actual:?}, expected={expected:?}" + ); + guard +} + +#[test_case::test_case(CompletedAgentUnloadBlocker::Disabled; "disabled")] +#[test_case::test_case(CompletedAgentUnloadBlocker::Root; "root")] +#[test_case::test_case(CompletedAgentUnloadBlocker::Ephemeral; "ephemeral")] +#[test_case::test_case(CompletedAgentUnloadBlocker::Sleep; "durable_sleep")] +#[test_case::test_case(CompletedAgentUnloadBlocker::Elicitation; "elicitation")] +#[test_case::test_case(CompletedAgentUnloadBlocker::CodeModeCell; "code_mode_cell")] +#[cfg_attr(unix, test_case::test_case(CompletedAgentUnloadBlocker::BackgroundTerminal; "background_terminal"))] +#[test_case::test_case(CompletedAgentUnloadBlocker::TriggerMail; "trigger_mail")] +#[test_case::test_case(CompletedAgentUnloadBlocker::ResumedTurn; "resumed_turn")] +#[test_case::test_case(CompletedAgentUnloadBlocker::ActiveDescendant; "active_descendant")] +#[test_case::test_case(CompletedAgentUnloadBlocker::TreeShutdown; "tree_shutdown")] +#[test_case::test_case(CompletedAgentUnloadBlocker::ExtremeDelay; "extreme_delay")] +#[tokio::test] +async fn completed_agent_unload_retains_required_runtime(blocker: CompletedAgentUnloadBlocker) { + let delay = match blocker { + CompletedAgentUnloadBlocker::TreeShutdown => 60, + CompletedAgentUnloadBlocker::ExtremeDelay => i64::MAX, + CompletedAgentUnloadBlocker::BackgroundTerminal => 2, + _ => 1, + }; + let overrides = if matches!(blocker, CompletedAgentUnloadBlocker::Disabled) { + Vec::new() + } else { + vec![( + "features.multi_agent_v2.completed_agent_unload_delay_secs".to_string(), + TomlValue::Integer(delay), + )] + }; + let (home, mut config) = test_config_with_cli_overrides(overrides).await; + let _ = config.features.enable(Feature::MultiAgentV2); + let _ = config.features.enable(Feature::Sqlite); + config.ephemeral = matches!(blocker, CompletedAgentUnloadBlocker::Ephemeral); + config.agent_max_depth = 2; + if matches!(blocker, CompletedAgentUnloadBlocker::BackgroundTerminal) { + config + .permissions + .set_permission_profile(PermissionProfile::Disabled) + .expect("test process permissions"); + config + .permissions + .approval_policy + .set(AskForApproval::Never) + .expect("test process approval"); + } + let harness = AgentControlHarness::new_with_config(home, config).await; + let (root_id, root) = if matches!(blocker, CompletedAgentUnloadBlocker::BackgroundTerminal) { + harness.start_thread().await + } else { + harness.start_paginated_thread().await + }; + let control = root + .session + .services + .local_agent_runtime + .control(root.session.session_id()); + let thread = if matches!(blocker, CompletedAgentUnloadBlocker::Root) { + Arc::clone(&root) + } else { + let child = + spawn_v2_reload_test_child(&control, harness.config.clone(), &root, "worker").await; + harness + .manager + .get_thread(child.thread_id) + .await + .expect("loaded child") + }; + let setup_guard = if !matches!(blocker, CompletedAgentUnloadBlocker::Ephemeral) { + persist_completed_agent_unload_fixture( + &thread, + &assistant_message("saved child history", Some(MessagePhase::FinalAnswer)), + ) + .await + } else { + None + }; + if matches!(blocker, CompletedAgentUnloadBlocker::Sleep) { + thread + .thread_extension_data() + .insert(codex_extension_items::sleep::SleepItem { + id: "durable-sleep".to_string(), + duration_ms: 60_000, + }); + } + let elicitation = matches!(blocker, CompletedAgentUnloadBlocker::Elicitation) + .then(|| thread.session.services.elicitations.register()); + if matches!(blocker, CompletedAgentUnloadBlocker::ActiveDescendant) { + let descendant = + spawn_v2_reload_test_child(&control, harness.config.clone(), &thread, "nested").await; + let descendant = harness + .manager + .get_thread(descendant.thread_id) + .await + .expect("loaded descendant"); + let turn = descendant.session.new_default_turn().await; + descendant + .session + .spawn_task(turn, Vec::new(), CompletedAgentUnloadTask::Pending) + .await; + } + let turn = thread.session.new_default_turn().await; + thread + .session + .spawn_task(turn, Vec::new(), CompletedAgentUnloadTask::Completed) + .await; + let expects_timer = !matches!( + blocker, + CompletedAgentUnloadBlocker::Disabled + | CompletedAgentUnloadBlocker::Root + | CompletedAgentUnloadBlocker::Ephemeral + | CompletedAgentUnloadBlocker::ExtremeDelay + ); + timeout(Duration::from_secs(10), async { + while !matches!(thread.agent_status().await, AgentStatus::Completed(_)) + || thread.session.active_turn.lock().await.is_some() + || (expects_timer + && !thread + .session + .services + .local_agent_runtime + .has_completed_agent_unload_timer(thread.session.thread_id)) + { + sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("terminal turn should finish and arm its unload timer when enabled"); + if matches!(blocker, CompletedAgentUnloadBlocker::BackgroundTerminal) { + let turn = thread.session.new_default_turn().await; + let manager = &thread.session.services.unified_exec_manager; + let context = crate::unified_exec::UnifiedExecContext::new( + Arc::clone(&thread.session), + StepContext::for_test(Arc::clone(&turn)), + CancellationToken::new(), + "unload-background-process".to_string(), + ); + let environment = turn + .initial_environments + .primary() + .cloned() + .expect("local environment"); + let request = crate::unified_exec::ExecCommandRequest { + command: vec!["sh".to_string(), "-c".to_string(), "sleep 30".to_string()], + shell: crate::shell::ShellInvocation { + shell: crate::shell::Shell { + shell_type: crate::shell::ShellType::Sh, + shell_path: "sh".into(), + }, + use_login_shell: false, + }, + hook_command: "sleep 30".to_string(), + process_id: manager.allocate_process_id().await, + yield_time_ms: 250, + max_output_tokens: Some(128), + cwd: environment.selection.cwd.clone(), + sandbox_cwd: environment.selection.cwd.clone(), + turn_environment: environment, + network: None, + tty: false, + sandbox_permissions: crate::sandboxing::SandboxPermissions::UseDefault, + additional_permissions: None, + additional_permissions_preapproved: false, + justification: None, + prefix_rule: None, + }; + let output = manager + .exec_command(request, &context) + .await + .expect("start real background process"); + assert!(output.process_id.is_some()); + assert!(!thread.list_background_terminals().await.is_empty()); + } + let cell_id = codex_code_mode::CellId::new("unload-guard-cell".to_string()); + if matches!(blocker, CompletedAgentUnloadBlocker::CodeModeCell) { + thread + .session + .services + .code_mode_service + .mark_cell_ready_for_dispatch(&cell_id, None); + } + if matches!(blocker, CompletedAgentUnloadBlocker::TriggerMail) { + thread + .session + .input_queue + .enqueue_mailbox_communication( + InterAgentCommunication::new( + AgentPath::root(), + AgentPath::try_from("/root/worker").expect("agent path"), + Vec::new(), + "pending followup".to_string(), + true, + ), + Default::default(), + ) + .await; + } + if matches!(blocker, CompletedAgentUnloadBlocker::ResumedTurn) { + let turn = thread.session.new_default_turn().await; + thread + .session + .spawn_task(turn, Vec::new(), CompletedAgentUnloadTask::Pending) + .await; + } + let runtime = &thread.session.services.local_agent_runtime; + let expects_attempt = expects_timer + && !matches!( + blocker, + CompletedAgentUnloadBlocker::ResumedTurn | CompletedAgentUnloadBlocker::TreeShutdown + ); + if expects_attempt { + // Replace any deadline that elapsed while the fixture held its setup pin. + runtime.schedule_completed_agent_unload(&thread.session, delay as u64); + } + // No await between arming the fresh deadline and releasing the setup pin. + drop(setup_guard); + if expects_attempt { + timeout(Duration::from_secs(15), async { + while runtime.completed_agent_unload_attempts(thread.session.thread_id) == 0 { + sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("the timer must check the installed blocker after setup releases its pin"); + assert!(runtime.has_completed_agent_unload_timer(thread.session.thread_id)); + } else if !matches!(blocker, CompletedAgentUnloadBlocker::TreeShutdown) { + assert!(!runtime.has_completed_agent_unload_timer(thread.session.thread_id)); + } + assert!( + harness + .manager + .get_thread(thread.session.thread_id) + .await + .is_ok(), + "runtime still has an unload blocker" + ); + thread + .session + .services + .code_mode_service + .finish_cell_dispatch(&cell_id); + drop(elicitation); + if matches!(blocker, CompletedAgentUnloadBlocker::Elicitation) { + timeout(Duration::from_secs(15), async { + while harness + .manager + .get_thread(thread.session.thread_id) + .await + .is_ok() + { + sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("clearing an elicitation must let the bounded retry unload the child"); + } + let shutdown = harness + .manager + .request_agent_tree_shutdown(root_id) + .await + .expect("request tree shutdown"); + timeout(Duration::from_secs(5), shutdown.wait()) + .await + .expect("shutdown must not wait for the unload deadline") + .expect("tree shutdown succeeds"); +} + +#[test_case::test_case(false; "turn_input")] +#[test_case::test_case(true; "thread_settings")] +#[tokio::test] +async fn completed_agent_unload_excludes_direct_turn_input(settings: bool) { + let (home, mut config) = test_config().await; + let _ = config.features.enable(Feature::MultiAgentV2); + let harness = AgentControlHarness::new_with_config(home, config).await; + let (_, root) = harness.start_paginated_thread().await; + let control = root + .session + .services + .local_agent_runtime + .control(root.session.session_id()); + let child = spawn_v2_reload_test_child(&control, harness.config.clone(), &root, "worker").await; + let thread = harness + .manager + .get_thread(child.thread_id) + .await + .expect("loaded child"); + let eviction_guard = Arc::clone(&thread.residency_gate).write_owned().await; + let mut input = Box::pin(async { + if settings { + thread.update_thread_settings(codex_protocol::protocol::ThreadSettingsOverrides::default()).await + } else { + thread + .start_or_steer_turn(TurnInputRequest::user_input(text_input("followup"))) + .await + .map(|_| ()) + } + }); + assert!( + timeout(Duration::from_millis(50), &mut input) + .await + .is_err(), + "turn input must wait for the residency gate" + ); + drop(eviction_guard); + timeout(Duration::from_secs(10), input) + .await + .expect("input admission completes") + .expect("input accepted"); + harness + .manager + .shutdown_all_threads_bounded(Duration::from_secs(10)) + .await; +} + +#[tokio::test] +async fn completed_agent_unload_interrupt_during_teardown_preserves_reload() { + let (home, mut config) = test_config().await; + let _ = config.features.enable(Feature::MultiAgentV2); + let _ = config.features.enable(Feature::Sqlite); + let harness = AgentControlHarness::new_with_config(home, config).await; + let (_, root) = harness.start_paginated_thread().await; + let control = root + .session + .services + .local_agent_runtime + .control(root.session.session_id()); + let child = spawn_v2_reload_test_child(&control, harness.config.clone(), &root, "worker").await; + let thread = harness + .manager + .get_thread(child.thread_id) + .await + .expect("loaded child"); + drop( + persist_completed_agent_unload_fixture( + &thread, + &assistant_message( + "persisted before interrupt", + Some(MessagePhase::FinalAnswer), + ), + ) + .await, + ); + let eviction_guard = Arc::clone(&thread.residency_gate).write_owned().await; + // Teardown has stopped the receiver, but has not removed the registered runtime. + thread.shutdown_and_wait().await.expect("stop old runtime"); + let mut interrupt = Box::pin(control.interrupt_agent(child.thread_id)); + assert!( + timeout(Duration::from_millis(50), &mut interrupt) + .await + .is_err(), + "interrupt must wait for teardown instead of submitting to its closed receiver" + ); + harness + .manager + .remove_thread(&child.thread_id) + .await + .expect("remove old runtime"); + control.forget_v2_residency(child.thread_id); + drop(eviction_guard); + let result = timeout(Duration::from_secs(10), interrupt) + .await + .expect("interrupt completes"); + assert_matches!(result, Err(err) if matches!(err.details(), CodexErrorDetails::ThreadNotFound(_))); + assert!(control.get_agent_metadata(child.thread_id).is_some()); + control + .ensure_v2_agent_loaded(harness.config.clone(), child.thread_id, None) + .await + .map(drop) + .expect("interrupted child remains reloadable"); + let reloaded = harness + .manager + .get_thread(child.thread_id) + .await + .expect("reloaded child"); + assert!(!Arc::ptr_eq(&thread, &reloaded)); + harness + .manager + .shutdown_all_threads_bounded(Duration::from_secs(10)) + .await; +} + +#[test_case::test_case(false, false; "resume")] +#[test_case::test_case(true, false; "queue_only_caller_without_unload")] +#[test_case::test_case(true, true; "terminal_result_during_teardown")] +#[tokio::test] +async fn completed_agent_unload_resume_waits_for_eviction_and_reloads( + queue_only: bool, + terminal_result: bool, +) { + let (home, mut config) = test_config().await; + let _ = config.features.enable(Feature::MultiAgentV2); + let _ = config.features.enable(Feature::Sqlite); + let harness = AgentControlHarness::new_with_config(home, config).await; + let (_, root) = harness.start_paginated_thread().await; + let control = root + .session + .services + .local_agent_runtime + .control(root.session.session_id()); + let mut child_config = harness.config.clone(); + child_config + .multi_agent_v2 + .completed_agent_unload_delay_secs = Some(60); + let child = spawn_v2_reload_test_child(&control, child_config, &root, "worker").await; + let thread = harness + .manager + .get_thread(child.thread_id) + .await + .expect("loaded child"); + drop( + persist_completed_agent_unload_fixture( + &thread, + &assistant_message("persisted child", Some(MessagePhase::FinalAnswer)), + ) + .await, + ); + + // Hold the same exclusion gate as eviction while a resume observes the old runtime. + let eviction_guard = Arc::clone(&thread.residency_gate).write_owned().await; + let mut resume = Box::pin(async { + if terminal_result { + control + .send_inter_agent_communication( + child.thread_id, + InterAgentCommunication::new( + AgentPath::root(), + AgentPath::try_from("/root/worker").expect("agent path"), + Vec::new(), + "delivered once".into(), + false, + ), + AgentCommunicationContext::new( + AgentCommunicationKind::Result, + root.session.thread_id, + ), + Default::default(), + ) + .await + .map(|_| ()) + } else if queue_only { + control + .send(crate::SendRequest { + caller: root.session.thread_id, + target: crate::AgentTarget::Id(child.thread_id), + resume_config: harness.config.clone(), + input: AgentInput::Message { + message: AgentMessage::Plaintext("delivered once".into()), + mode: MessageDeliveryMode::QueueOnly, + }, + start_options: Default::default(), + }) + .await + .map(|_| ()) + } else { + control + .ensure_v2_agent_loaded(harness.config.clone(), child.thread_id, None) + .await + .map(drop) + } + }); + assert!( + timeout(Duration::from_millis(50), &mut resume) + .await + .is_err() + ); + thread + .shutdown_and_wait() + .await + .expect("finish resident shutdown"); + harness + .manager + .remove_thread(&child.thread_id) + .await + .expect("remove old runtime"); + control.forget_v2_residency(child.thread_id); + drop(eviction_guard); + + timeout(Duration::from_secs(10), resume) + .await + .expect("resume completes") + .expect("resume reloads after lost pin"); + if queue_only { + assert!(harness.manager.get_thread(child.thread_id).await.is_err()); + let expected = if terminal_result { + InterAgentCommunication::new( + AgentPath::root(), + AgentPath::try_from("/root/worker").expect("agent path"), + Vec::new(), + "delivered once".into(), + false, + ) + } else { + AgentMessage::Plaintext("delivered once".into()).into_communication( + AgentPath::root(), + AgentPath::try_from("/root/worker").expect("agent path"), + MessageDeliveryMode::QueueOnly, + ) + }; + assert_eq!(control.take_mailbox(child.thread_id), vec![expected]); + } else { + let reloaded = harness + .manager + .get_thread(child.thread_id) + .await + .expect("resume must publish a runtime"); + assert!(!Arc::ptr_eq(&thread, &reloaded)); + } + harness + .manager + .shutdown_all_threads_bounded(Duration::from_secs(10)) + .await; +} + +#[test_case::test_case(CompletedAgentUnloadTask::Completed, 0, false, false; "completed_zero_code_mode_off")] +#[test_case::test_case(CompletedAgentUnloadTask::Completed, 2, false, false; "completed_delayed_code_mode_off")] +#[test_case::test_case(CompletedAgentUnloadTask::Errored, 0, false, false; "errored")] +#[test_case::test_case(CompletedAgentUnloadTask::Interrupted, 0, false, false; "interrupted")] +#[test_case::test_case(CompletedAgentUnloadTask::Pending, 0, false, false; "explicit_interrupt")] +#[test_case::test_case(CompletedAgentUnloadTask::Completed, 0, true, false; "completed_zero_code_mode_on")] +#[test_case::test_case(CompletedAgentUnloadTask::Completed, 2, true, false; "completed_delayed_code_mode_on")] +#[test_case::test_case(CompletedAgentUnloadTask::Completed, 0, false, true; "archived_child_reloaded")] +#[tokio::test] +async fn completed_agent_unload_preserves_history_and_root( + outcome: CompletedAgentUnloadTask, + delay: i64, + code_mode: bool, + archived: bool, +) { + let (home, mut config) = test_config_with_cli_overrides(vec![( + "features.multi_agent_v2.completed_agent_unload_delay_secs".to_string(), + TomlValue::Integer(delay), + )]) + .await; + let _ = config.features.enable(Feature::MultiAgentV2); + let _ = config.features.enable(Feature::Sqlite); + if code_mode { + config.features.enable(Feature::CodeMode).unwrap(); + } else { + config.features.disable(Feature::CodeMode).unwrap(); + } + assert_eq!(config.features.enabled(Feature::CodeMode), code_mode); + config.multi_agent_v2.max_concurrent_threads_per_session = 8; + let harness = AgentControlHarness::new_with_config(home, config).await; + let root = harness + .manager + .start_thread(StartThreadOptions { + history_mode: Some(ThreadHistoryMode::Paginated), + ..StartThreadOptions::new(harness.config.clone()) + }) + .await + .expect("start root with local environment"); + let root_id = root.thread_id; + let root = root.thread; + let control = root + .session + .services + .local_agent_runtime + .control(root.session.session_id()); + let child = spawn_v2_reload_test_child(&control, harness.config.clone(), &root, "worker").await; + let mut thread = harness + .manager + .get_thread(child.thread_id) + .await + .expect("loaded child"); + let history_item = assistant_message("history before unload", Some(MessagePhase::FinalAnswer)); + let environments = thread.environment_selections().await; + assert!(!environments.is_empty()); + let mut setup_guard = persist_completed_agent_unload_fixture(&thread, &history_item).await; + if archived { + drop(setup_guard); + thread + .shutdown_and_wait() + .await + .expect("stop child before archive"); + harness + .manager + .remove_thread(&child.thread_id) + .await + .expect("remove stopped child"); + control.forget_v2_residency(child.thread_id); + let store = LocalThreadStore::new( + LocalThreadStoreConfig::from_config(&harness.config), + harness.state_db.clone(), + ); + store + .archive_thread(ArchiveThreadParams { + thread_id: child.thread_id, + }) + .await + .expect("archive saved V2 child"); + setup_guard = control + .ensure_v2_agent_loaded(harness.config.clone(), child.thread_id, None) + .await + .expect("reload archived V2 child through collaboration"); + let reloaded = harness + .manager + .get_thread(child.thread_id) + .await + .expect("archived child is loaded"); + assert!(!Arc::ptr_eq(&thread, &reloaded)); + assert_eq!(reloaded.multi_agent_version(), Some(MultiAgentVersion::V2)); + let stored = store + .read_thread(ReadThreadParams { + thread_id: child.thread_id, + include_archived: true, + include_history: false, + }) + .await + .expect("read archived child after reload"); + assert!( + stored.archived_at.is_some(), + "collaboration reload must keep archive state" + ); + thread = reloaded; + } + let turn = thread.session.new_default_turn().await; + thread.session.spawn_task(turn, Vec::new(), outcome).await; + drop(setup_guard); + if matches!(outcome, CompletedAgentUnloadTask::Pending) { + thread + .session + .abort_all_tasks(TurnAbortReason::Interrupted) + .await; + } + + let queued_mail = AgentMessage::Plaintext("queue-only mail survives unload".to_string()) + .into_communication( + AgentPath::root(), + AgentPath::try_from("/root/worker").expect("agent path"), + MessageDeliveryMode::QueueOnly, + ); + if delay > 0 { + timeout(Duration::from_secs(1), async { + while !matches!(thread.agent_status().await, AgentStatus::Completed(_)) + || thread.session.active_turn.lock().await.is_some() + || !thread + .session + .services + .local_agent_runtime + .has_completed_agent_unload_timer(child.thread_id) + { + sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("terminal result before unload deadline"); + control + .send(crate::SendRequest { + caller: root_id, + target: crate::AgentTarget::Id(child.thread_id), + resume_config: harness.config.clone(), + input: AgentInput::Message { + message: AgentMessage::Plaintext("queue-only mail survives unload".to_string()), + mode: MessageDeliveryMode::QueueOnly, + }, + start_options: Default::default(), + }) + .await + .expect("deliver queue-only mail after the unload timer is armed"); + assert!(harness.manager.get_thread(child.thread_id).await.is_ok()); + } + + timeout(Duration::from_secs(10), async { + while harness.manager.get_thread(child.thread_id).await.is_ok() { + sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("completed child should unload without capacity pressure"); + assert!(harness.manager.get_thread(root_id).await.is_ok()); + assert_matches!( + control + .inspect_agent(child.thread_id) + .await + .expect("known child"), + AgentInfo::Unloaded(_) + ); + if delay > 0 { + assert_eq!(control.take_mailbox(child.thread_id), vec![queued_mail]); + } + control + .ensure_v2_agent_loaded(harness.config.clone(), child.thread_id, None) + .await + .map(drop) + .expect("reload completed child"); + let reloaded = harness + .manager + .get_thread(child.thread_id) + .await + .expect("reloaded child"); + assert!(!Arc::ptr_eq(&thread, &reloaded)); + assert_eq!(reloaded.environment_selections().await, environments); + let history = reloaded.session.clone_history().await; + let actual = strip_response_item_ids(&strip_metadata_from_items( + &history.raw_items().cloned().collect::>(), + )); + let expected = strip_response_item_ids(&strip_metadata_from_items(&[history_item])); + assert!( + actual.contains(&expected[0]), + "reloaded history must retain the marker: actual={actual:?}, expected={expected:?}" + ); + harness + .manager + .shutdown_all_threads_bounded(Duration::from_secs(10)) + .await; +} + #[tokio::test] async fn ensure_v2_child_loaded_preserves_evicted_parent_authority() { check_v2_agent_reload(V2ReloadRoute::NestedParent).await; diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index a63dc8f4d231..a1dadbb6ec01 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/codex-rs/core/src/codex_thread.rs @@ -256,7 +256,7 @@ impl CodexThread { } pub async fn submit(&self, op: Op) -> CodexResult { - self.io.submit(op).await + self.submit_with_trace(op, None).await } /// Returns the session telemetry handle for thread-scoped production instrumentation. @@ -401,10 +401,22 @@ impl CodexThread { op: Op, trace: Option, ) -> CodexResult { + let residency_guard = if matches!(op, Op::Shutdown) { + None + } else { + self.session + .services + .local_agent_runtime + .pin_v2_residency(self) + .await? + }; self.io .submit_with_trace( - op, trace, /*parent_turn_id*/ None, /*root_turn_id*/ None, - /*residency_guard*/ None, + op, + trace, + /*parent_turn_id*/ None, + /*root_turn_id*/ None, + residency_guard, ) .await } @@ -475,6 +487,12 @@ impl CodexThread { &self, request: impl Into>, ) -> CodexResult { + let _residency_guard = self + .session + .services + .local_agent_runtime + .pin_v2_residency(self) + .await?; self.ensure_execution_capacity_for_turn_start(self.session.services.agent_control.as_ref()) .await?; let WithTurnExtensionData { @@ -588,6 +606,12 @@ impl CodexThread { request: impl Into>, mode: TurnInputMode, ) -> CodexResult { + let _residency_guard = self + .session + .services + .local_agent_runtime + .pin_v2_residency(self) + .await?; if !matches!(mode, TurnInputMode::Steer { .. }) { self.ensure_execution_capacity_for_turn_start( self.session.services.agent_control.as_ref(), @@ -686,6 +710,12 @@ impl CodexThread { &self, thread_settings: impl Into>, ) -> CodexResult<()> { + let _residency_guard = self + .session + .services + .local_agent_runtime + .pin_v2_residency(self) + .await?; let (reply, result) = oneshot::channel(); let WithTurnExtensionData { request: thread_settings, diff --git a/codex-rs/core/src/config/mod.rs b/codex-rs/core/src/config/mod.rs index c5ab36546e18..7cdcedaed8ef 100644 --- a/codex-rs/core/src/config/mod.rs +++ b/codex-rs/core/src/config/mod.rs @@ -1368,6 +1368,7 @@ impl Default for CurrentTimeReminderConfig { #[derive(Debug, Clone, PartialEq, Eq, Serialize)] pub struct MultiAgentV2Config { pub max_concurrent_threads_per_session: usize, + pub completed_agent_unload_delay_secs: Option, pub min_wait_timeout_ms: i64, pub max_wait_timeout_ms: i64, pub default_wait_timeout_ms: i64, @@ -1392,6 +1393,7 @@ impl MultiAgentV2Config { fn defaults_for_max_concurrency(max_concurrent_threads_per_session: usize) -> Self { Self { max_concurrent_threads_per_session, + completed_agent_unload_delay_secs: None, min_wait_timeout_ms: DEFAULT_MULTI_AGENT_V2_MIN_WAIT_TIMEOUT_MS, max_wait_timeout_ms: DEFAULT_MULTI_AGENT_V2_MAX_WAIT_TIMEOUT_MS, default_wait_timeout_ms: DEFAULT_MULTI_AGENT_V2_DEFAULT_WAIT_TIMEOUT_MS, @@ -2881,6 +2883,8 @@ fn resolve_multi_agent_v2_config(config_toml: &ConfigToml) -> MultiAgentV2Config MultiAgentV2Config { max_concurrent_threads_per_session, + completed_agent_unload_delay_secs: base + .and_then(|config| config.completed_agent_unload_delay_secs), min_wait_timeout_ms, max_wait_timeout_ms, default_wait_timeout_ms, diff --git a/codex-rs/core/src/session/handlers.rs b/codex-rs/core/src/session/handlers.rs index 91bf8de9bf91..1f85e32e9148 100644 --- a/codex-rs/core/src/session/handlers.rs +++ b/codex-rs/core/src/session/handlers.rs @@ -292,6 +292,9 @@ pub(super) async fn shutdown_session_runtime(sess: &Arc) { let mut state = sess.state.lock().await; // Stop admission and take the current warmup together so resume cannot replace it. state.shutting_down = true; + sess.services + .local_agent_runtime + .cancel_completed_agent_unload(sess.thread_id); state.take_session_startup_prewarm() }; if let Some(startup_prewarm) = startup_prewarm { diff --git a/codex-rs/core/src/session/turn_input.rs b/codex-rs/core/src/session/turn_input.rs index 74a7592de8ea..1cf2fe3a869a 100644 --- a/codex-rs/core/src/session/turn_input.rs +++ b/codex-rs/core/src/session/turn_input.rs @@ -696,11 +696,41 @@ impl Session { ) { let mut state = self.state.lock().await; state.last_started_turn_id = Some(turn_id.to_owned()); + self.services + .local_agent_runtime + .cancel_completed_agent_unload(self.thread_id); if let Some(attribution) = attribution { state.turn_attribution = Some(attribution); } } + #[expect( + clippy::await_holding_invalid_type, + reason = "The active-turn lock serializes scheduling with turn admission." + )] + pub(crate) async fn schedule_completed_agent_unload(self: &Arc, turn: &TurnContext) { + let Some(delay) = turn.config.multi_agent_v2.completed_agent_unload_delay_secs else { + return; + }; + if turn.multi_agent_version != codex_protocol::protocol::MultiAgentVersion::V2 + || !matches!(turn.session_source, SessionSource::SubAgent(_)) + || turn.config.ephemeral + { + return; + } + // Serialize the deadline with turn admission, including turns that have no task yet. + let active = self.active_turn.lock().await; + let state = self.state.lock().await; + if active.is_none() + && !state.shutting_down + && state.last_started_turn_id.as_deref() == Some(turn.sub_id.as_str()) + { + self.services + .local_agent_runtime + .schedule_completed_agent_unload(self, delay); + } + } + #[expect( clippy::await_holding_invalid_type, reason = "The admission lock keeps the previous turn stable while capturing its attribution." diff --git a/codex-rs/core/src/tasks/mod.rs b/codex-rs/core/src/tasks/mod.rs index 5612045b974b..4c002a1ece05 100644 --- a/codex-rs/core/src/tasks/mod.rs +++ b/codex-rs/core/src/tasks/mod.rs @@ -550,17 +550,23 @@ impl Session { pub async fn abort_all_tasks(self: &Arc, reason: TurnAbortReason) { let mut aborted_turn = false; let mut active_turn_to_clear = None; + let mut completed_turn = None; if let Some(mut active_turn) = self.take_active_turn(&reason).await { let task = active_turn.task.take(); aborted_turn = task.is_some(); if let Some(task) = task { - self.handle_task_abort( - task, - reason.clone(), - &active_turn.turn_state, - /*error*/ None, - ) - .await; + let turn = Arc::clone(&task.turn_context); + if self + .handle_task_abort( + task, + reason.clone(), + &active_turn.turn_state, + /*error*/ None, + ) + .await + { + completed_turn = Some(turn); + } } if aborted_turn { active_turn_to_clear = Some(active_turn); @@ -575,6 +581,13 @@ impl Session { if reason == TurnAbortReason::Interrupted && aborted_turn { self.maybe_start_turn_for_pending_work().await; } + if matches!( + reason, + TurnAbortReason::Interrupted | TurnAbortReason::BudgetLimited + ) && let Some(turn) = completed_turn + { + self.schedule_completed_agent_unload(&turn).await; + } } pub(crate) async fn abort_turn_if_active( @@ -616,9 +629,15 @@ impl Session { error: Option, ) { let task = active_turn.task.take(); + let mut completed_turn = None; if let Some(task) = task { - self.handle_task_abort(task, reason.clone(), &active_turn.turn_state, error) - .await; + let turn = Arc::clone(&task.turn_context); + if self + .handle_task_abort(task, reason.clone(), &active_turn.turn_state, error) + .await + { + completed_turn = Some(turn); + } } // Let interrupted tasks observe cancellation before dropping pending approvals, or an // in-flight approval wait can surface as a model-visible rejection before TurnAborted. @@ -627,6 +646,13 @@ impl Session { if reason == TurnAbortReason::Interrupted { self.maybe_start_turn_for_pending_work().await; } + if matches!( + reason, + TurnAbortReason::Interrupted | TurnAbortReason::BudgetLimited + ) && let Some(turn) = completed_turn + { + self.schedule_completed_agent_unload(&turn).await; + } } pub async fn on_task_finished( @@ -887,11 +913,22 @@ impl Session { } // Private reviewers already flushed the terminal event before delivering it. // Other buffering writers still need a barrier for the terminal event. - if !saved_guardian_completion && let Err(err) = self.flush_rollout().await { - warn!("failed to flush rollout after emitting terminal turn event: {err}"); - } + let flushed = if saved_guardian_completion { + true + } else { + match self.flush_rollout().await { + Ok(()) => true, + Err(err) => { + warn!("failed to flush rollout after emitting terminal turn event: {err}"); + false + } + } + }; if cleared_active_turn { self.maybe_start_turn_for_pending_work().await; + if flushed { + self.schedule_completed_agent_unload(&turn_context).await; + } } } @@ -933,12 +970,12 @@ impl Session { reason: TurnAbortReason, turn_state: &Mutex, error: Option, - ) { + ) -> bool { let sub_id = task.turn_context.sub_id.clone(); if task.cancellation_token.is_cancelled() { self.emit_turn_abort_lifecycle(reason, task.turn_context.extension_data.as_ref()) .await; - return; + return false; } trace!(task_kind = ?task.kind, sub_id, "aborting running task"); @@ -1029,8 +1066,12 @@ impl Session { self.send_event(task.turn_context.as_ref(), event).await; // Regular items were flushed before this terminal event was appended; buffering // thread writers may not flush it without another explicit barrier. - if let Err(err) = self.flush_rollout().await { - warn!("failed to flush rollout after emitting terminal turn event: {err}"); + match self.flush_rollout().await { + Ok(()) => true, + Err(err) => { + warn!("failed to flush rollout after emitting terminal turn event: {err}"); + false + } } } } diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index 2fb6da8d7157..6a7eca63ae0c 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -1798,12 +1798,12 @@ impl ThreadManagerState { root_turn_id: Option, ) -> CodexResult { let thread = self.get_thread(thread_id).await?; - let residency_guard = if matches!(op, Op::InterAgentCommunication { .. }) { + let residency_guard = if !matches!(op, Op::Shutdown) { thread .session .services .local_agent_runtime - .pin_v2_residency(self, &thread) + .pin_v2_residency(&thread) .await? } else { None diff --git a/codex-rs/core/src/tools/code_mode/mod.rs b/codex-rs/core/src/tools/code_mode/mod.rs index 2b241ffc7d03..0d2bae48b993 100644 --- a/codex-rs/core/src/tools/code_mode/mod.rs +++ b/codex-rs/core/src/tools/code_mode/mod.rs @@ -107,6 +107,10 @@ impl CodeModeService { self.availability.is_ok() } + pub(crate) fn has_active_cells(&self) -> bool { + !self.dispatch_broker.active_cell_ids().is_empty() + } + pub(crate) fn take_unavailable_warning(&self, tool_mode: ToolMode) -> Option { let error = self.availability.as_ref().err()?; let behavior = match tool_mode { diff --git a/codex-rs/features/src/feature_configs.rs b/codex-rs/features/src/feature_configs.rs index e3ae224609dd..81acdb5bdd7d 100644 --- a/codex-rs/features/src/feature_configs.rs +++ b/codex-rs/features/src/feature_configs.rs @@ -289,6 +289,9 @@ where pub struct MultiAgentV2ConfigToml { #[serde(skip_serializing_if = "Option::is_none")] pub enabled: Option, + /// Unload idle terminal children after this delay. Omit to disable automatic unload. + #[serde(skip_serializing_if = "Option::is_none")] + pub completed_agent_unload_delay_secs: Option, #[serde(skip_serializing_if = "Option::is_none")] #[schemars(range(min = 1))] pub max_concurrent_threads_per_session: Option, diff --git a/codex-rs/features/src/tests.rs b/codex-rs/features/src/tests.rs index a8a8fca058a7..bffc10ebaee4 100644 --- a/codex-rs/features/src/tests.rs +++ b/codex-rs/features/src/tests.rs @@ -740,6 +740,27 @@ multi_agent_v2 = true assert_eq!(features.multi_agent_v2, Some(FeatureToml::Enabled(true))); } +#[test] +fn multi_agent_v2_completed_agent_unload_delay_deserializes() { + for delay in [0, 2] { + let config: crate::MultiAgentV2ConfigToml = + toml::from_str(&format!("completed_agent_unload_delay_secs = {delay}")) + .expect("nonnegative unload delay should deserialize"); + let serialized = toml::Value::try_from(config).expect("serialize feature config"); + assert_eq!( + serialized["completed_agent_unload_delay_secs"].as_integer(), + Some(delay) + ); + } + let default = toml::Value::try_from(crate::MultiAgentV2ConfigToml::default()) + .expect("serialize default feature config"); + assert!(default.get("completed_agent_unload_delay_secs").is_none()); + assert!( + toml::from_str::("completed_agent_unload_delay_secs = -1") + .is_err() + ); +} + #[test] fn multi_agent_v2_feature_config_deserializes_table() { let features: FeaturesToml = toml::from_str( @@ -777,6 +798,7 @@ non_code_mode_only = true features.multi_agent_v2, Some(crate::FeatureToml::Config(crate::MultiAgentV2ConfigToml { enabled: Some(true), + completed_agent_unload_delay_secs: None, max_concurrent_threads_per_session: Some(4), min_wait_timeout_ms: Some(2500), max_wait_timeout_ms: Some(120000), diff --git a/docs/config.md b/docs/config.md index d35b0a87e3c2..97de86e3e0cf 100644 --- a/docs/config.md +++ b/docs/config.md @@ -13,3 +13,38 @@ Admins can set top-level `allow_managed_hooks_only = true` in still allowing managed hooks from requirements and managed config layers. This setting is only supported in `requirements.toml`; putting it in `config.toml` does not enable managed-hooks-only mode. + +## Unload completed multi-agent V2 children + +An optional delay releases idle child runtimes after completion, error, or +interruption. The root runtime remains loaded. Omit this key to disable automatic +unload. A value of `0` requests unload as soon as the terminal result is delivered +and saved history is ready for reload. + +```toml +[features.multi_agent_v2] +enabled = true +completed_agent_unload_delay_secs = 60 +``` + +Enable this setting only in a profile whose MCP servers and tools can restart +without losing required state. Unload stops the child session and its MCP clients. +Saved conversation history, queue-only mail, and environment selections survive. +Codex does not restore transient REPL variables, subscriptions, or server-side tasks. + +An active turn, durable sleep, trigger mail, background terminal, active code-mode +cell, pending elicitation, or active descendant blocks unload. Codex retries a blocked +unload once after five seconds, then keeps the runtime loaded. Resume or a new turn +cancels the pending deadline. A later terminal turn can set a new deadline. + +Follow-up input waits for an unload in progress, then reloads the child from saved +history. Codex does not replay input that the old session already accepted. +This setting does not change concurrency limits or capacity-driven eviction. + +The unload setting works with `features.code_mode` enabled or disabled. +An active code-mode cell blocks unload only when such a cell exists. + +`list_agents` lists loaded runtimes only. It omits children after unload. +Core's `get_status` returns `NotFound` for an unloaded runtime. +The child identity remains registered. Core's `inspect_agent` returns +`AgentInfo::Unloaded`, and follow-up input can reload the child.