openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-21 17:48:44 +00:00
Code Issues Packages Projects Releases Wiki Activity

Preserve agent state when unloading idle multi-agent v2 children (#52034)

Browse Source 
## What changed

- Route app-server idle unloading of multi-agent v2 children through local eviction that preserves unread mail and environment selections for reload.
- Hold residency pins through input submission so newly reloaded children cannot be evicted before receiving their work.
- Guard listener cleanup against replacement runtimes and subscriptions, and reject local eviction of host-controlled threads.

## Testing

Add coverage for unsubscribe-driven eviction followed by `followup_task`, verifying restored conversation context and exactly-once delivery of queued mail. Extend tests to check that residency pins block eviction and that rejected host eviction preserves the thread and mailbox.

GitOrigin-RevId: e36723730af683d2081e495e637f0e6d23548b8d
This commit is contained in:
jif
2026-10-08 09:39:33 +00:00
committed by copyberry
parent ea27864f99
commit d2a6dac20b
12 changed files with 493 additions and 124 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/agent/control.rs
+1
View File
@@ -73,6 +73,7 @@ mod interrupt;
mod legacy;
mod mailbox;
mod residency;
pub use residency::ThreadEvictionOutcome;
mod resume;
mod root_handoff;
mod runtime;
codex-rs/core/src/agent/control/api.rs
+11 -6
View File
@@ -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)
})
}
codex-rs/core/src/agent/control/residency.rs
+100 -82
View File
@@ -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<Self>,
thread: Arc<CodexThread>,
membership: &AgentTreeMembership,
) -> CodexResult<ThreadEvictionOutcome> {
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(&registered, &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<V2ResidencyState>,
@@ -65,16 +148,6 @@ impl LocalAgentControl {
.await
}
pub(super) async fn touch_loaded_v2_residency(
&self,
state: &Arc<ThreadManagerState>,
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(&registered, &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
codex-rs/core/src/agent/control/spawn.rs
+39 -30
View File
@@ -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<Arc<CodexThread>>,
) -> CodexResult<()> {
) -> CodexResult<Option<tokio::sync::OwnedRwLockReadGuard<()>>> {
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)
}
codex-rs/core/src/agent/control_tests.rs
+14 -2
View File
@@ -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 {