diff --git a/codex-rs/core/src/agent/control/residency.rs b/codex-rs/core/src/agent/control/residency.rs index 828c07c1b7..9e8d5ad331 100644 --- a/codex-rs/core/src/agent/control/residency.rs +++ b/codex-rs/core/src/agent/control/residency.rs @@ -65,13 +65,29 @@ impl LocalAgentControl { state: &Arc, thread_id: ThreadId, ) { - if let Ok(thread) = state.get_thread(thread_id).await - && is_resident_candidate(thread.as_ref()) - { - self.v2_residency.touch(thread_id); + if let Ok(thread) = state.get_thread(thread_id).await { + let _ = self.pin_v2_residency(state, &thread).await; } } + /// Pins and touches the registered runtime without waiting for unrelated eviction. + pub(crate) async fn pin_v2_residency( + &self, + state: &ThreadManagerState, + thread: &Arc, + ) -> 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?) { + return Err(CodexErr::ThreadNotFound(thread_id)); + } + self.v2_residency.touch(thread_id); + Ok(Some(guard)) + } + pub(super) fn forget_v2_residency(&self, thread_id: ThreadId) { self.v2_residency.remove(thread_id); } @@ -115,71 +131,85 @@ impl V2Residency { } async fn try_unload_one_resident( - &self, + self: &Arc, manager: &Arc, protected_thread_id: Option, ) -> bool { - let candidates_to_scan = self.resident_count(); - for _ in 0..candidates_to_scan { - let Some(candidate_thread_id) = self.pop_lru_candidate(protected_thread_id) else { - return false; - }; - let Some(candidate_thread) = manager - .get_thread(candidate_thread_id) - .await - .ok() - .filter(|thread| is_resident_candidate(thread)) - else { - continue; - }; - if !is_unloadable(candidate_thread.as_ref()).await { - self.touch(candidate_thread_id); - continue; - } - candidate_thread.ensure_rollout_materialized().await; - if let Err(err) = candidate_thread.shutdown_and_wait().await { - warn!( - "failed to shut down v2 resident thread before unloading {candidate_thread_id}: {err}" - ); - self.touch(candidate_thread_id); - continue; - } - let environments = candidate_thread.environment_selections().await; - candidate_thread - .session - .services - .agent_control - .state - .save_evicted_environments(candidate_thread_id, environments); - let _ = manager.remove_thread(&candidate_thread_id).await; - return true; - } - false - } - - fn resident_count(&self) -> usize { - self.state + // Keep shutting-down workers counted until removal. Each runtime's write guard + // excludes delivery and competing evictions without blocking unrelated workers. + let candidates = self + .state .lock() .unwrap_or_else(std::sync::PoisonError::into_inner) .residents - .len() - } - - fn pop_lru_candidate(&self, protected_thread_id: Option) -> Option { - let mut state = self - .state - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); - let candidates_to_scan = state.residents.len(); - for _ in 0..candidates_to_scan { - let candidate_thread_id = state.residents.pop_front()?; + .clone(); + for candidate_thread_id in candidates { if Some(candidate_thread_id) == protected_thread_id { - state.residents.push_back(candidate_thread_id); continue; } - return Some(candidate_thread_id); + let candidate_thread = { + let threads = manager.threads.read().await; + match threads.get(&candidate_thread_id) { + Some(thread) if is_resident_candidate(thread) => Arc::clone(thread), + Some(_) | None => { + // A reload cannot publish between the lookup and stale-entry removal. + self.remove(candidate_thread_id); + return true; + } + } + }; + 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 + { + 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 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 { + warn!( + "failed to shut down v2 resident thread before unloading {candidate_thread_id}: {err}" + ); + return false; + } + 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)) + { + return false; + } + candidate_thread + .session + .services + .agent_control + .state + .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); + true + }); + match eviction.await { + Ok(true) => return true, + Ok(false) => {} + Err(err) => warn!("v2 resident eviction task failed: {err}"), + } } - None + false } fn touch(&self, thread_id: ThreadId) { diff --git a/codex-rs/core/src/codex_delegate.rs b/codex-rs/core/src/codex_delegate.rs index 1aca5777c8..49d8ee4c18 100644 --- a/codex-rs/core/src/codex_delegate.rs +++ b/codex-rs/core/src/codex_delegate.rs @@ -1,5 +1,6 @@ use std::sync::Arc; +use crate::session::Submission; use async_channel::Receiver; use async_channel::Sender; use codex_async_utils::OrCancelExt; @@ -9,7 +10,6 @@ use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::Op; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::SubAgentSource; -use codex_protocol::protocol::Submission; use codex_protocol::protocol::ThreadSource; use codex_protocol::user_input::UserInput; use serde_json::Value; @@ -260,6 +260,7 @@ pub(crate) async fn run_codex_thread_one_shot( trace: None, parent_turn_id: None, root_turn_id: None, + residency_guard: None, }) .await; child_cancel.cancel(); diff --git a/codex-rs/core/src/codex_delegate_tests.rs b/codex-rs/core/src/codex_delegate_tests.rs index 7b145dc07d..d364ef184f 100644 --- a/codex-rs/core/src/codex_delegate_tests.rs +++ b/codex-rs/core/src/codex_delegate_tests.rs @@ -154,6 +154,7 @@ async fn forward_ops_preserves_submission_trace_context() { }), parent_turn_id: Some("parent-turn".to_string()), root_turn_id: Some("root-turn".to_string()), + residency_guard: None, }; tx_ops.send(submission).await.unwrap(); drop(tx_ops); diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index da0e02801f..29c37de904 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/codex-rs/core/src/codex_thread.rs @@ -6,6 +6,7 @@ use crate::elicitation::ElicitationRegistration; use crate::environment_selection::TurnEnvironmentState; use crate::session::SessionIo; use crate::session::SessionSettingsUpdate; +use crate::session::Submission; use crate::session::new_submission_id; use crate::session::session::Session; use crate::session::step_settings::StepSettingsUpdate; @@ -42,7 +43,6 @@ use codex_protocol::protocol::MultiAgentVersion; use codex_protocol::protocol::Op; use codex_protocol::protocol::SandboxPolicy; use codex_protocol::protocol::SessionSource; -use codex_protocol::protocol::Submission; use codex_protocol::protocol::ThreadHistoryMode; use codex_protocol::protocol::ThreadMemoryMode; use codex_protocol::protocol::ThreadSettingsSnapshot; @@ -72,6 +72,7 @@ use rmcp::model::ReadResourceRequestParams; use std::path::PathBuf; use std::sync::Arc; use tokio::sync::Mutex; +use tokio::sync::RwLock; use tokio::sync::oneshot; use tokio::sync::watch; use tokio_util::sync::CancellationToken; @@ -184,6 +185,8 @@ pub struct GuardianRootSnapshot { pub struct CodexThread { pub(crate) session: Arc, pub(crate) io: SessionIo, + // Queued agent mail owns a read guard until handled or dropped; eviction needs a write guard. + pub(crate) residency_gate: Arc>, // Registration source controls live access and lifecycle hooks. Managed Guardian // reviewers keep their existing subagent identity inside the session. pub(crate) session_source: SessionSource, @@ -220,6 +223,7 @@ impl CodexThread { Self { session, io, + residency_gate: Arc::default(), session_source, startup_metadata, rollout_path, @@ -327,6 +331,7 @@ impl CodexThread { self.io .submit_with_trace( op, trace, /*parent_turn_id*/ None, /*root_turn_id*/ None, + /*residency_guard*/ None, ) .await } @@ -461,6 +466,7 @@ impl CodexThread { trace: current_span_w3c_trace_context(), parent_turn_id: None, root_turn_id: None, + residency_guard: None, }) .await .map_err(|_| CodexErr::Fatal("thread session has stopped".to_string()))?; diff --git a/codex-rs/core/src/guardian/review_session_tests.rs b/codex-rs/core/src/guardian/review_session_tests.rs index ae8ecc30d3..455c5709e5 100644 --- a/codex-rs/core/src/guardian/review_session_tests.rs +++ b/codex-rs/core/src/guardian/review_session_tests.rs @@ -2,6 +2,7 @@ use super::*; use crate::agents_md_manager::AgentsMdManager; use crate::context::ContextualUserFragment; use crate::context_manager::ContextManager; +use crate::session::Submission; use codex_guardian_reviewer::ReviewerRequest; use codex_guardian_reviewer::guardian_output_contract_prompt; use codex_history::CodexHarnessMetadata; @@ -13,7 +14,6 @@ use codex_protocol::openai_models::AutoReviewMessages; use codex_protocol::openai_models::ModelMessages; use codex_protocol::protocol::AgentStatus; use codex_protocol::protocol::ErrorEvent; -use codex_protocol::protocol::Submission; use codex_protocol::protocol::TurnAbortReason; use codex_protocol::protocol::TurnAbortedEvent; use codex_protocol::protocol::TurnCompleteEvent; diff --git a/codex-rs/core/src/session/handlers.rs b/codex-rs/core/src/session/handlers.rs index dcea9a99d2..4ee789353c 100644 --- a/codex-rs/core/src/session/handlers.rs +++ b/codex-rs/core/src/session/handlers.rs @@ -1,3 +1,4 @@ +use super::Submission; use crate::realtime_conversation::handle_audio as handle_realtime_conversation_audio; use crate::realtime_conversation::handle_close as handle_realtime_conversation_close; use crate::realtime_conversation::handle_speech as handle_realtime_conversation_speech; @@ -5,7 +6,6 @@ use crate::realtime_conversation::handle_start as handle_realtime_conversation_s use crate::realtime_conversation::handle_text as handle_realtime_conversation_text; use async_channel::Receiver; use codex_otel::set_parent_from_w3c_trace_context; -use codex_protocol::protocol::Submission; use tracing::Instrument; use tracing::debug_span; use tracing::info_span; @@ -598,6 +598,7 @@ pub(super) async fn submission_loop( } .instrument(dispatch_span) .await; + drop(sub.residency_guard); if should_exit { shutdown_received = true; break; diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index bde763e8dc..e5b393704b 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -235,7 +235,9 @@ mod guardian_checkpoint; mod handlers; mod inject; mod reasoning_effort; +mod submission; pub(crate) use reasoning_effort::RequestEffortUsage; +pub(crate) use submission::Submission; mod input_queue; mod mcp; mod mcp_prewarm; @@ -380,7 +382,6 @@ use codex_protocol::protocol::SandboxPolicy; use codex_protocol::protocol::SessionConfiguredEvent; use codex_protocol::protocol::SessionNetworkProxyRuntime; use codex_protocol::protocol::StreamErrorEvent; -use codex_protocol::protocol::Submission; use codex_protocol::protocol::ThreadMemoryMode; use codex_protocol::protocol::TokenCountEvent; use codex_protocol::protocol::TokenUsage; @@ -962,6 +963,7 @@ impl SessionIo { pub(crate) async fn submit(&self, op: Op) -> CodexResult { self.submit_with_trace( op, /*trace*/ None, /*parent_turn_id*/ None, /*root_turn_id*/ None, + /*residency_guard*/ None, ) .await } @@ -972,6 +974,7 @@ impl SessionIo { trace: Option, parent_turn_id: Option, root_turn_id: Option, + residency_guard: Option>, ) -> CodexResult { let id = new_submission_id(); let sub = Submission { @@ -980,6 +983,7 @@ impl SessionIo { trace, parent_turn_id, root_turn_id, + residency_guard, }; self.submit_with_id(sub).await?; Ok(id) @@ -1019,6 +1023,7 @@ impl SessionIo { trace, parent_turn_id: None, root_turn_id: None, + residency_guard: None, }) .await?; reply_rx.await.unwrap_or(Err(CodexErr::InternalAgentDied)) @@ -1042,6 +1047,7 @@ impl SessionIo { trace, parent_turn_id: None, root_turn_id: None, + residency_guard: None, }) .await?; reply_rx.await.unwrap_or(Err(CodexErr::InternalAgentDied)) diff --git a/codex-rs/core/src/session/step_activation_tests.rs b/codex-rs/core/src/session/step_activation_tests.rs index 27b9a31d4b..cae3f507c7 100644 --- a/codex-rs/core/src/session/step_activation_tests.rs +++ b/codex-rs/core/src/session/step_activation_tests.rs @@ -1,6 +1,7 @@ use super::*; use crate::agents_md_manager::AgentsMdManager; use crate::agents_md_manager::SessionInstructions; +use crate::session::Submission; use crate::session::handlers::submission_loop; use crate::session::step_context::StepContext; use crate::session::step_settings::StepSettings; @@ -42,7 +43,6 @@ use codex_protocol::openai_models::ModelsResponse; use codex_protocol::openai_models::ReasoningEffort; use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::Op; -use codex_protocol::protocol::Submission; use codex_protocol::protocol::TurnAbortReason; use pretty_assertions::assert_eq; use std::collections::BTreeSet; @@ -457,6 +457,7 @@ fn settings_submission( trace: None, parent_turn_id: None, root_turn_id: None, + residency_guard: None, }, receiver, ) diff --git a/codex-rs/core/src/session/submission.rs b/codex-rs/core/src/session/submission.rs new file mode 100644 index 0000000000..9bd9162354 --- /dev/null +++ b/codex-rs/core/src/session/submission.rs @@ -0,0 +1,18 @@ +//! Core-owned queue metadata; forwarding transfers the residency guard with the operation. + +use codex_protocol::protocol::Op; +use codex_protocol::protocol::W3cTraceContext; +use tokio::sync::OwnedRwLockReadGuard; + +#[derive(Debug)] +#[expect(dead_code, reason = "Turn ancestry is retained in Debug diagnostics.")] +pub(crate) struct Submission { + pub id: String, + pub op: Op, + /// Optional W3C trace carrier propagated across async submission handoffs. + pub trace: Option, + pub parent_turn_id: Option, + pub root_turn_id: Option, + /// Keeps a V2 recipient resident until this submission is handled or dropped. + pub residency_guard: Option>, +} diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index bfa037d6d8..9a2c8d8d11 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -111,6 +111,7 @@ use crate::connectors::AppInfo; use crate::responses_metadata::CodexResponsesRequestKind; use crate::responses_metadata::CompactionTurnMetadata; use crate::rollout::recorder::RolloutRecorder; +use crate::session::Submission; use crate::state::ActiveTurn; use crate::state::TaskKind; use crate::tasks::SessionTask; @@ -177,7 +178,6 @@ use codex_protocol::protocol::RealtimeVoice; use codex_protocol::protocol::RealtimeVoicesList; use codex_protocol::protocol::SessionMeta; use codex_protocol::protocol::SessionMetaLine; -use codex_protocol::protocol::Submission; use codex_protocol::protocol::ThreadSettingsOverrides; use codex_protocol::protocol::TokenCountEvent; use codex_protocol::protocol::TokenUsage; @@ -7568,6 +7568,7 @@ async fn submit_with_trace_captures_current_span_trace_context() { /*trace*/ None, /*parent_turn_id*/ None, /*root_turn_id*/ None, + /*residency_guard*/ None, ) .await .expect("submit should succeed"); @@ -7640,6 +7641,7 @@ fn submission_dispatch_span_prefers_submission_trace_context() { op: Op::Interrupt, parent_turn_id: None, root_turn_id: None, + residency_guard: None, trace: Some(submission_trace), }) }); @@ -7668,6 +7670,7 @@ fn submission_dispatch_span_uses_debug_for_realtime_audio() { }), parent_turn_id: None, root_turn_id: None, + residency_guard: None, trace: None, }); @@ -8025,6 +8028,7 @@ async fn spawn_task_turn_span_inherits_dispatch_trace_context() { op: Op::Interrupt, parent_turn_id: None, root_turn_id: None, + residency_guard: None, trace: Some(submission_trace.clone()), }); let dispatch_span_id = dispatch_span.context().span().span_context().span_id(); diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index a0d7a8cdd0..fcf5c1b173 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -395,7 +395,8 @@ pub(crate) struct ResumeThreadWithHistoryOptions { /// `Arc` reference that can be downgraded to by `LocalAgentControl` while preventing every single /// function to require an `Arc<&Self>`. pub(crate) struct ThreadManagerState { - threads: Arc>>>, + // Eviction updates this registry and residency together, locking the registry first. + pub(crate) threads: Arc>>>, shared_thread_instructions: shared_instructions::SharedThreadInstructionsProviders, thread_created_tx: broadcast::Sender, thread_id_generator: ThreadIdGenerator, @@ -1690,6 +1691,16 @@ impl ThreadManagerState { root_turn_id: Option, ) -> CodexResult { let thread = self.get_thread(thread_id).await?; + let residency_guard = if matches!(op, Op::InterAgentCommunication { .. }) { + thread + .session + .services + .agent_control + .pin_v2_residency(self, &thread) + .await? + } else { + None + }; if let Some(ops_log) = &self.ops_log && let Ok(mut log) = ops_log.lock() && let Some(captured_op) = capture_test_op(&op) @@ -1698,7 +1709,13 @@ impl ThreadManagerState { } thread .io - .submit_with_trace(op, /*trace*/ None, parent_turn_id, root_turn_id) + .submit_with_trace( + op, + /*trace*/ None, + parent_turn_id, + root_turn_id, + residency_guard, + ) .await } diff --git a/codex-rs/core/tests/suite/agent_eviction_tests.rs b/codex-rs/core/tests/suite/agent_eviction_tests.rs new file mode 100644 index 0000000000..8e2d12ced9 --- /dev/null +++ b/codex-rs/core/tests/suite/agent_eviction_tests.rs @@ -0,0 +1,243 @@ +//! Accepted mail pins a worker through dispatch; cancelled eviction still owns teardown. + +use super::*; +use crate::suite::settings_commits::COMMITTED_MODEL; +use crate::suite::settings_commits::PauseAfterCommit; +use codex_core::TurnInputRequest; +use codex_core::config::Config; +use codex_extension_api::ExtensionFuture; +use codex_extension_api::ExtensionRegistryBuilder; +use codex_extension_api::ThreadLifecycleContributor; +use codex_extension_api::ThreadStopInput; +use codex_extension_api::ToolLifecycleContributor; +use codex_extension_api::ToolLifecycleFuture; +use codex_extension_api::ToolStartInput; +use codex_protocol::protocol::Op; +use codex_protocol::user_input::UserInput; +use core_test_support::ThreadIdle; +use pretty_assertions::assert_eq; +use std::sync::Arc; +use std::sync::Mutex; +use std::sync::mpsc; +use tokio::sync::Notify; +use tokio::sync::oneshot; +use tokio::time::timeout; + +#[derive(Default)] +struct ShutdownGate { + entered: Notify, + release: Notify, +} + +#[derive(Default)] +struct PauseShutdown { + delivery_started: Notify, +} + +impl ToolLifecycleContributor for PauseShutdown { + fn on_tool_start<'a>(&'a self, input: ToolStartInput<'a>) -> ToolLifecycleFuture<'a> { + Box::pin(async move { + if input.call_id == "late-mail" { + self.delivery_started.notify_one(); + } + }) + } +} + +impl ThreadLifecycleContributor for PauseShutdown { + fn on_thread_stop<'a>(&'a self, input: ThreadStopInput<'a>) -> ExtensionFuture<'a, ()> { + Box::pin(async move { + if let Some(gate) = input.thread_store.get::>() { + gate.entered.notify_one(); + gate.release.notified().await; + } + }) + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn queued_mail_and_cancelled_eviction_keep_worker_ownership() -> Result<()> { + const QUEUED_TASK: &str = "handle this accepted task after dispatch resumes"; + let server = start_mock_server().await; + mount_root_collaboration_call( + &server, + FIRST_PROMPT, + "first-call", + "spawn_agent", + json!({ "message": FIRST_TASK, "task_name": "first", "fork_turns": "none" }), + ) + .await; + mount_completed_worker(&server, FIRST_TASK, "first-call").await; + let delivery = mount_root_collaboration_call( + &server, + "queue a followup", + "queued-mail", + "followup_task", + json!({ "target": "first", "message": QUEUED_TASK }), + ) + .await; + let worker = mount_completed_worker(&server, QUEUED_TASK, "queued-mail").await; + let queued_message_count = || { + worker + .requests() + .iter() + .flat_map(|request| request.inputs_of_type("agent_message")) + .filter(|item| item.to_string().contains(QUEUED_TASK)) + .count() + }; + let blocked_spawn = mount_root_collaboration_call( + &server, + "reclaim the queued worker", + "blocked-spawn", + "spawn_agent", + json!({ "message": SECOND_TASK, "task_name": "replacement", "fork_turns": "none" }), + ) + .await; + + let (entered_tx, entered_rx) = oneshot::channel(); + let (release_tx, release_rx) = mpsc::channel(); + let mut extensions = ExtensionRegistryBuilder::::new(); + extensions.config_contributor(Arc::new(PauseAfterCommit { + gate: Mutex::new(Some((entered_tx, release_rx))), + })); + extensions.thread_lifecycle_contributor(Arc::new(ThreadIdle)); + let lifecycle = Arc::new(PauseShutdown::default()); + extensions.thread_lifecycle_contributor(lifecycle.clone()); + extensions.tool_lifecycle_contributor(lifecycle.clone()); + let test = test_codex() + .with_model("gpt-5.6-sol") + .with_extensions(Arc::new(extensions.build())) + .with_config(|config| { + config.features.enable(Feature::Collab).unwrap(); + config.features.enable(Feature::MultiAgentV2).unwrap(); + config.multi_agent_v2.max_concurrent_threads_per_session = 2; + }) + .build_with_auto_env(&server) + .await?; + let mut created = test.thread_manager.subscribe_thread_created(); + test.submit_turn(FIRST_PROMPT).await?; + ThreadIdle::wait(&test.codex).await; + let first_id = created.recv().await?; + let first = test.thread_manager.get_thread(first_id).await?; + wait_for_event(&first, |event| matches!(event, EventMsg::TurnComplete(_))).await; + ThreadIdle::wait(&first).await; + + // Reuse the settings fixture to pause the submission loop without starting a turn. + first + .submit(Op::ThreadSettings { + thread_settings: ThreadSettingsOverrides { + model: Some(COMMITTED_MODEL.to_string()), + ..Default::default() + }, + }) + .await?; + timeout(Duration::from_secs(10), entered_rx).await??; + test.submit_turn("queue a followup").await?; + ThreadIdle::wait(&test.codex).await; + assert_eq!( + delivery.function_call_output_text("queued-mail"), + Some(String::new()), + ); + assert_eq!(queued_message_count(), 0); + + // The sender has finished, but accepted mail still owns its place in the queue. + timeout( + Duration::from_secs(10), + test.submit_turn("reclaim the queued worker"), + ) + .await + .expect("queued mail must prevent eviction without waiting for dispatch")?; + ThreadIdle::wait(&test.codex).await; + assert_eq!( + blocked_spawn.function_call_output_text("blocked-spawn"), + Some("collab spawn failed: agent thread limit reached".to_string()), + ); + assert!(test.thread_manager.get_thread(first_id).await.is_ok()); + + release_tx.send(())?; + wait_for_event(&first, |event| matches!(event, EventMsg::TurnComplete(_))).await; + ThreadIdle::wait(&first).await; + assert_eq!(queued_message_count(), 1); + + // Reaching shutdown proves dispatch released the mail's residency guard. + let gate = Arc::new(ShutdownGate::default()); + first.thread_extension_data().insert(Arc::clone(&gate)); + + mount_sse_once_match( + &server, + |request: &wiremock::Request| body_contains(request, "start eviction"), + sse(vec![ + ev_response_created("eviction-response"), + ev_function_call_with_namespace( + "cancelled-spawn", + MULTI_AGENT_V2_NAMESPACE, + "spawn_agent", + &json!({ "message": SECOND_TASK, "task_name": "replacement", "fork_turns": "none" }).to_string(), + ), + ev_completed("eviction-response"), + ]), + ) + .await; + test.codex + .start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text { + text: "start eviction".into(), + text_elements: Vec::new(), + }])) + .await?; + timeout(Duration::from_secs(10), gate.entered.notified()).await?; + test.codex.submit(Op::Interrupt).await?; + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnAborted(_)) + }) + .await; + + let delivery = mount_root_collaboration_call( + &server, + "send during eviction", + "late-mail", + "send_message", + json!({ "target": "first", "message": "must not disappear behind shutdown" }), + ) + .await; + let send = test.submit_turn("send during eviction"); + tokio::pin!(send); + tokio::select! { + result = &mut send => panic!("delivery finished before shutdown was released: {result:?}"), + started = timeout(Duration::from_secs(10), lifecycle.delivery_started.notified()) => started?, + } + assert!( + timeout(Duration::from_secs(1), &mut send).await.is_err(), + "mail must not be accepted by the closing runtime after its eviction caller is cancelled", + ); + gate.release.notify_one(); + timeout(Duration::from_secs(10), send).await??; + assert_eq!( + delivery.function_call_output_text("late-mail"), + Some(format!("agent with id {first_id} not found")), + ); + assert!(test.thread_manager.get_thread(first_id).await.is_err()); + + // The cancelled caller leaves no reservation behind once teardown finishes. + mount_root_collaboration_call( + &server, + "retry replacement", + "replacement-call", + "spawn_agent", + json!({ "message": SECOND_TASK, "task_name": "replacement", "fork_turns": "none" }), + ) + .await; + mount_completed_worker(&server, SECOND_TASK, "replacement-call").await; + test.submit_turn("retry replacement").await?; + let replacement = test + .thread_manager + .get_thread(created.recv().await?) + .await?; + wait_for_event(&replacement, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + test.thread_manager + .shutdown_all_threads_bounded(Duration::from_secs(5)) + .await; + Ok(()) +} diff --git a/codex-rs/core/tests/suite/agent_execution.rs b/codex-rs/core/tests/suite/agent_execution.rs index 0096a0f52d..0eddad8b5b 100644 --- a/codex-rs/core/tests/suite/agent_execution.rs +++ b/codex-rs/core/tests/suite/agent_execution.rs @@ -56,7 +56,7 @@ async fn mount_root_collaboration_call( call_id: &'static str, tool_name: &'static str, arguments: serde_json::Value, -) { +) -> ResponseMock { let response_id = format!("resp-{call_id}"); mount_sse_once_match( server, @@ -84,7 +84,7 @@ async fn mount_root_collaboration_call( ev_completed(&completion_id), ]), ) - .await; + .await } async fn mount_completed_worker( @@ -563,3 +563,6 @@ async fn v2_residency_reload_preserves_inherited_environment_and_tools( Ok(()) } + +#[path = "agent_eviction_tests.rs"] +mod eviction_tests; diff --git a/codex-rs/core/tests/suite/settings_commits.rs b/codex-rs/core/tests/suite/settings_commits.rs index fd686b7435..9c8a9db591 100644 --- a/codex-rs/core/tests/suite/settings_commits.rs +++ b/codex-rs/core/tests/suite/settings_commits.rs @@ -35,7 +35,7 @@ use tokio::sync::oneshot; use tokio::time::timeout; const INITIAL_MODEL: &str = "gpt-5.4"; -const COMMITTED_MODEL: &str = "gpt-5.2"; +pub(super) const COMMITTED_MODEL: &str = "gpt-5.2"; const TIMEOUT: Duration = Duration::from_secs(10); fn assert_checkpoints(test: &TestCodex, expected: &[ThreadSettingsSnapshot]) -> Result<()> { @@ -138,8 +138,8 @@ async fn initial_plugin_ids_use_turn_context_without_extra_settings_checkpoints( Ok(()) } -struct PauseAfterCommit { - gate: Mutex, mpsc::Receiver<()>)>>, +pub(super) struct PauseAfterCommit { + pub(super) gate: Mutex, mpsc::Receiver<()>)>>, } impl ConfigContributor for PauseAfterCommit { diff --git a/codex-rs/protocol/src/protocol.rs b/codex-rs/protocol/src/protocol.rs index ad793df081..e76aa077ac 100644 --- a/codex-rs/protocol/src/protocol.rs +++ b/codex-rs/protocol/src/protocol.rs @@ -188,23 +188,6 @@ impl GitSha { } } -/// Submission Queue Entry - requests from user -#[derive(Debug)] -pub struct Submission { - /// Unique id for this Submission to correlate with Events - pub id: String, - /// Payload - pub op: Op, - /// Optional W3C trace carrier propagated across async submission handoffs. - pub trace: Option, - /// Core-provided ID of the parent turn that directly initiated this submission. - /// - /// This is only used for inter-agent communication. - pub parent_turn_id: Option, - /// Core-provided ID of the top-level turn that causally initiated this submission. - pub root_turn_id: Option, -} - #[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq, JsonSchema, TS)] pub struct W3cTraceContext { #[serde(default, skip_serializing_if = "Option::is_none")]