diff --git a/codex-rs/core/src/codex_delegate.rs b/codex-rs/core/src/codex_delegate.rs index f9304384c4..57d69f1b0a 100644 --- a/codex-rs/core/src/codex_delegate.rs +++ b/codex-rs/core/src/codex_delegate.rs @@ -87,6 +87,7 @@ pub(crate) async fn run_codex_thread_interactive( thread_extension_init.insert(isolation); thread_extension_init.insert(crate::guardian::GuardianReviewSessionHost::default()); let (session, io) = Session::spawn(SessionSpawnArgs { + startup: None, config, allow_provider_model_fallback: false, instructions, diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 4743503917..b35cc3905e 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -1,3 +1,5 @@ +pub(crate) mod startup; + use std::borrow::Cow; use std::collections::HashMap; use std::collections::HashSet; @@ -397,6 +399,7 @@ use codex_utils_stream_parser::ProposedPlanSegment; /// Runtime state lives on `Session`; keeping these endpoints separate lets all /// submission senders be dropped to terminate the session loop. The shared /// completion future observes that shutdown. +#[derive(Clone)] pub(crate) struct SessionIo { pub(crate) tx_sub: Sender, pub(crate) rx_event: Receiver, @@ -425,6 +428,7 @@ pub(crate) enum ForkPersistence { } pub(crate) struct SessionSpawnArgs { + pub(crate) startup: Option>, pub(crate) config: Config, pub(crate) allow_provider_model_fallback: bool, pub(crate) instructions: SessionInstructions, @@ -530,6 +534,7 @@ impl Session { async fn spawn_internal(args: SessionSpawnArgs) -> CodexResult<(Arc, SessionIo)> { let SessionSpawnArgs { + startup, config, allow_provider_model_fallback, instructions, @@ -822,6 +827,7 @@ impl Session { let (agent_status_tx, agent_status_rx) = watch::channel(AgentStatus::PendingInit); let session = Box::pin(Session::new( + startup.clone(), session_configuration, &environment_selections, config.clone(), @@ -889,6 +895,9 @@ impl Session { session_loop_termination: session_loop_termination_from_handle(session_loop_handle), }; + if let Some(startup) = startup { + let _ = startup.io.set(io.clone()); + } Ok((session, io)) } } diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index c2b62185ae..66c82df206 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -725,6 +725,7 @@ impl Session { #[instrument(name = "session_init", level = "info", skip_all)] #[allow(clippy::too_many_arguments)] pub(crate) async fn new( + startup: Option>, mut session_configuration: SessionConfiguration, environment_selections: &[TurnEnvironmentSelection], config: Arc, @@ -960,10 +961,15 @@ impl Session { // - load history metadata (skipped for subagents) let thread_persistence_fut = async { if config.ephemeral { - Ok::<_, anyhow::Error>(LiveThreadInitGuard::new(/*live_thread*/ None)) + Ok::<_, anyhow::Error>((None, LiveThreadInitGuard::new(/*live_thread*/ None))) } else { - let mut guard = LiveThreadInitGuard::default(); - match &initial_history { + let mut local_guard = LiveThreadInitGuard::default(); + let mut managed_guard = match &startup { + Some(startup) => Some(startup.persistence.lock().await), + None => None, + }; + let guard = managed_guard.as_deref_mut().unwrap_or(&mut local_guard); + let live_thread = match &initial_history { InitialHistory::New | InitialHistory::Cleared | InitialHistory::Forked(_) => { let params = CreateThreadParams { session_id, @@ -1009,7 +1015,7 @@ impl Session { Arc::clone(&thread_store), params, items, - &mut guard, + guard, ) .await? } else { @@ -1043,8 +1049,7 @@ impl Session { .await? } }; - // The completed result can wait in join! while the other startup work is pending. - Ok(guard) + Ok((Some(live_thread), local_guard)) } } .instrument(info_span!( @@ -1126,11 +1131,10 @@ impl Session { let (thread_persistence_result, state_db_ctx, (auth, mcp_projection)) = tokio::join!(thread_persistence_fut, state_db_fut, auth_and_mcp_fut); - let mut live_thread_init = thread_persistence_result.map_err(|e| { + let (live_thread, mut live_thread_init) = thread_persistence_result.map_err(|e| { error!("failed to initialize thread persistence: {e:#}"); e })?; - let live_thread = live_thread_init.as_ref().cloned(); let session_result: anyhow::Result> = async { let rollout_path = if let Some(live_thread) = live_thread.as_ref() { live_thread.local_rollout_path().await? @@ -1710,6 +1714,9 @@ impl Session { forked_from_ordinal_exclusive, next_internal_sub_id: AtomicU64::new(0), }); + if let Some(startup) = &startup { + let _ = startup.session.set(Arc::clone(&sess)); + } if let Some(network_policy_decider_session) = network_policy_decider_session { let mut guard = network_policy_decider_session.write().await; *guard = Arc::downgrade(&sess); diff --git a/codex-rs/core/src/session/startup.rs b/codex-rs/core/src/session/startup.rs new file mode 100644 index 0000000000..cc1a18064c --- /dev/null +++ b/codex-rs/core/src/session/startup.rs @@ -0,0 +1,37 @@ +//! Retains persistence acquisition and resources throughout managed startup. +//! The thread manager drops initialization, then joins acquisition and resource cleanup here. + +use std::sync::Arc; +use std::sync::OnceLock; + +use codex_protocol::protocol::Op; +use codex_thread_store::LiveThreadInitGuard; +use tokio::sync::Mutex; + +use super::SessionIo; +use super::session::Session; + +#[derive(Default)] +pub(crate) struct SessionStartup { + pub(crate) persistence: Mutex, + pub(crate) session: OnceLock>, + pub(crate) io: OnceLock, +} + +impl SessionStartup { + pub(crate) async fn cleanup(&self) { + if let Some(io) = self.io.get() { + // The session loop owns persistence now. Preserve its shutdown semantics even + // if registration or the caller's handoff was interrupted after the loop started. + self.persistence.lock().await.commit(); + let _ = io.submit(Op::Interrupt).await; + let _ = io.shutdown_and_wait().await; + } else { + if let Some(session) = self.session.get() { + super::handlers::shutdown_session_runtime(session).await; + } + let mut persistence = std::mem::take(&mut *self.persistence.lock().await); + persistence.discard().await; + } + } +} diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index de620e2265..df2914b17a 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -5787,6 +5787,7 @@ async fn session_new_fails_when_zsh_fork_enabled_without_packaged_zsh() { )); let environment_manager = Arc::new(EnvironmentManager::default_for_tests()); let result = Session::new( + /*startup*/ None, session_configuration, /*environment_selections*/ &[], Arc::clone(&config), @@ -6270,6 +6271,7 @@ async fn make_session_with_config_and_rx( let environment_manager = Arc::new(EnvironmentManager::default_for_tests()); let session = Session::new( + /*startup*/ None, session_configuration, &default_environments, Arc::clone(&config), @@ -6399,6 +6401,7 @@ async fn make_session_with_history_source_and_agent_control_and_rx( let environment_manager = Arc::new(EnvironmentManager::default_for_tests()); let session = Session::new( + /*startup*/ None, session_configuration, &default_environments, Arc::clone(&config), diff --git a/codex-rs/core/src/session/tests/guardian_tests.rs b/codex-rs/core/src/session/tests/guardian_tests.rs index e474b32621..1361983074 100644 --- a/codex-rs/core/src/session/tests/guardian_tests.rs +++ b/codex-rs/core/src/session/tests/guardian_tests.rs @@ -1260,6 +1260,7 @@ async fn guardian_subagent_does_not_inherit_parent_exec_policy_rules() { let mut thread_extension_init = codex_extension_api::ExtensionDataInit::default(); thread_extension_init.insert(codex_extension_api::SessionIsolation::Isolated); let (session, io) = Session::spawn(SessionSpawnArgs { + startup: None, config, allow_provider_model_fallback: false, instructions: Default::default(), diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index 3c6bfccf88..e5e6606936 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -1,3 +1,5 @@ +mod managed; + use crate::CodexAppsToolsCache; use crate::agent::AgentControl; use crate::agents_md_manager::SessionInstructions; @@ -291,6 +293,7 @@ impl StartThreadOptions { } struct ThreadSpawnRequest { + startup: Option>, options: StartThreadOptions, auth_manager: Arc, agent_control: AgentControl, @@ -310,6 +313,7 @@ impl ThreadSpawnRequest { agent_control: AgentControl, ) -> Self { Self { + startup: None, options, auth_manager, agent_control, @@ -987,7 +991,10 @@ impl ThreadManager { } pub async fn start_thread(&self, options: StartThreadOptions) -> CodexResult { - Box::pin(self.start_thread_inner(options, /*forked_from_thread_id*/ None)).await + Box::pin(self.start_thread_inner( + options, /*forked_from_thread_id*/ None, /*startup*/ None, + )) + .await } /// Starts a fresh internal session associated with an existing parent thread. @@ -1059,6 +1066,7 @@ impl ThreadManager { &self, mut options: StartThreadOptions, forked_from_thread_id: Option, + startup: Option>, ) -> CodexResult { let agent_control = self.agent_control_for_config(&options.config); let (resumed_session_source, resumed_thread_source) = options @@ -1075,6 +1083,7 @@ impl ThreadManager { let mut request = ThreadSpawnRequest::new(options, Arc::clone(&self.state.auth_manager), agent_control); request.forked_from_thread_id = forked_from_thread_id; + request.startup = startup; Box::pin(self.state.spawn_thread(request)).await } @@ -1111,7 +1120,7 @@ impl ThreadManager { inherited_multi_agent_version, ), ); - self.start_thread_inner(options, Some(forked_from_thread_id)) + self.start_thread_inner(options, Some(forked_from_thread_id), /*startup*/ None) .await } @@ -1924,6 +1933,7 @@ impl ThreadManagerState { /// Spawn a new thread with optional history and register it with the manager. async fn spawn_thread(&self, request: ThreadSpawnRequest) -> CodexResult { let ThreadSpawnRequest { + startup, options, auth_manager, agent_control, @@ -2094,6 +2104,7 @@ impl ThreadManagerState { codex_sandboxing::WindowsSandboxProxySettingsMode::Reconcile }; let (session, io) = Session::spawn(SessionSpawnArgs { + startup, config, allow_provider_model_fallback, instructions, diff --git a/codex-rs/core/src/thread_manager/managed.rs b/codex-rs/core/src/thread_manager/managed.rs new file mode 100644 index 0000000000..27fe8c9a27 --- /dev/null +++ b/codex-rs/core/src/thread_manager/managed.rs @@ -0,0 +1,109 @@ +//! Runs a thread for its owner's chosen lifetime, including cancellation during startup. +//! Cleanup remains tracked when the caller drops its startup future or stops waiting. + +use std::future::Future; +use std::sync::Arc; + +use codex_protocol::error::CodexErr; +use codex_protocol::error::Result as CodexResult; +use tokio::sync::oneshot; +use tokio_util::sync::CancellationToken; +use tokio_util::task::TaskTracker; + +use super::NewThread; +use super::StartThreadOptions; +use super::ThreadManager; +use crate::session::startup::SessionStartup; + +impl ThreadManager { + /// Starts a new isolated thread and owns it until `until` completes or the thread exits. + /// Dropping the startup future also cancels initialization. `tasks` tracks the + /// entire lifetime, including partial startup cleanup and deregistration. + /// Cleanup waits for any in-flight persistence acquisition before releasing its writer. + /// Once the session runtime starts, cleanup follows normal shutdown history rules, + /// even if the caller has not received the thread yet. + /// Resumed threads are excluded because they can refer to another owner's live runtime. + pub async fn start_thread_until( + self: &Arc, + options: StartThreadOptions, + until: impl Future + Send + 'static, + tasks: &TaskTracker, + ) -> CodexResult { + if matches!( + options.initial_history, + codex_history::InitialHistory::Resumed(_) + ) { + return Err(CodexErr::InvalidRequest( + "owned thread startup requires new or forked history".to_owned(), + )); + } + if options + .thread_extension_init + .get::() + .as_deref() + != Some(&codex_extension_api::SessionIsolation::Isolated) + { + return Err(CodexErr::InvalidRequest( + "owned thread startup requires explicit session isolation".to_owned(), + )); + } + let task = tasks.token(); + let manager = Arc::clone(self); + let state = Arc::downgrade(&self.state); + let (sender, receiver) = oneshot::channel(); + let abandoned = CancellationToken::new(); + let handoff = abandoned.clone().drop_guard(); + tokio::spawn(async move { + let _task = task; + let startup = Arc::new(SessionStartup::default()); + tokio::pin!(until); + // Drop initialization before cleanup; startup still owns any unfinished acquisition. + let result = { + let start = manager.start_thread_inner( + options, + /*forked_from_thread_id*/ None, + Some(Arc::clone(&startup)), + ); + tokio::select! { + biased; + _ = &mut until => Err(CodexErr::TurnAborted), + _ = abandoned.cancelled() => Err(CodexErr::TurnAborted), + result = start => result, + } + }; + // The lifetime task must not keep the manager and its parent threads alive. + drop(manager); + match result { + Ok(thread) => { + let running = Arc::clone(&thread.thread); + if sender.send(Ok(thread)).is_ok() { + tokio::select! { + _ = abandoned.cancelled() => {} + _ = &mut until => {} + _ = running.wait_until_terminated() => {} + } + } + } + Err(error) => { + let _ = sender.send(Err(error)); + } + } + startup.cleanup().await; + if let Some(state) = state.upgrade() + && let Some(session) = startup.session.get() + { + let mut threads = state.threads.write().await; + let thread_id = session.thread_id(); + if threads + .get(&thread_id) + .is_some_and(|thread| Arc::ptr_eq(&thread.session, session)) + { + threads.remove(&thread_id); + } + } + }); + let thread = receiver.await.map_err(|_| CodexErr::InternalAgentDied)??; + handoff.disarm(); + Ok(thread) + } +} diff --git a/codex-rs/core/tests/suite/managed_threads_tests.rs b/codex-rs/core/tests/suite/managed_threads_tests.rs new file mode 100644 index 0000000000..21e9388414 --- /dev/null +++ b/codex-rs/core/tests/suite/managed_threads_tests.rs @@ -0,0 +1,244 @@ +//! Exercises owned thread startup and shutdown through the real runtime and persistence store. + +use std::future::Future; +use std::sync::Arc; +use std::task::Context as TaskContext; +use std::task::Wake; +use std::time::Duration; + +use anyhow::Context; +use codex_core::StartThreadOptions; +use codex_core::TurnInputRequest; +use codex_extension_api::SessionIsolation; +use codex_protocol::protocol::EventMsg; +use codex_protocol::user_input::UserInput; +use codex_thread_store::ThreadStoreError; +use core_test_support::apps_test_server::AppsTestServer; +use core_test_support::responses; +use core_test_support::skip_if_no_network; +use core_test_support::test_codex::test_codex; +use core_test_support::wait_for_event; +use pretty_assertions::assert_eq; +use serde_json::json; +use tokio_util::sync::CancellationToken; +use tokio_util::task::TaskTracker; + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn dropping_startup_cleans_up_while_required_mcp_is_stalled() -> anyhow::Result<()> { + skip_if_no_network!(Ok(())); + + let server = responses::start_mock_server().await; + let fixture = test_codex().build_with_auto_env(&server).await?; + let mcp_server = responses::start_mock_server().await; + let (http_server, control) = AppsTestServer::mount_with_startup_control(&mcp_server).await?; + let release = control.hold_next_successful_initialize(); + let mut config = fixture.config.clone(); + let mut servers = config.mcp_servers.get().clone(); + servers.insert( + "stalled".to_owned(), + serde_json::from_value(json!({ + "url": format!("{}/api/codex/ps/mcp", http_server.chatgpt_base_url), + "http_headers": { "Authorization": "Bearer synthetic-test-token" }, + "required": true, + "startup_timeout_sec": 120, + }))?, + ); + config.mcp_servers.set(servers)?; + let thread_id = fixture.thread_manager.reserve_thread_id(); + let mut options = StartThreadOptions::new(config); + options.reserved_thread_id = Some(thread_id); + options.environments = Some(fixture.codex.environment_selections().await); + options + .thread_extension_init + .insert(SessionIsolation::Isolated); + let tasks = TaskTracker::new(); + let mut start = Box::pin(fixture.thread_manager.start_thread_until( + options, + std::future::pending(), + &tasks, + )); + + tokio::select! { + _ = &mut start => panic!("startup must wait for the required MCP server"), + result = tokio::time::timeout(Duration::from_secs(10), async { + while control.initialize_attempts() == 0 { + tokio::task::yield_now().await; + } + }) => result.context("required MCP initialization did not begin")?, + } + // Persistence is already open while initialization is blocked on the server. + fixture.thread_store.flush_thread(thread_id).await?; + drop(start); + tasks.close(); + tokio::time::timeout(Duration::from_secs(10), tasks.wait()) + .await + .context("dropping startup must finish cleanup without releasing MCP initialization")?; + assert_eq!( + fixture.thread_manager.list_thread_ids().await, + vec![fixture.session_configured.thread_id], + ); + assert!(matches!( + fixture.thread_store.flush_thread(thread_id).await, + Err(ThreadStoreError::ThreadNotFound { .. }) + )); + drop(release); + fixture.codex.shutdown_and_wait().await?; + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn dropping_startup_before_receiving_the_agent_finishes_cleanup() -> anyhow::Result<()> { + skip_if_no_network!(Ok(())); + + // Startup first suspends on its result channel. Observe that channel's wake without + // polling again, so the result stays sent but unconsumed until we drop the caller. + struct ResultReady(tokio::sync::Notify); + impl Wake for ResultReady { + fn wake(self: Arc) { + self.0.notify_one(); + } + } + let server = responses::start_mock_server().await; + let fixture = test_codex().build_with_auto_env(&server).await?; + let mut options = StartThreadOptions::new(fixture.config.clone()); + let thread_id = fixture.thread_manager.reserve_thread_id(); + options.reserved_thread_id = Some(thread_id); + options.environments = Some(fixture.codex.environment_selections().await); + options + .thread_extension_init + .insert(SessionIsolation::Isolated); + let tasks = TaskTracker::new(); + let mut start = Box::pin(fixture.thread_manager.start_thread_until( + options, + std::future::pending(), + &tasks, + )); + let ready = Arc::new(ResultReady(tokio::sync::Notify::new())); + let waker = Arc::clone(&ready).into(); + assert!( + start + .as_mut() + .poll(&mut TaskContext::from_waker(&waker)) + .is_pending() + ); + tokio::time::timeout(Duration::from_secs(10), ready.0.notified()).await?; + fixture.thread_manager.get_thread(thread_id).await?; + drop(start); + tasks.close(); + tokio::time::timeout(Duration::from_secs(10), tasks.wait()) + .await + .context("dropping an unconsumed result must also finish cleanup")?; + assert_eq!( + fixture.thread_manager.list_thread_ids().await, + vec![fixture.session_configured.thread_id] + ); + assert!(matches!( + fixture.thread_store.flush_thread(thread_id).await, + Err(ThreadStoreError::ThreadNotFound { .. }) + )); + fixture.codex.shutdown_and_wait().await?; + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn owner_cancellation_closes_agent_and_preserves_history_and_parent() -> anyhow::Result<()> { + skip_if_no_network!(Ok(())); + + let server = responses::start_mock_server().await; + let fixture = test_codex().build_with_auto_env(&server).await?; + let mut options = StartThreadOptions::new(fixture.config.clone()); + options.environments = Some(fixture.codex.environment_selections().await); + options + .thread_extension_init + .insert(SessionIsolation::Isolated); + let cancelled = CancellationToken::new(); + let tasks = TaskTracker::new(); + let agent = fixture + .thread_manager + .start_thread_until(options, cancelled.clone().cancelled_owned(), &tasks) + .await?; + // A failed startup with the same ID must leave the original agent's writer intact. + let mut duplicate = StartThreadOptions::new(fixture.config.clone()); + duplicate.reserved_thread_id = Some(agent.thread_id); + duplicate.environments = Some(fixture.codex.environment_selections().await); + duplicate + .thread_extension_init + .insert(SessionIsolation::Isolated); + let failed_tasks = TaskTracker::new(); + assert!( + fixture + .thread_manager + .start_thread_until(duplicate, std::future::pending(), &failed_tasks,) + .await + .is_err() + ); + failed_tasks.close(); + tokio::time::timeout(Duration::from_secs(10), failed_tasks.wait()).await?; + fixture.thread_store.flush_thread(agent.thread_id).await?; + + let response = responses::mount_sse_once( + &server, + responses::sse(vec![ + responses::ev_response_created("agent-response"), + responses::ev_assistant_message("agent-message", "agent finished"), + responses::ev_completed("agent-response"), + ]), + ) + .await; + agent + .thread + .start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text { + text: "run the owned agent".to_owned(), + text_elements: Vec::new(), + }])) + .await?; + wait_for_event(&agent.thread, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + assert!( + response + .single_request() + .body_contains_text("run the owned agent") + ); + + cancelled.cancel(); + tasks.close(); + tokio::time::timeout(Duration::from_secs(10), tasks.wait()) + .await + .context("the owner must be able to join shutdown and deregistration")?; + tokio::time::timeout(Duration::from_secs(1), agent.thread.wait_until_terminated()).await?; + assert_eq!( + fixture.thread_manager.list_thread_ids().await, + vec![fixture.session_configured.thread_id], + ); + assert!(matches!( + fixture.thread_store.flush_thread(agent.thread_id).await, + Err(ThreadStoreError::ThreadNotFound { .. }) + )); + let history = agent + .thread + .load_history(/*include_archived*/ false) + .await?; + assert!(history.items.iter().any(|item| matches!( + item, + codex_history::RolloutItem::EventMsg(EventMsg::TurnComplete(_)) + ))); + + let parent_response = responses::mount_sse_once( + &server, + responses::sse(vec![ + responses::ev_response_created("parent-response"), + responses::ev_completed("parent-response"), + ]), + ) + .await; + fixture.submit_turn("parent still works").await?; + assert!( + parent_response + .single_request() + .body_contains_text("parent still works") + ); + fixture.codex.shutdown_and_wait().await?; + Ok(()) +} diff --git a/codex-rs/core/tests/suite/mod.rs b/codex-rs/core/tests/suite/mod.rs index e757b5a8b9..78a793a46e 100644 --- a/codex-rs/core/tests/suite/mod.rs +++ b/codex-rs/core/tests/suite/mod.rs @@ -99,6 +99,8 @@ mod interrupt_hooks; mod items; mod json_result; mod live_cli; +#[path = "managed_threads_tests.rs"] +mod managed_threads; mod mcp_auth_elicitation; mod mcp_auth_refresh; mod mcp_ema_config;