diff --git a/codex-rs/app-server/src/request_processors/thread_lifecycle.rs b/codex-rs/app-server/src/request_processors/thread_lifecycle.rs index 34fff277d4..1c7535455a 100644 --- a/codex-rs/app-server/src/request_processors/thread_lifecycle.rs +++ b/codex-rs/app-server/src/request_processors/thread_lifecycle.rs @@ -1,8 +1,11 @@ use super::*; use crate::extensions::send_thread_warning; use codex_app_server_protocol::ThreadQueueChangedNotification; +use codex_core::ThreadEvictionOutcome; use codex_extension_api::ThreadIdleCause; use codex_protocol::config_types::MultiAgentMode; +use codex_protocol::protocol::MultiAgentVersion; +use codex_protocol::protocol::SessionSource; #[derive(Clone)] pub(super) struct ListenerTaskContext { @@ -93,6 +96,7 @@ impl UnloadingState { if let Some(target) = unloading_target && target <= Instant::now() { + tokio::task::yield_now().await; return true; } let unloading_sleep = async { @@ -222,6 +226,10 @@ pub(super) fn log_listener_attach_result( } } +#[expect( + clippy::await_holding_invalid_type, + reason = "listener removal must be serialized against new subscriptions" +)] pub(super) async fn ensure_listener_task_running( listener_task_context: ListenerTaskContext, conversation_id: ThreadId, @@ -251,6 +259,8 @@ pub(super) async fn ensure_listener_task_running( ) .await; let config_snapshot = conversation.config_snapshot().await; + let is_v2_child = conversation.multi_agent_version() == Some(MultiAgentVersion::V2) + && matches!(config_snapshot.session_source, SessionSource::SubAgent(_)); let thread_settings_baseline = thread_settings_from_config_snapshot(&config_snapshot); let (mut listener_command_rx, listener_generation) = { let mut thread_state = thread_state.lock().await; @@ -392,6 +402,60 @@ pub(super) async fn ensure_listener_task_running( } pending_thread_unloads.insert(conversation_id); } + if is_v2_child { + let eviction = thread_manager.try_evict_v2_thread(Arc::clone(&conversation)); + tokio::pin!(eviction); + let result = match tokio::time::timeout( + Duration::from_secs(/*secs*/ 10), &mut eviction, + ).await { + Ok(result) => result, + Err(_) => { + warn!(event.name = "codex.app_server.thread_shutdown_slow", "thread {conversation_id} shutdown is taking longer than expected; continuing to wait"); + eviction.await + } + }; + let outcome = match result { + Ok(ThreadEvictionOutcome::Busy) => { + pending_thread_unloads.lock().await.remove(&conversation_id); + unloading_state.note_thread_activity_observed(); + continue; + } + Ok(outcome @ (ThreadEvictionOutcome::Evicted | ThreadEvictionOutcome::NotCurrent)) => outcome, + Err(err) => { + warn!("failed to evict thread {conversation_id}: {err}"); + pending_thread_unloads.lock().await.remove(&conversation_id); + unloading_state.note_thread_activity_observed(); + continue; + } + }; + { + let mut pending = pending_thread_unloads.lock().await; + if !thread_state_manager.remove_unsubscribed_listener( + conversation_id, &thread_state, listener_generation, + ).await { + pending.remove(&conversation_id); + break; + } + } + // Release the status receiver before removing its watch entry. + drop(unloading_state); + outgoing_for_task + .cancel_requests_for_thread(conversation_id, /*error*/ None) + .await; + // A runtime removed by someone else may already have a replacement. + if outcome == ThreadEvictionOutcome::Evicted { + thread_watch_manager + .remove_thread(&conversation_id.to_string()) + .await; + outgoing_for_task + .send_server_notification(ServerNotification::ThreadClosed( + ThreadClosedNotification { thread_id: conversation_id.to_string() }, + )) + .await; + } + pending_thread_unloads.lock().await.remove(&conversation_id); + return; + } unload_thread_without_subscribers( thread_manager.clone(), outgoing_for_task.clone(), diff --git a/codex-rs/app-server/src/thread_state.rs b/codex-rs/app-server/src/thread_state.rs index a19c41b19f..7ef7e0f20e 100644 --- a/codex-rs/app-server/src/thread_state.rs +++ b/codex-rs/app-server/src/thread_state.rs @@ -501,6 +501,32 @@ impl ThreadStateManager { } } + /// Removes an unsubscribed listener's entry without clearing a replacement's state. + #[expect( + clippy::await_holding_invalid_type, + reason = "the listener generation must stay fixed while removing its map entry" + )] + pub(crate) async fn remove_unsubscribed_listener( + &self, + thread_id: ThreadId, + expected: &Arc>, + generation: u64, + ) -> bool { + let mut thread_state = expected.lock().await; + let mut state = self.state.lock().await; + if thread_state.listener_generation != generation + || !state.threads.get(&thread_id).is_some_and(|entry| { + Arc::ptr_eq(&entry.state, expected) && entry.connection_ids.is_empty() + }) + { + return false; + } + state.threads.remove(&thread_id); + self.unregister_listener_command_tx(thread_id); + thread_state.clear_listener(); + true + } + pub(crate) async fn clear_all_listeners(&self) { let thread_states = { let state = self.state.lock().await; diff --git a/codex-rs/app-server/tests/suite/v2/thread_unsubscribe.rs b/codex-rs/app-server/tests/suite/v2/thread_unsubscribe.rs index cddbccb317..289859b26a 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_unsubscribe.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_unsubscribe.rs @@ -2,6 +2,7 @@ use anyhow::Result; use app_test_support::MockResponsesConfig; use app_test_support::TestAppServer; use app_test_support::create_mock_responses_server_repeating_assistant; +use app_test_support::write_models_cache; use codex_app_server_protocol::ClientRequest; use codex_app_server_protocol::DynamicToolCallOutputContentItem; use codex_app_server_protocol::DynamicToolCallParams; @@ -28,6 +29,7 @@ use codex_app_server_protocol::ThreadUnsubscribeStatus; use codex_app_server_protocol::TurnStartParams; use codex_app_server_protocol::TurnStartResponse; use codex_app_server_protocol::UserInput as V2UserInput; +use codex_features::Feature; use core_test_support::responses; use core_test_support::streaming_sse::StreamingSseChunk; use core_test_support::streaming_sse::start_streaming_sse_server; @@ -522,3 +524,188 @@ async fn wait_for_dynamic_tool_started( } } } + +#[tokio::test] +async fn v2_child_idle_eviction_preserves_mail_and_followup_context() -> Result<()> { + let server = responses::start_mock_server().await; + let codex_home = TempDir::new()?; + MockResponsesConfig::new(&server.uri()) + .with_model("gpt-5.5") + .enable_feature(Feature::MultiAgentV2) + .with_root_config("thread_unload_delay_secs = 0") + .write(codex_home.path())?; + write_models_cache(codex_home.path()).await?; + let mut mcp = TestAppServer::builder() + .with_codex_home(codex_home.path()) + .build_initialized() + .await?; + let parent = mcp.start_thread(ThreadStartParams::default()).await?.thread; + let mut child_id = None; + + for (tool, args) in [ + ( + "spawn_agent", + json!({ + "task_name": "worker", "fork_turns": "none", "message": "Initial child task" + }), + ), + ( + "send_message", + json!({"target": "worker", "message": "Unread note for later"}), + ), + ( + "followup_task", + json!({"target": "worker", "message": "Continue the child task"}), + ), + ] { + let call_id = format!("idle-eviction-{tool}"); + let _call = responses::mount_sse_once_match( + &server, + wiremock::matchers::header("thread-id", parent.id.clone()), + responses::sse(vec![ + responses::ev_function_call_with_namespace( + &call_id, + "collaboration", + tool, + &args.to_string(), + ), + responses::ev_completed("parent-call"), + ]), + ) + .await; + let _parent_done = responses::mount_sse_once_match( + &server, + move |request: &wiremock::Request| { + String::from_utf8_lossy(&request.body).contains(&call_id) + }, + responses::sse(vec![responses::ev_completed("parent-done")]), + ) + .await; + if tool != "send_message" { + let parent_id = parent.id.clone(); + let _child = responses::mount_sse_once_match( + &server, + move |request: &wiremock::Request| { + request + .headers + .get("thread-id") + .is_some_and(|id| id != parent_id.as_str()) + }, + responses::sse(vec![ + responses::ev_assistant_message( + "child-result", + "Remember the initial child result", + ), + responses::ev_completed("child-done"), + ]), + ) + .await; + } + timeout( + DEFAULT_READ_TIMEOUT, + mcp.start_turn_and_wait_for_completion(TurnStartParams { + thread_id: parent.id.clone(), + input: vec![V2UserInput::Text { + text: format!("Run {tool}"), + text_elements: Vec::new(), + }], + ..Default::default() + }), + ) + .await??; + + if tool != "send_message" { + let completed = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_matching_notification( + "child turn completed", + |notification| { + notification.method == "turn/completed" + && notification + .params + .as_ref() + .is_some_and(|params| params["threadId"] != parent.id) + }, + ), + ) + .await??; + let params = completed.params.expect("child completion params"); + let completed_child = params["threadId"].as_str().expect("child ID").to_string(); + if let Some(child_id) = &child_id { + assert_eq!(&completed_child, child_id); + } else { + child_id = Some(completed_child); + } + } else { + // Mail is queued while subscribed. Unsubscribe lets the actual child listener + // evict it; the parent stays loaded so the follow-up uses the same agent tree. + let child_id = child_id.clone().expect("spawned child"); + let unsubscribed: ThreadUnsubscribeResponse = mcp + .request(|request_id| ClientRequest::ThreadUnsubscribe { + request_id, + params: ThreadUnsubscribeParams { + thread_id: child_id.clone(), + }, + }) + .await?; + assert_eq!(unsubscribed.status, ThreadUnsubscribeStatus::Unsubscribed); + let closed: ThreadClosedNotification = + timeout(DEFAULT_READ_TIMEOUT, mcp.read_notification("thread/closed")).await??; + assert_eq!( + closed, + ThreadClosedNotification { + thread_id: child_id.clone() + } + ); + let loaded: ThreadLoadedListResponse = mcp + .request(|request_id| ClientRequest::ThreadLoadedList { + request_id, + params: ThreadLoadedListParams::default(), + }) + .await?; + assert_eq!(loaded.data, vec![parent.id.clone()]); + assert_eq!( + responses::received_responses_requests(&server) + .await + .iter() + .filter( + |request| request.header("thread-id").as_deref() == Some(child_id.as_str()) + ) + .count(), + 1, + "queue-only mail must not start another child turn", + ); + } + } + let child_id = child_id.expect("spawned child"); + let requests = responses::received_responses_requests(&server).await; + let child_requests: Vec<_> = requests + .iter() + .filter(|request| request.header("thread-id").as_deref() == Some(child_id.as_str())) + .collect(); + assert_eq!(child_requests.len(), 2); + assert_eq!( + child_requests[1] + .inputs_of_type("message") + .into_iter() + .filter(|item| item["role"] == "assistant") + .collect::>(), + vec![json!({ + "type": "message", "role": "assistant", + "content": [{"type": "output_text", "text": "Remember the initial child result"}], + })], + ); + for message in ["Unread note for later", "Continue the child task"] { + assert_eq!( + child_requests[1] + .inputs_of_type("agent_message") + .iter() + .filter(|item| item.to_string().contains(message)) + .count(), + 1, + "resumed child must receive {message} exactly once", + ); + } + mcp.shutdown_gracefully().await?; + Ok(()) +} diff --git a/codex-rs/core/src/agent/control.rs b/codex-rs/core/src/agent/control.rs index a8d57e19ce..895cd26802 100644 --- a/codex-rs/core/src/agent/control.rs +++ b/codex-rs/core/src/agent/control.rs @@ -73,6 +73,7 @@ mod interrupt; mod legacy; mod mailbox; mod residency; +pub use residency::ThreadEvictionOutcome; mod resume; mod root_handoff; mod runtime; diff --git a/codex-rs/core/src/agent/control/api.rs b/codex-rs/core/src/agent/control/api.rs index ef0a699486..4087cde92b 100644 --- a/codex-rs/core/src/agent/control/api.rs +++ b/codex-rs/core/src/agent/control/api.rs @@ -110,10 +110,12 @@ impl AgentControl for LocalAgentControl { let (metadata, submission_id) = match input { AgentInput::UserInput(input) => { let receiver = self.get_agent_metadata(target); - if receiver.is_some() { + let _residency_pin = if receiver.is_some() { self.ensure_v2_agent_loaded(resume_config, target, /*parent*/ None) - .await?; - } + .await? + } else { + None + }; let submission_id = self.send_input(target, input, start_options).await?; (receiver.unwrap_or_default(), submission_id) } @@ -140,13 +142,15 @@ impl AgentControl for LocalAgentControl { // locally evicted recipients can retain mail without reloading. // Loaded recipients go straight to delivery, which rejects sends // racing an in-progress eviction when it acquires the residency pin. - if mode == MessageDeliveryMode::TriggerTurn + let _residency_pin = if mode == MessageDeliveryMode::TriggerTurn || (self.runtime.upgrade()?.get_thread(target).await.is_err() && self.runtime.registry.evicted_environments(target).is_none()) { self.ensure_v2_agent_loaded(resume_config, target, /*parent*/ None) - .await?; - } + .await? + } else { + None + }; let communication = message.into_communication(author, receiver_path, mode); let kind = match mode { MessageDeliveryMode::QueueOnly => { @@ -191,6 +195,7 @@ impl AgentControl for LocalAgentControl { let config = parent.session.get_config().await.as_ref().clone(); self.ensure_v2_agent_loaded(config, child, Some(parent)) .await + .map(drop) }) } diff --git a/codex-rs/core/src/agent/control/residency.rs b/codex-rs/core/src/agent/control/residency.rs index dd53eb6f2c..4755c221f8 100644 --- a/codex-rs/core/src/agent/control/residency.rs +++ b/codex-rs/core/src/agent/control/residency.rs @@ -1,3 +1,5 @@ +//! Accounts for resident MAv2 children and unloads exact idle runtimes while preserving their state. + use super::LocalAgentControl; use super::LocalAgentRuntime; use super::runtime::AgentTreeMembership; @@ -18,6 +20,87 @@ use std::sync::Arc; use std::sync::Mutex; use tracing::warn; +/// The result of trying to unload one exact MAv2 child runtime. +#[derive(Debug, PartialEq, Eq)] +pub enum ThreadEvictionOutcome { + /// Work or another teardown owns the runtime; retry later. + Busy, + /// This operation stopped and removed the runtime. + Evicted, + /// The runtime was removed or replaced; leave its replacement alone. + NotCurrent, +} + +impl ThreadManagerState { + pub(crate) async fn try_evict_v2_thread( + self: &Arc, + thread: Arc, + membership: &AgentTreeMembership, + ) -> CodexResult { + let Ok(residency_guard) = Arc::clone(&thread.residency_gate).try_write_owned() else { + return Ok(ThreadEvictionOutcome::Busy); + }; + let thread_id = thread.session.thread_id; + if !self + .get_thread(thread_id) + .await + .is_ok_and(|registered| Arc::ptr_eq(®istered, &thread)) + { + return Ok(ThreadEvictionOutcome::NotCurrent); + } + if !is_resident_candidate(&thread) || !is_unloadable(&thread).await { + return Ok(ThreadEvictionOutcome::Busy); + } + let manager = Arc::clone(self); + let teardown = membership + .clone() + .into_teardown_guard("resident_eviction", Some(thread_id)); + // Once shutdown is submitted, cancellation cannot release delivery or capacity early. + tokio::spawn(async move { + let _residency_guard = residency_guard; + thread.ensure_rollout_materialized().await; + if let Err(err) = thread.shutdown_and_wait().await { + teardown.record_shutdown_failure("stop_resident", CodexErrKind::from(&err).into()); + teardown.complete(); + return Err(err); + } + // Dispatch has stopped and the residency guard excludes senders. A tree + // shutdown deliberately discards its mailbox instead of retaining it. + let mail = thread + .session + .input_queue + .drain_mailbox() + .await + .into_iter() + .map(|mail| mail.communication) + .collect(); + let runtime = &thread.session.services.local_agent_runtime; + let _ = runtime.mailboxes.enqueue(thread_id, /*id*/ None, mail); + let environments = thread.environment_selections().await; + let mut threads = manager.threads.write().await; + if threads + .get(&thread_id) + .is_some_and(|registered| !Arc::ptr_eq(registered, &thread)) + { + teardown.complete(); + return Ok(ThreadEvictionOutcome::NotCurrent); + } + runtime + .registry + .save_evicted_environments(thread_id, environments); + let removed = threads.remove(&thread_id).is_some(); + runtime.residency.remove(thread_id); + teardown.complete(); + Ok(if removed { + ThreadEvictionOutcome::Evicted + } else { + ThreadEvictionOutcome::NotCurrent + }) + }) + .await? + } +} + #[derive(Default)] pub(super) struct V2Residency { state: Mutex, @@ -65,16 +148,6 @@ impl LocalAgentControl { .await } - pub(super) async fn touch_loaded_v2_residency( - &self, - state: &Arc, - thread_id: ThreadId, - ) { - if let Ok(thread) = state.get_thread(thread_id).await { - let _ = self.runtime.pin_v2_residency(state, &thread).await; - } - } - pub(super) fn forget_v2_residency(&self, thread_id: ThreadId) { self.runtime.residency.remove(thread_id); } @@ -168,80 +241,22 @@ impl V2Residency { } } }; - let Ok(residency_guard) = - Arc::clone(&candidate_thread.residency_gate).try_write_owned() - else { - continue; - }; - if !manager - .get_thread(candidate_thread_id) - .await - .is_ok_and(|registered| Arc::ptr_eq(®istered, &candidate_thread)) - || !is_unloadable(candidate_thread.as_ref()).await - { + // Capacity eviction keeps its existing policy of retaining unstarted workers. + if matches!( + candidate_thread.agent_status().await, + AgentStatus::PendingInit + ) { continue; } - // Once shutdown is submitted, cancellation cannot revoke it. The eviction task - // must keep delivery excluded and capacity reserved through registry removal. - let manager = Arc::clone(manager); - let residency = Arc::clone(self); - let teardown = membership - .clone() - .into_teardown_guard("resident_eviction", Some(candidate_thread_id)); - let eviction = tokio::spawn(async move { - let _residency_guard = residency_guard; - candidate_thread.ensure_rollout_materialized().await; - if let Err(err) = candidate_thread.shutdown_and_wait().await { - teardown - .record_shutdown_failure("stop_resident", CodexErrKind::from(&err).into()); - warn!( - "failed to shut down v2 resident thread before unloading {candidate_thread_id}: {err}" - ); - teardown.complete(); - return false; + match manager + .try_evict_v2_thread(candidate_thread, membership) + .await + { + Ok(ThreadEvictionOutcome::Evicted | ThreadEvictionOutcome::NotCurrent) => { + return true; } - // The submission loop has stopped and the residency guard excludes senders. - // Preserve unread queue-only mail before dropping the session that held it. - let mail = candidate_thread - .session - .input_queue - .drain_mailbox() - .await - .into_iter() - .map(|mail| mail.communication) - .collect(); - // A concurrent tree shutdown deliberately discards its unread mail. - let _ = candidate_thread - .session - .services - .local_agent_runtime - .mailboxes - .enqueue(candidate_thread_id, /*id*/ None, mail); - let environments = candidate_thread.environment_selections().await; - let mut threads = manager.threads.write().await; - if threads - .get(&candidate_thread_id) - .is_some_and(|registered| !Arc::ptr_eq(registered, &candidate_thread)) - { - teardown.complete(); - return false; - } - candidate_thread - .session - .services - .local_agent_runtime - .registry - .save_evicted_environments(candidate_thread_id, environments); - // Keep publication excluded until both entries have been removed. - threads.remove(&candidate_thread_id); - residency.remove(candidate_thread_id); - teardown.complete(); - true - }); - match eviction.await { - Ok(true) => return true, - Ok(false) => {} - Err(err) => warn!("v2 resident eviction task failed: {err}"), + Ok(ThreadEvictionOutcome::Busy) => {} + Err(err) => warn!("v2 resident eviction failed: {err}"), } } false @@ -298,7 +313,10 @@ pub(super) fn is_v2_resident_session_source(session_source: &SessionSource) -> b async fn is_unloadable(thread: &CodexThread) -> bool { matches!( thread.agent_status().await, - AgentStatus::Completed(_) | AgentStatus::Errored(_) | AgentStatus::Interrupted + AgentStatus::PendingInit + | AgentStatus::Completed(_) + | AgentStatus::Errored(_) + | AgentStatus::Interrupted ) && thread.session.active_turn.lock().await.is_none() && !thread.session.has_outstanding_durable_sleep() && !thread diff --git a/codex-rs/core/src/agent/control/spawn.rs b/codex-rs/core/src/agent/control/spawn.rs index 20e28a8bfe..c370b49211 100644 --- a/codex-rs/core/src/agent/control/spawn.rs +++ b/codex-rs/core/src/agent/control/spawn.rs @@ -320,12 +320,13 @@ impl LocalAgentControl { } /// A provided parent enables owner-validated reloads; `None` preserves sender-driven reloads. + /// Keep the returned residency pin until the caller has submitted its input. pub(crate) async fn ensure_v2_agent_loaded( &self, mut config: Config, thread_id: ThreadId, parent: Option>, - ) -> CodexResult<()> { + ) -> CodexResult>> { let membership = self.runtime.admit_start()?; let state = self.runtime.upgrade()?; let owner_thread_id = parent.as_ref().map(|parent| parent.session.thread_id); @@ -347,9 +348,10 @@ impl LocalAgentControl { ))); } } - if owner_thread_id.is_none() && state.get_thread(thread_id).await.is_ok() { - self.touch_loaded_v2_residency(&state, thread_id).await; - return Ok(()); + if owner_thread_id.is_none() + && let Ok(thread) = state.get_thread(thread_id).await + { + return self.runtime.pin_v2_residency(&state, &thread).await; } if self .runtime @@ -402,8 +404,7 @@ impl LocalAgentControl { } if let Ok(thread) = state.get_thread(thread_id).await { self.validate_loaded_v2_child(&thread, parent_thread_id)?; - self.touch_loaded_v2_residency(&state, thread_id).await; - return Ok(()); + return self.runtime.pin_v2_residency(&state, &thread).await; } } let parent = if let Some(parent) = parent { @@ -633,28 +634,37 @@ impl LocalAgentControl { if let Some(parent_thread_id) = owner_thread_id { self.validate_loaded_v2_child(&reloaded_thread.thread, parent_thread_id)?; } - self.runtime.registry.clear_evicted_environments(thread_id); - residency_slot.commit(reloaded_thread.thread_id); - // Register before listeners can forward events from the resumed thread. - if let Some((client_metadata, source)) = subagent_analytics { - let thread_config = reloaded_thread.thread.config_snapshot().await; - emit_subagent_session_started( - &reloaded_thread - .thread - .session - .services - .analytics_events_client, - client_metadata, - reloaded_thread.thread.session.session_id(), - reloaded_thread.thread_id, - thread_config.parent_thread_id, - thread_config, - source, - /*resumed_created_at*/ Some(stored_thread.created_at), - ); - } - state.notify_thread_created(reloaded_thread.thread_id); - Ok(()) + let runtime = self.runtime.clone(); + // 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?; + runtime.registry.clear_evicted_environments(thread_id); + residency_slot.commit(reloaded_thread.thread_id); + // Register before listeners can forward events from the resumed thread. + if let Some((client_metadata, source)) = subagent_analytics { + let thread_config = reloaded_thread.thread.config_snapshot().await; + emit_subagent_session_started( + &reloaded_thread + .thread + .session + .services + .analytics_events_client, + client_metadata, + reloaded_thread.thread.session.session_id(), + reloaded_thread.thread_id, + thread_config.parent_thread_id, + thread_config, + source, + /*resumed_created_at*/ Some(stored_thread.created_at), + ); + } + state.notify_thread_created(reloaded_thread.thread_id); + Ok(residency_pin) + }) + .await? } Err(err) => { if let Ok(thread) = state.get_thread(thread_id).await { @@ -663,8 +673,7 @@ impl LocalAgentControl { } self.runtime.registry.clear_evicted_environments(thread_id); drop(residency_slot); - self.touch_loaded_v2_residency(&state, thread_id).await; - return Ok(()); + return self.runtime.pin_v2_residency(&state, &thread).await; } Err(err) } diff --git a/codex-rs/core/src/agent/control_tests.rs b/codex-rs/core/src/agent/control_tests.rs index 801d4d9f4b..902f86dc39 100644 --- a/codex-rs/core/src/agent/control_tests.rs +++ b/codex-rs/core/src/agent/control_tests.rs @@ -1168,7 +1168,7 @@ async fn check_v2_agent_reload(route: V2ReloadRoute) { .expect("ollama provider should be configured"); let mut parent_turn = parent_thread.session.new_default_turn().await; - match route { + let residency_pin = match route { V2ReloadRoute::Sender => control .ensure_v2_agent_loaded(sender_config, spawned_agent.thread_id, /*parent*/ None) .await @@ -1208,13 +1208,24 @@ async fn check_v2_agent_reload(route: V2ReloadRoute) { .expect("known child should reload through its parent"); parent_turn = parent_thread.session.new_default_turn().await; assert!(harness.manager.get_thread(parent_thread_id).await.is_err()); + None } - } + }; let reloaded_child = harness .manager .get_thread(spawned_agent.thread_id) .await .expect("reloaded child thread should exist"); + if matches!(route, V2ReloadRoute::Sender) { + assert_eq!( + harness + .manager + .try_evict_v2_thread(Arc::clone(&reloaded_child)) + .await + .expect("attempt eviction before input submission"), + crate::ThreadEvictionOutcome::Busy, + ); + } let reloaded_instructions = reloaded_child.session.inherited_instructions().await; assert_eq!( (reloaded_instructions.user, reloaded_instructions.thread), @@ -1279,6 +1290,7 @@ async fn check_v2_agent_reload(route: V2ReloadRoute) { ) .await .expect("send_inter_agent_communication should succeed after reload"); + drop(residency_pin); let expected = ( spawned_agent.thread_id, Op::InterAgentCommunication { diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index a2b44e3f1f..56f64aeeb6 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -147,6 +147,7 @@ pub use codex_prompts as review_prompts; mod thread_manager; pub(crate) mod web_search; pub(crate) mod windows_sandbox_read_grants; +pub use agent::control::ThreadEvictionOutcome; pub use thread_manager::AgentTreeShutdown; pub use thread_manager::AgentTreeShutdownFailure; pub use thread_manager::AgentTreeShutdownFailureReason; diff --git a/codex-rs/core/src/thread_manager/shutdown.rs b/codex-rs/core/src/thread_manager/shutdown.rs index 11f89bf39a..cfa4af1ca6 100644 --- a/codex-rs/core/src/thread_manager/shutdown.rs +++ b/codex-rs/core/src/thread_manager/shutdown.rs @@ -1,8 +1,9 @@ -//! Fences and signals one local agent tree, then exposes its teardown result. +//! Unloads locally controlled idle MAv2 children and exposes tracked agent-tree teardown. use super::ThreadManager; use crate::agent::control::AgentTreeShutdownState; use codex_protocol::ThreadId; +use codex_protocol::error::CodexErr; use codex_protocol::error::Result as CodexResult; use codex_thread_store::ThreadStoreError; use std::fmt; @@ -126,6 +127,26 @@ impl AgentTreeShutdown { } impl ThreadManager { + /// Unloads this exact idle MAv2 child using the manager's local controller backend, + /// preserving mail and environments for reloading. Host controllers own their eviction + /// and recovery; managers configured with a host controller factory are rejected. + /// Dropping the future after shutdown starts does not cancel cleanup. + pub async fn try_evict_v2_thread( + &self, + thread: Arc, + ) -> CodexResult { + // Backend selection is fixed before the manager is shared or starts any threads. + if self.state.agent_control_factory.is_some() { + return Err(CodexErr::InvalidRequest( + "idle eviction requires the local agent controller".to_owned(), + )); + } + let Ok(membership) = thread.session.services.local_agent_runtime.admit_start() else { + return Ok(crate::ThreadEvictionOutcome::Busy); + }; + self.state.try_evict_v2_thread(thread, &membership).await + } + /// Requests shutdown of a loaded thread and every session sharing its local runtime. /// /// Returns after fencing new starts and signalling admitted sessions. Callers can wait on the diff --git a/codex-rs/core/src/tools/handlers/multi_agents/send_input.rs b/codex-rs/core/src/tools/handlers/multi_agents/send_input.rs index 3d2942f119..9998a17367 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents/send_input.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents/send_input.rs @@ -54,14 +54,16 @@ impl Handler { .local_agent_runtime .control(session.session_id()); let receiver_agent = local_agent_control.get_agent_metadata(receiver_thread_id); - if receiver_agent.is_some() { + let residency_pin = if receiver_agent.is_some() { let resume_config = build_agent_resume_config(turn.as_ref()) .map_err(FunctionCallError::RespondToModel)?; local_agent_control .ensure_v2_agent_loaded(resume_config, receiver_thread_id, /*parent*/ None) .await - .map_err(|err| collab_agent_error(receiver_thread_id, err))?; - } + .map_err(|err| collab_agent_error(receiver_thread_id, err))? + } else { + None + }; let receiver_agent = receiver_agent.unwrap_or_default(); if args.interrupt { session @@ -106,6 +108,7 @@ impl Handler { ) .await .map_err(|err| collab_agent_error(receiver_thread_id, err)); + drop(residency_pin); let status = local_agent_control.get_status(receiver_thread_id).await; session .emit_turn_item_completed( diff --git a/codex-rs/core/tests/suite/agent_control_tests.rs b/codex-rs/core/tests/suite/agent_control_tests.rs index 41162eaea0..165165c7bf 100644 --- a/codex-rs/core/tests/suite/agent_control_tests.rs +++ b/codex-rs/core/tests/suite/agent_control_tests.rs @@ -278,6 +278,28 @@ async fn host_threads_preserve_lineage_settings_and_resume_routing() -> anyhow:: (controller.identity(), Some(root_id)) ); assert_eq!(controller.service_tier(), Some("priority".to_string())); + let unread_mail = vec![codex_protocol::protocol::InterAgentCommunication::new( + AgentPath::root(), + AgentPath::root().join("worker").expect("valid child path"), + Vec::new(), + "Retain this in the host mailbox".into(), + /*trigger_turn*/ false, + )]; + *controller.mail.lock().expect("mail lock") = unread_mail.clone(); + let error = test + .thread_manager + .try_evict_v2_thread(Arc::clone(&child.thread)) + .await + .expect_err("local eviction must reject a host-controlled child"); + assert!(matches!( + error.details(), + CodexErrorDetails::InvalidRequest(_) + )); + assert!(Arc::ptr_eq( + &test.thread_manager.get_thread(child.thread_id).await?, + &child.thread, + )); + assert_eq!(controller.take_mailbox(child.thread_id), unread_mail); child.thread.ensure_rollout_materialized().await; child.thread.flush_rollout().await?; let saved = test