diff --git a/codex-rs/core/src/agent/control.rs b/codex-rs/core/src/agent/control.rs index d2f49ae3ecb7..9aab4f2bfda4 100644 --- a/codex-rs/core/src/agent/control.rs +++ b/codex-rs/core/src/agent/control.rs @@ -591,7 +591,6 @@ impl AgentControl { let parent_thread = state.get_thread(*parent_thread_id).await.ok()?; Some( parent_thread - .codex .session .services .turn_environments @@ -614,14 +613,12 @@ impl AgentControl { }; let parent_thread = state.get_thread(*parent_thread_id).await.ok()?; - let parent_config = parent_thread.codex.session.get_config().await; + let parent_config = parent_thread.session.get_config().await; if !crate::exec_policy::child_uses_parent_exec_policy(&parent_config, child_config) { return None; } - Some(Arc::clone( - &parent_thread.codex.session.services.exec_policy, - )) + Some(Arc::clone(&parent_thread.session.services.exec_policy)) } async fn open_thread_spawn_children( diff --git a/codex-rs/core/src/agent/control/execution.rs b/codex-rs/core/src/agent/control/execution.rs index d8559c1499cc..2c4fc76f6fd1 100644 --- a/codex-rs/core/src/agent/control/execution.rs +++ b/codex-rs/core/src/agent/control/execution.rs @@ -46,10 +46,10 @@ impl AgentControl { } let state = self.upgrade()?; let thread = state.get_thread(thread_id).await?; - if thread.codex.session.active_turn.lock().await.is_some() { + if thread.session.active_turn.lock().await.is_some() { return Ok(()); } - let config = thread.codex.session.get_config().await; + let config = thread.session.get_config().await; let multi_agent_version = thread .multi_agent_version() .unwrap_or_else(|| config.multi_agent_version_from_features()); diff --git a/codex-rs/core/src/agent/control/legacy.rs b/codex-rs/core/src/agent/control/legacy.rs index c6c6a6c382cf..9264833cbb57 100644 --- a/codex-rs/core/src/agent/control/legacy.rs +++ b/codex-rs/core/src/agent/control/legacy.rs @@ -6,8 +6,8 @@ impl AgentControl { pub(crate) async fn shutdown_live_agent(&self, agent_id: ThreadId) -> CodexResult { let state = self.upgrade()?; let result = if let Ok(thread) = state.get_thread(agent_id).await { - thread.codex.session.ensure_rollout_materialized().await; - thread.codex.session.flush_rollout().await?; + thread.session.ensure_rollout_materialized().await; + thread.session.flush_rollout().await?; let result = if matches!(thread.agent_status().await, AgentStatus::Shutdown) { Ok(String::new()) } else { diff --git a/codex-rs/core/src/agent/control/residency.rs b/codex-rs/core/src/agent/control/residency.rs index 901f2c93bd3c..0f069a8e7f61 100644 --- a/codex-rs/core/src/agent/control/residency.rs +++ b/codex-rs/core/src/agent/control/residency.rs @@ -226,13 +226,8 @@ async fn is_unloadable(thread: &CodexThread) -> bool { matches!( thread.agent_status().await, AgentStatus::Completed(_) | AgentStatus::Errored(_) | AgentStatus::Interrupted - ) && thread.codex.session.active_turn.lock().await.is_none() - && !thread - .codex - .session - .input_queue - .has_pending_mailbox_items() - .await + ) && thread.session.active_turn.lock().await.is_none() + && !thread.session.input_queue.has_pending_mailbox_items().await } #[cfg(test)] diff --git a/codex-rs/core/src/agent/control/residency_tests.rs b/codex-rs/core/src/agent/control/residency_tests.rs index be3210dfec2a..e3162d56c61c 100644 --- a/codex-rs/core/src/agent/control/residency_tests.rs +++ b/codex-rs/core/src/agent/control/residency_tests.rs @@ -151,9 +151,8 @@ async fn spawn_v2_subagent( } async fn mark_thread_completed(thread: &CodexThread) { - let turn = thread.codex.session.new_default_turn().await; + let turn = thread.session.new_default_turn().await; thread - .codex .session .send_event( turn.as_ref(), @@ -172,9 +171,8 @@ async fn mark_thread_completed(thread: &CodexThread) { } async fn mark_thread_interrupted(thread: &CodexThread) { - let turn = thread.codex.session.new_default_turn().await; + let turn = thread.session.new_default_turn().await; thread - .codex .session .send_event( turn.as_ref(), @@ -192,5 +190,5 @@ async fn mark_thread_interrupted(thread: &CodexThread) { async fn clear_active_turn(thread: &CodexThread) { // The fixture has no task runner to clear the turn after the terminal event. - *thread.codex.session.active_turn.lock().await = None; + *thread.session.active_turn.lock().await = None; } diff --git a/codex-rs/core/src/agent/control/spawn.rs b/codex-rs/core/src/agent/control/spawn.rs index 649c104f3015..df9a63be7f27 100644 --- a/codex-rs/core/src/agent/control/spawn.rs +++ b/codex-rs/core/src/agent/control/spawn.rs @@ -411,13 +411,7 @@ impl AgentControl { )) = notification_source.as_ref() { let client_metadata = match state.get_thread(*parent_thread_id).await { - Ok(parent_thread) => { - parent_thread - .codex - .session - .app_server_client_metadata() - .await - } + Ok(parent_thread) => parent_thread.session.app_server_client_metadata().await, Err(error) => { tracing::warn!( error = %error, @@ -430,17 +424,12 @@ impl AgentControl { } } }; - let thread_config = new_thread.thread.codex.thread_config_snapshot().await; + let thread_config = new_thread.thread.config_snapshot().await; let parent_thread_id = thread_config.parent_thread_id; emit_subagent_session_started( - &new_thread - .thread - .codex - .session - .services - .analytics_events_client, + &new_thread.thread.session.services.analytics_events_client, client_metadata, - new_thread.thread.codex.session.session_id(), + new_thread.thread.session.session_id(), new_thread.thread_id, parent_thread_id, thread_config, @@ -569,7 +558,7 @@ impl AgentControl { let multi_agent_v2_usage_hint_texts_to_filter: Vec = if let Some(parent_thread) = parent_thread.as_ref() { if multi_agent_version == MultiAgentVersion::V2 { - let parent_config = parent_thread.codex.session.get_config().await; + let parent_config = parent_thread.session.get_config().await; [ parent_config .multi_agent_v2 diff --git a/codex-rs/core/src/agent/control_tests.rs b/codex-rs/core/src/agent/control_tests.rs index d46ffdcb648e..13dd644c2d97 100644 --- a/codex-rs/core/src/agent/control_tests.rs +++ b/codex-rs/core/src/agent/control_tests.rs @@ -242,7 +242,6 @@ async fn wait_for_subagent_notification(parent_thread: &Arc) -> boo let wait = async { loop { let history_items = parent_thread - .codex .session .clone_history() .await @@ -263,9 +262,8 @@ async fn persist_thread_for_tree_resume(thread: &Arc, message: &str thread .inject_user_message_without_turn(message.to_string()) .await; - thread.codex.session.ensure_rollout_materialized().await; + thread.session.ensure_rollout_materialized().await; thread - .codex .session .flush_rollout() .await @@ -562,10 +560,9 @@ async fn send_inter_agent_communication_without_turn_queues_message_without_trig timeout(Duration::from_secs(5), async { loop { if thread - .codex .session .input_queue - .has_pending_input(&thread.codex.session.active_turn) + .has_pending_input(&thread.session.active_turn) .await { break; @@ -576,13 +573,7 @@ async fn send_inter_agent_communication_without_turn_queues_message_without_trig .await .expect("inter-agent communication should stay pending"); - let history_items = thread - .codex - .session - .clone_history() - .await - .raw_items() - .to_vec(); + let history_items = thread.session.clone_history().await.raw_items().to_vec(); assert!(!history_contains_assistant_inter_agent_communication( &history_items, &communication @@ -875,7 +866,6 @@ async fn spawn_agent_can_fork_parent_thread_history_with_sanitized_items() { .inject_user_message_without_turn("parent seed context".to_string()) .await; let expected_parent_seed = parent_thread - .codex .session .clone_history() .await @@ -883,7 +873,7 @@ async fn spawn_agent_can_fork_parent_thread_history_with_sanitized_items() { .first() .cloned() .expect("parent seed should be recorded"); - let turn_context = parent_thread.codex.session.new_default_turn().await; + let turn_context = parent_thread.session.new_default_turn().await; let parent_spawn_call_id = "spawn-call-history".to_string(); let trigger_message = InterAgentCommunication::new( AgentPath::root(), @@ -893,7 +883,6 @@ async fn spawn_agent_can_fork_parent_thread_history_with_sanitized_items() { /*trigger_turn*/ true, ); parent_thread - .codex .session .record_conversation_items( turn_context.as_ref(), @@ -933,19 +922,13 @@ async fn spawn_agent_can_fork_parent_thread_history_with_sanitized_items() { .await; let parent_reference_context_item = turn_context.to_turn_context_item(); parent_thread - .codex .session .persist_rollout_items(&[RolloutItem::TurnContext( parent_reference_context_item.clone(), )]) .await; + parent_thread.session.ensure_rollout_materialized().await; parent_thread - .codex - .session - .ensure_rollout_materialized() - .await; - parent_thread - .codex .session .flush_rollout() .await @@ -978,7 +961,7 @@ async fn spawn_agent_can_fork_parent_thread_history_with_sanitized_items() { .await .expect("child thread should be registered"); assert_ne!(child_thread_id, parent_thread_id); - let history = child_thread.codex.session.clone_history().await; + let history = child_thread.session.clone_history().await; let mut expected_final_answer = assistant_message("parent final answer", Some(MessagePhase::FinalAnswer)); expected_final_answer.set_turn_id_if_missing(&turn_context.sub_id); @@ -1001,7 +984,7 @@ async fn spawn_agent_can_fork_parent_thread_history_with_sanitized_items() { "full-history forked child history should replace parent usage hints with the child subagent hint while filtering non-final assistant/tool chatter" ); assert_eq!( - serde_json::to_value(child_thread.codex.session.reference_context_item().await) + serde_json::to_value(child_thread.session.reference_context_item().await) .expect("serialize child reference context item"), serde_json::to_value(Some(parent_reference_context_item)) .expect("serialize expected reference context item"), @@ -1037,7 +1020,7 @@ async fn spawn_agent_can_fork_parent_thread_history_with_sanitized_items() { .get_thread(no_hint_child_thread_id) .await .expect("no-hint child thread should be registered"); - let no_hint_history = no_hint_child_thread.codex.session.clone_history().await; + let no_hint_history = no_hint_child_thread.session.clone_history().await; assert!( !history_contains_text(no_hint_history.raw_items(), "Child subagent guidance."), "full-history forked child should not add empty subagent guidance" @@ -1101,7 +1084,7 @@ async fn spawn_agent_fork_strips_parent_usage_hints_from_compacted_history() { .expect("start parent thread"); let parent_thread_id = new_thread.thread_id; let parent_thread = new_thread.thread; - let turn_context = parent_thread.codex.session.new_default_turn().await; + let turn_context = parent_thread.session.new_default_turn().await; let parent_spawn_call_id = "spawn-call-compacted-usage-hints".to_string(); let replacement_history = vec![ ResponseItem::Message { @@ -1124,7 +1107,6 @@ async fn spawn_agent_fork_strips_parent_usage_hints_from_compacted_history() { }, ]; parent_thread - .codex .session .persist_rollout_items(&[ RolloutItem::Compacted(CompactedItem { @@ -1139,13 +1121,8 @@ async fn spawn_agent_fork_strips_parent_usage_hints_from_compacted_history() { RolloutItem::ResponseItem(spawn_agent_call(&parent_spawn_call_id)), ]) .await; + parent_thread.session.ensure_rollout_materialized().await; parent_thread - .codex - .session - .ensure_rollout_materialized() - .await; - parent_thread - .codex .session .flush_rollout() .await @@ -1178,7 +1155,7 @@ async fn spawn_agent_fork_strips_parent_usage_hints_from_compacted_history() { .get_thread(child_thread_id) .await .expect("child thread should be registered"); - let history = child_thread.codex.session.clone_history().await; + let history = child_thread.session.clone_history().await; assert!( history_contains_text(history.raw_items(), "compacted parent summary"), "forked child history should retain compacted non-hint content" @@ -1207,10 +1184,9 @@ async fn spawn_agent_fork_strips_parent_usage_hints_from_compacted_history() { async fn spawn_agent_fork_flushes_parent_rollout_before_loading_history() { let harness = AgentControlHarness::new().await; let (parent_thread_id, parent_thread) = harness.start_thread().await; - let turn_context = parent_thread.codex.session.new_default_turn().await; + let turn_context = parent_thread.session.new_default_turn().await; let parent_spawn_call_id = "spawn-call-unflushed".to_string(); parent_thread - .codex .session .record_conversation_items( turn_context.as_ref(), @@ -1248,7 +1224,7 @@ async fn spawn_agent_fork_flushes_parent_rollout_before_loading_history() { .get_thread(child_thread_id) .await .expect("child thread should be registered"); - let history = child_thread.codex.session.clone_history().await; + let history = child_thread.session.clone_history().await; assert!( history_contains_text(history.raw_items(), "unflushed final answer"), "forked child history should include unflushed assistant final answers after flushing the parent rollout" @@ -1280,9 +1256,8 @@ async fn spawn_agent_fork_last_n_turns_keeps_only_recent_turns() { "queued message".to_string(), /*trigger_turn*/ false, ); - let queued_turn_context = parent_thread.codex.session.new_default_turn().await; + let queued_turn_context = parent_thread.session.new_default_turn().await; parent_thread - .codex .session .record_conversation_items( queued_turn_context.as_ref(), @@ -1297,9 +1272,8 @@ async fn spawn_agent_fork_last_n_turns_keeps_only_recent_turns() { "triggered context".to_string(), /*trigger_turn*/ true, ); - let triggered_turn_context = parent_thread.codex.session.new_default_turn().await; + let triggered_turn_context = parent_thread.session.new_default_turn().await; parent_thread - .codex .session .record_conversation_items( triggered_turn_context.as_ref(), @@ -1309,10 +1283,9 @@ async fn spawn_agent_fork_last_n_turns_keeps_only_recent_turns() { parent_thread .inject_user_message_without_turn("current parent task".to_string()) .await; - let spawn_turn_context = parent_thread.codex.session.new_default_turn().await; + let spawn_turn_context = parent_thread.session.new_default_turn().await; let parent_spawn_call_id = "spawn-call-last-n".to_string(); parent_thread - .codex .session .record_conversation_items( spawn_turn_context.as_ref(), @@ -1320,19 +1293,13 @@ async fn spawn_agent_fork_last_n_turns_keeps_only_recent_turns() { ) .await; parent_thread - .codex .session .persist_rollout_items(&[RolloutItem::TurnContext( spawn_turn_context.to_turn_context_item(), )]) .await; + parent_thread.session.ensure_rollout_materialized().await; parent_thread - .codex - .session - .ensure_rollout_materialized() - .await; - parent_thread - .codex .session .flush_rollout() .await @@ -1365,7 +1332,7 @@ async fn spawn_agent_fork_last_n_turns_keeps_only_recent_turns() { .get_thread(child_thread_id) .await .expect("child thread should be registered"); - let history = child_thread.codex.session.clone_history().await; + let history = child_thread.session.clone_history().await; assert!( !history_contains_text(history.raw_items(), "old parent context"), @@ -1385,7 +1352,6 @@ async fn spawn_agent_fork_last_n_turns_keeps_only_recent_turns() { ); assert!( child_thread - .codex .session .reference_context_item() .await @@ -1436,9 +1402,8 @@ async fn spawn_agent_fork_last_n_turns_drops_parent_startup_prefix_when_under_li .expect("start parent thread"); let parent_thread_id = parent.thread_id; let parent_thread = parent.thread; - let startup_turn_context = parent_thread.codex.session.new_default_turn().await; + let startup_turn_context = parent_thread.session.new_default_turn().await; parent_thread - .codex .session .record_conversation_items( startup_turn_context.as_ref(), @@ -1456,23 +1421,17 @@ async fn spawn_agent_fork_last_n_turns_drops_parent_startup_prefix_when_under_li parent_thread .inject_user_message_without_turn("current parent task".to_string()) .await; - let spawn_turn_context = parent_thread.codex.session.new_default_turn().await; + let spawn_turn_context = parent_thread.session.new_default_turn().await; let parent_spawn_call_id = "spawn-call-last-n-under-limit".to_string(); parent_thread - .codex .session .record_conversation_items( spawn_turn_context.as_ref(), &[spawn_agent_call(&parent_spawn_call_id)], ) .await; + parent_thread.session.ensure_rollout_materialized().await; parent_thread - .codex - .session - .ensure_rollout_materialized() - .await; - parent_thread - .codex .session .flush_rollout() .await @@ -1505,7 +1464,7 @@ async fn spawn_agent_fork_last_n_turns_drops_parent_startup_prefix_when_under_li .get_thread(child_thread_id) .await .expect("child thread should be registered"); - let history = child_thread.codex.session.clone_history().await; + let history = child_thread.session.clone_history().await; assert!( history_contains_text(history.raw_items(), "current parent task"), "bounded fork should retain the requested recent parent turn" @@ -1515,16 +1474,11 @@ async fn spawn_agent_fork_last_n_turns_drops_parent_startup_prefix_when_under_li "bounded fork should drop parent startup context even when fewer turns exist than requested" ); assert_eq!( - &child_thread - .codex - .session - .services - .selected_capability_roots, + &child_thread.session.services.selected_capability_roots, &selected_capability_roots ); assert!( child_thread - .codex .session .reference_context_item() .await @@ -1564,10 +1518,9 @@ async fn spawn_agent_fork_last_n_turns_strips_parent_usage_hints() { parent_thread .inject_user_message_without_turn("parent task".to_string()) .await; - let turn_context = parent_thread.codex.session.new_default_turn().await; + let turn_context = parent_thread.session.new_default_turn().await; let parent_spawn_call_id = "spawn-call-last-n-usage-hints".to_string(); parent_thread - .codex .session .record_conversation_items( turn_context.as_ref(), @@ -1585,13 +1538,8 @@ async fn spawn_agent_fork_last_n_turns_strips_parent_usage_hints() { ], ) .await; + parent_thread.session.ensure_rollout_materialized().await; parent_thread - .codex - .session - .ensure_rollout_materialized() - .await; - parent_thread - .codex .session .flush_rollout() .await @@ -1624,7 +1572,7 @@ async fn spawn_agent_fork_last_n_turns_strips_parent_usage_hints() { .get_thread(child_thread_id) .await .expect("child thread should be registered"); - let history = child_thread.codex.session.clone_history().await; + let history = child_thread.session.clone_history().await; assert!( history_contains_text(history.raw_items(), "parent task"), "bounded fork should retain the requested recent parent turn" @@ -1961,9 +1909,8 @@ async fn multi_agent_v2_completion_ignores_dead_direct_parent() { .get_thread(tester_thread_id) .await .expect("tester thread should exist"); - let tester_turn = tester_thread.codex.session.new_default_turn().await; + let tester_turn = tester_thread.session.new_default_turn().await; tester_thread - .codex .session .send_event( tester_turn.as_ref(), @@ -1999,7 +1946,6 @@ async fn multi_agent_v2_completion_ignores_dead_direct_parent() { ); let root_history_items = root_thread - .codex .session .clone_history() .await @@ -2050,9 +1996,8 @@ async fn multi_agent_v2_completion_queues_message_for_direct_parent() { tester_path.to_string(), Some(tester_path.clone()), ); - let tester_turn = tester_thread.codex.session.new_default_turn().await; + let tester_turn = tester_thread.session.new_default_turn().await; tester_thread - .codex .session .send_event( tester_turn.as_ref(), @@ -2104,7 +2049,6 @@ async fn multi_agent_v2_completion_queues_message_for_direct_parent() { .expect("completion watcher should queue a direct-parent message"); let root_history_items = root_thread - .codex .session .clone_history() .await @@ -2144,7 +2088,6 @@ async fn completion_watcher_notifies_parent_when_child_is_missing() { assert_eq!(wait_for_subagent_notification(&parent_thread).await, true); let history_items = parent_thread - .codex .session .clone_history() .await @@ -2383,13 +2326,8 @@ async fn resume_thread_subagent_restores_stored_metadata() { .get_thread(child_thread_id) .await .expect("child thread should exist"); + child_thread.session.ensure_rollout_materialized().await; child_thread - .codex - .session - .ensure_rollout_materialized() - .await; - child_thread - .codex .session .flush_rollout() .await diff --git a/codex-rs/core/src/codex_delegate.rs b/codex-rs/core/src/codex_delegate.rs index d2911df4073f..4d2953f3e9c9 100644 --- a/codex-rs/core/src/codex_delegate.rs +++ b/codex-rs/core/src/codex_delegate.rs @@ -45,10 +45,9 @@ use crate::mcp_tool_call::build_guardian_mcp_tool_review_request; use crate::mcp_tool_call::is_mcp_tool_approval_question_id; use crate::mcp_tool_call::lookup_mcp_tool_metadata; use crate::mcp_tool_call::mcp_approvals_reviewer; -use crate::session::Codex; -use crate::session::CodexSpawnArgs; -use crate::session::CodexSpawnOk; use crate::session::SUBMISSION_CHANNEL_CAPACITY; +use crate::session::SessionIo; +use crate::session::SessionSpawnArgs; use crate::session::emit_subagent_session_started; use crate::session::session::Session; use crate::session::turn_context::TurnContext; @@ -67,11 +66,11 @@ struct PendingMcpInvocation { metadata: Option, } -/// Start an interactive sub-Codex thread and return IO channels. +/// Start an interactive sub-Codex thread and return its runtime and IO channels. /// -/// The returned `events_rx` yields non-approval events emitted by the sub-agent. +/// The returned IO yields non-approval events emitted by the sub-agent. /// Approval requests are handled via `parent_session` and are not surfaced. -/// The returned `ops_tx` allows the caller to submit additional `Op`s to the sub-agent. +/// Its submission channel accepts additional `Op`s for the sub-agent. #[allow(clippy::too_many_arguments)] pub(crate) async fn run_codex_thread_interactive( config: Config, @@ -82,7 +81,7 @@ pub(crate) async fn run_codex_thread_interactive( cancel_token: CancellationToken, subagent_source: SubAgentSource, initial_history: Option, -) -> Result { +) -> Result<(Arc, SessionIo), CodexErr> { let (tx_sub, rx_sub) = async_channel::bounded(SUBMISSION_CHANNEL_CAPACITY); let (tx_ops, rx_ops) = async_channel::bounded(SUBMISSION_CHANNEL_CAPACITY); let conversation_history = initial_history.unwrap_or(InitialHistory::New); @@ -91,7 +90,7 @@ pub(crate) async fn run_codex_thread_interactive( instructions: parent_session.user_instructions().await, warnings: Vec::new(), }; - let CodexSpawnOk { codex, .. } = Box::pin(Codex::spawn(CodexSpawnArgs { + let (session, io) = Box::pin(Session::spawn(SessionSpawnArgs { config, allow_provider_model_fallback: false, user_instructions, @@ -136,19 +135,17 @@ pub(crate) async fn run_codex_thread_interactive( })) .or_cancel(&cancel_token) .await??; - let thread_config = codex.thread_config_snapshot().await; + let thread_config = session.thread_config_snapshot().await; let client_metadata = parent_session.app_server_client_metadata().await; emit_subagent_session_started( &parent_session.services.analytics_events_client, client_metadata, - codex.session.session_id(), - codex.session.thread_id, + session.session_id(), + session.thread_id(), Some(parent_session.thread_id), thread_config, subagent_source, ); - let codex = Arc::new(codex); - // Use a child token so parent cancel cascades but we can scope it to this task let cancel_token_events = cancel_token.child_token(); let cancel_token_ops = cancel_token.child_token(); @@ -157,14 +154,23 @@ pub(crate) async fn run_codex_thread_interactive( // routing them to the parent session for decisions. let parent_session_clone = Arc::clone(&parent_session); let parent_ctx_clone = Arc::clone(&parent_ctx); - let codex_for_events = Arc::clone(&codex); + let session_for_events = Arc::clone(&session); + let io = Arc::new(io); // Cache the child call's MCP metadata at begin time. The later legacy // RequestUserInput approval event only carries a call_id and question metadata. let pending_mcp_invocations = Arc::new(Mutex::new(HashMap::::new())); + let caller_io = SessionIo { + tx_sub: tx_ops, + rx_event: rx_sub, + agent_status: io.agent_status.clone(), + session_loop_termination: io.session_loop_termination.clone(), + }; + let io_for_events = Arc::clone(&io); tokio::spawn(async move { forward_events( - codex_for_events, + io_for_events, + session_for_events, tx_sub, parent_session_clone, parent_ctx_clone, @@ -175,18 +181,11 @@ pub(crate) async fn run_codex_thread_interactive( }); // Forward ops from the caller to the sub-agent. - let codex_for_ops = Arc::clone(&codex); tokio::spawn(async move { - forward_ops(codex_for_ops, rx_ops, cancel_token_ops).await; + forward_ops(io, rx_ops, cancel_token_ops).await; }); - Ok(Codex { - tx_sub: tx_ops, - rx_event: rx_sub, - agent_status: codex.agent_status.clone(), - session: Arc::clone(&codex.session), - session_loop_termination: codex.session_loop_termination.clone(), - }) + Ok((session, caller_io)) } /// Convenience wrapper for one-time use with an initial prompt. @@ -204,11 +203,11 @@ pub(crate) async fn run_codex_thread_one_shot( subagent_source: SubAgentSource, final_output_json_schema: Option, initial_history: Option, -) -> Result { +) -> Result<(Arc, SessionIo), CodexErr> { // Use a child token so we can stop the delegate after completion without // requiring the caller to cancel the parent token. let child_cancel = cancel_token.child_token(); - let io = Box::pin(run_codex_thread_interactive( + let (session, io) = Box::pin(run_codex_thread_interactive( config, auth_manager, models_manager, @@ -234,7 +233,6 @@ pub(crate) async fn run_codex_thread_one_shot( let (tx_bridge, rx_bridge) = async_channel::bounded(SUBMISSION_CHANNEL_CAPACITY); let ops_tx = io.tx_sub.clone(); let agent_status = io.agent_status.clone(); - let session = Arc::clone(&io.session); let session_loop_termination = io.session_loop_termination.clone(); let io_for_bridge = io; tokio::spawn(async move { @@ -265,17 +263,20 @@ pub(crate) async fn run_codex_thread_one_shot( let (tx_closed, rx_closed) = async_channel::bounded(SUBMISSION_CHANNEL_CAPACITY); drop(rx_closed); - Ok(Codex { - rx_event: rx_bridge, - tx_sub: tx_closed, - agent_status, + Ok(( session, - session_loop_termination, - }) + SessionIo { + rx_event: rx_bridge, + tx_sub: tx_closed, + agent_status, + session_loop_termination, + }, + )) } async fn forward_events( - codex: Arc, + io: Arc, + session: Arc, tx_sub: Sender, parent_session: Arc, parent_ctx: Arc, @@ -288,10 +289,10 @@ async fn forward_events( loop { tokio::select! { _ = &mut cancelled => { - shutdown_delegate(&codex).await; + shutdown_delegate(&io).await; break; } - event = codex.next_event() => { + event = io.next_event() => { let event = match event { Ok(event) => event, Err(_) => break, @@ -311,7 +312,7 @@ async fn forward_events( } => { // Initiate approval via parent session; do not surface to consumer. handle_exec_approval( - &codex, + &io, id, &parent_session, &parent_ctx, @@ -325,7 +326,7 @@ async fn forward_events( msg: EventMsg::ApplyPatchApprovalRequest(event), } => { handle_patch_approval( - &codex, + &io, id, &parent_session, &parent_ctx, @@ -339,7 +340,7 @@ async fn forward_events( .. } => { handle_request_permissions( - &codex, + &io, &parent_session, &parent_ctx, event, @@ -352,7 +353,7 @@ async fn forward_events( msg: EventMsg::RequestUserInput(event), } => { handle_request_user_input( - &codex, + &io, id, &parent_session, &parent_ctx, @@ -370,11 +371,11 @@ async fn forward_events( // the child runtime at call begin is the one executing this invocation. // Cache its metadata now; the later approval event has only a call ID. let metadata = if let Some(turn_context) = - codex.session.turn_context_for_sub_id(&id).await + session.turn_context_for_sub_id(&id).await { - let mcp = codex.session.services.latest_mcp_runtime(); + let mcp = session.services.latest_mcp_runtime(); lookup_mcp_tool_metadata( - codex.session.as_ref(), + session.as_ref(), turn_context.as_ref(), mcp.manager(), &event.invocation.server, @@ -395,7 +396,7 @@ async fn forward_events( }, ); if !forward_event_or_shutdown( - &codex, + &io, &tx_sub, &cancel_token, Event { @@ -414,7 +415,7 @@ async fn forward_events( } => { pending_mcp_invocations.lock().await.remove(&event.call_id); if !forward_event_or_shutdown( - &codex, + &io, &tx_sub, &cancel_token, Event { @@ -428,7 +429,7 @@ async fn forward_events( } } other => { - if !forward_event_or_shutdown(&codex, &tx_sub, &cancel_token, other).await + if !forward_event_or_shutdown(&io, &tx_sub, &cancel_token, other).await { break; } @@ -440,12 +441,12 @@ async fn forward_events( } /// Ask the delegate to stop and drain its events so background sends do not hit a closed channel. -async fn shutdown_delegate(codex: &Codex) { - let _ = codex.submit(Op::Interrupt).await; - let _ = codex.submit(Op::Shutdown {}).await; +async fn shutdown_delegate(io: &SessionIo) { + let _ = io.submit(Op::Interrupt).await; + let _ = io.submit(Op::Shutdown {}).await; let _ = timeout(Duration::from_millis(500), async { - while let Ok(event) = codex.next_event().await { + while let Ok(event) = io.next_event().await { if matches!( event.msg, EventMsg::TurnAborted(_) | EventMsg::TurnComplete(_) @@ -458,7 +459,7 @@ async fn shutdown_delegate(codex: &Codex) { } async fn forward_event_or_shutdown( - codex: &Codex, + io: &SessionIo, tx_sub: &Sender, cancel_token: &CancellationToken, event: Event, @@ -466,7 +467,7 @@ async fn forward_event_or_shutdown( match tx_sub.send(event).or_cancel(cancel_token).await { Ok(Ok(())) => true, _ => { - shutdown_delegate(codex).await; + shutdown_delegate(io).await; false } } @@ -474,7 +475,7 @@ async fn forward_event_or_shutdown( /// Forward ops from a caller to a sub-agent, respecting cancellation. async fn forward_ops( - codex: Arc, + io: Arc, rx_ops: Receiver, cancel_token_ops: CancellationToken, ) { @@ -483,13 +484,13 @@ async fn forward_ops( Ok(Ok(submission)) => submission, Ok(Err(_)) | Err(_) => break, }; - let _ = codex.submit_with_id(submission).await; + let _ = io.submit_with_id(submission).await; } } /// Handle an ExecApprovalRequest by consulting the parent session and replying. async fn handle_exec_approval( - codex: &Codex, + io: &SessionIo, turn_id: String, parent_session: &Arc, parent_ctx: &Arc, @@ -563,7 +564,7 @@ async fn handle_exec_approval( .await }; - let _ = codex + let _ = io .submit(Op::ExecApproval { id: approval_id_for_op, turn_id: Some(turn_id), @@ -574,7 +575,7 @@ async fn handle_exec_approval( /// Handle an ApplyPatchApprovalRequest by consulting the parent session and replying. async fn handle_patch_approval( - codex: &Codex, + io: &SessionIo, _id: String, parent_session: &Arc, parent_ctx: &Arc, @@ -667,7 +668,7 @@ async fn handle_patch_approval( ) .await }; - let _ = codex + let _ = io .submit(Op::PatchApproval { id: approval_id, decision, @@ -676,7 +677,7 @@ async fn handle_patch_approval( } async fn handle_request_user_input( - codex: &Codex, + io: &SessionIo, id: String, parent_session: &Arc, parent_ctx: &Arc, @@ -693,7 +694,7 @@ async fn handle_request_user_input( ) .await { - let _ = codex.submit(Op::UserInputAnswer { id, response }).await; + let _ = io.submit(Op::UserInputAnswer { id, response }).await; return; } @@ -710,7 +711,7 @@ async fn handle_request_user_input( cancel_token, ) .await; - let _ = codex.submit(Op::UserInputAnswer { id, response }).await; + let _ = io.submit(Op::UserInputAnswer { id, response }).await; } /// Intercepts delegated legacy MCP approval prompts on the RequestUserInput @@ -793,7 +794,7 @@ async fn maybe_auto_review_mcp_request_user_input( } async fn handle_request_permissions( - codex: &Codex, + io: &SessionIo, parent_session: &Arc, parent_ctx: &Arc, event: RequestPermissionsEvent, @@ -819,7 +820,7 @@ async fn handle_request_permissions( let response = await_request_permissions_with_cancel(response_fut, parent_session, &call_id, cancel_token) .await; - let _ = codex + let _ = io .submit(Op::RequestPermissionsResponse { id: call_id, response, diff --git a/codex-rs/core/src/codex_delegate_tests.rs b/codex-rs/core/src/codex_delegate_tests.rs index 6a6997472ba6..5cfbf252d014 100644 --- a/codex-rs/core/src/codex_delegate_tests.rs +++ b/codex-rs/core/src/codex_delegate_tests.rs @@ -42,11 +42,10 @@ async fn forward_events_filters_private_events_before_blocked_send_is_cancelled( let (tx_sub, rx_sub) = bounded(SUBMISSION_CHANNEL_CAPACITY); let (_agent_status_tx, agent_status) = watch::channel(AgentStatus::PendingInit); let (session, ctx, _rx_evt) = crate::session::tests::make_session_and_context_with_rx().await; - let codex = Arc::new(Codex { + let io = Arc::new(SessionIo { tx_sub, rx_event: rx_events, agent_status, - session: Arc::clone(&session), session_loop_termination: completed_session_loop_termination(), }); @@ -67,7 +66,8 @@ async fn forward_events_filters_private_events_before_blocked_send_is_cancelled( let cancel = CancellationToken::new(); let forward = tokio::spawn(forward_events( - Arc::clone(&codex), + Arc::clone(&io), + Arc::clone(&session), tx_out.clone(), session, ctx, @@ -141,17 +141,15 @@ async fn forward_ops_preserves_submission_trace_context() { let (tx_sub, rx_sub) = bounded(SUBMISSION_CHANNEL_CAPACITY); let (_tx_events, rx_events) = bounded(SUBMISSION_CHANNEL_CAPACITY); let (_agent_status_tx, agent_status) = watch::channel(AgentStatus::PendingInit); - let (session, _ctx, _rx_evt) = crate::session::tests::make_session_and_context_with_rx().await; - let codex = Arc::new(Codex { + let io = Arc::new(SessionIo { tx_sub, rx_event: rx_events, agent_status, - session, session_loop_termination: completed_session_loop_termination(), }); let (tx_ops, rx_ops) = bounded(1); let cancel = CancellationToken::new(); - let forward = tokio::spawn(forward_ops(Arc::clone(&codex), rx_ops, cancel)); + let forward = tokio::spawn(forward_ops(Arc::clone(&io), rx_ops, cancel)); let submission = Submission { id: "sub-1".to_string(), @@ -218,11 +216,10 @@ async fn handle_request_permissions_uses_tool_call_id_for_round_trip() { let (tx_sub, rx_sub) = bounded(SUBMISSION_CHANNEL_CAPACITY); let (_tx_events, rx_events_child) = bounded(SUBMISSION_CHANNEL_CAPACITY); let (_agent_status_tx, agent_status) = watch::channel(AgentStatus::PendingInit); - let codex = Arc::new(Codex { + let io = Arc::new(SessionIo { tx_sub, rx_event: rx_events_child, agent_status, - session: Arc::clone(&parent_session), session_loop_termination: completed_session_loop_termination(), }); @@ -244,13 +241,13 @@ async fn handle_request_permissions_uses_tool_call_id_for_round_trip() { let request_cwd = delegated_cwd.clone(); let handle = tokio::spawn({ - let codex = Arc::clone(&codex); + let io = Arc::clone(&io); let parent_session = Arc::clone(&parent_session); let parent_ctx = Arc::clone(&parent_ctx); let cancel_token = cancel_token.clone(); async move { handle_request_permissions( - codex.as_ref(), + io.as_ref(), &parent_session, &parent_ctx, RequestPermissionsEvent { @@ -323,23 +320,22 @@ async fn handle_exec_approval_uses_call_id_for_guardian_review_and_approval_id_f let (tx_sub, rx_sub) = bounded(SUBMISSION_CHANNEL_CAPACITY); let (_tx_events, rx_events_child) = bounded(SUBMISSION_CHANNEL_CAPACITY); let (_agent_status_tx, agent_status) = watch::channel(AgentStatus::PendingInit); - let codex = Arc::new(Codex { + let io = Arc::new(SessionIo { tx_sub, rx_event: rx_events_child, agent_status, - session: Arc::clone(&parent_session), session_loop_termination: completed_session_loop_termination(), }); let cancel_token = CancellationToken::new(); let handle = tokio::spawn({ - let codex = Arc::clone(&codex); + let io = Arc::clone(&io); let parent_session = Arc::clone(&parent_session); let parent_ctx = Arc::clone(&parent_ctx); let cancel_token = cancel_token.clone(); async move { handle_exec_approval( - codex.as_ref(), + io.as_ref(), "child-turn-1".to_string(), &parent_session, &parent_ctx, diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index f0c4244586f9..2b22f37c1cc3 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/codex-rs/core/src/codex_thread.rs @@ -1,9 +1,10 @@ use crate::agent::AgentStatus; use crate::config::ConstraintResult; use crate::elicitation::ElicitationRegistration; -use crate::session::Codex; +use crate::session::SessionIo; use crate::session::SessionSettingsUpdate; use crate::session::SteerInputError; +use crate::session::session::Session; use codex_exec_server::SelectedCapabilityRootsStatus; use codex_features::Feature; use codex_otel::SessionTelemetry; @@ -159,7 +160,8 @@ pub struct CodexThreadSettingsOverrides { } pub struct CodexThread { - pub(crate) codex: Codex, + pub(crate) session: Arc, + pub(crate) io: SessionIo, pub(crate) session_source: SessionSource, session_configured: SessionConfiguredEvent, rollout_path: Option, @@ -184,13 +186,15 @@ pub struct BackgroundTerminalInfo { /// (formerly called a conversation) in Codex. impl CodexThread { pub(crate) fn new( - codex: Codex, + session: Arc, + io: SessionIo, session_configured: SessionConfiguredEvent, rollout_path: Option, session_source: SessionSource, ) -> Self { Self { - codex, + session, + io, session_source, session_configured, rollout_path, @@ -199,26 +203,25 @@ impl CodexThread { } pub async fn submit(&self, op: Op) -> CodexResult { - self.codex.submit(op).await + self.io.submit(op).await } /// Returns the session telemetry handle for thread-scoped production instrumentation. pub fn session_telemetry(&self) -> SessionTelemetry { - self.codex.session.services.session_telemetry.clone() + self.session.services.session_telemetry.clone() } pub async fn shutdown_and_wait(&self) -> CodexResult<()> { - self.codex.shutdown_and_wait().await + self.io.shutdown_and_wait().await } /// Wait until the underlying session loop has terminated. pub async fn wait_until_terminated(&self) { - self.codex.session_loop_termination.clone().await; + self.io.session_loop_termination.clone().await; } pub(crate) async fn emit_thread_resume_lifecycle(&self) { for contributor in self - .codex .session .services .extensions @@ -226,28 +229,25 @@ impl CodexThread { { contributor .on_thread_resume(codex_extension_api::ThreadResumeInput { - session_store: &self.codex.session.services.session_extension_data, - thread_store: &self.codex.session.services.thread_extension_data, + session_store: &self.session.services.session_extension_data, + thread_store: &self.session.services.thread_extension_data, }) .await; } } pub async fn emit_thread_idle_lifecycle_if_idle(&self) { - self.codex - .session - .emit_thread_idle_lifecycle_if_idle() - .await; + self.session.emit_thread_idle_lifecycle_if_idle().await; } #[doc(hidden)] pub async fn ensure_rollout_materialized(&self) { - self.codex.session.ensure_rollout_materialized().await; + self.session.ensure_rollout_materialized().await; } #[doc(hidden)] pub async fn flush_rollout(&self) -> std::io::Result<()> { - self.codex.session.flush_rollout().await + self.session.flush_rollout().await } pub async fn submit_with_trace( @@ -255,7 +255,7 @@ impl CodexThread { op: Op, trace: Option, ) -> CodexResult { - self.codex.submit_with_trace(op, trace).await + self.io.submit_with_trace(op, trace).await } pub async fn submit_user_input_with_client_user_message_id( @@ -264,20 +264,19 @@ impl CodexThread { trace: Option, client_user_message_id: Option, ) -> CodexResult { - self.codex - .session + self.session .services .agent_control - .ensure_execution_capacity_for_op(self.session_configured.thread_id, &op) + .ensure_execution_capacity_for_op(self.session.thread_id(), &op) .await?; - self.codex + self.io .submit_user_input_with_client_user_message_id(op, trace, client_user_message_id) .await } /// Persist whether this thread is eligible for future memory generation. pub async fn set_thread_memory_mode(&self, mode: ThreadMemoryMode) -> anyhow::Result<()> { - self.codex.set_thread_memory_mode(mode).await + self.session.set_thread_memory_mode(mode).await } pub async fn steer_input( @@ -288,7 +287,7 @@ impl CodexThread { client_user_message_id: Option, responsesapi_client_metadata: Option>, ) -> Result { - self.codex + self.session .steer_input( input, additional_context, @@ -308,7 +307,7 @@ impl CodexThread { &self, items: Vec, ) -> Result<(), Vec> { - self.codex.session.inject_if_running(items).await + self.session.inject_if_running(items).await } /// Starts an automatic regular turn with model-visible items only when idle @@ -328,7 +327,7 @@ impl CodexThread { &self, items: Vec, ) -> Result<(), TryStartTurnIfIdleError> { - self.codex.session.try_start_turn_if_idle(items).await + self.session.try_start_turn_if_idle(items).await } pub async fn set_app_server_client_info( @@ -337,7 +336,7 @@ impl CodexThread { app_server_client_version: Option, mcp_elicitations_auto_deny: bool, ) -> ConstraintResult<()> { - self.codex + self.session .set_app_server_client_info( app_server_client_name, app_server_client_version, @@ -347,8 +346,7 @@ impl CodexThread { } pub async fn set_openai_form_elicitation_support(&self, supported: bool) -> anyhow::Result<()> { - self.codex - .session + self.session .set_openai_form_elicitation_support(supported) .await } @@ -359,7 +357,7 @@ impl CodexThread { overrides: CodexThreadSettingsOverrides, ) -> ConstraintResult { let updates = self.thread_settings_update(overrides).await; - self.codex.session.preview_settings(&updates).await + self.session.preview_settings(&updates).await } async fn thread_settings_update( @@ -385,8 +383,7 @@ impl CodexThread { let collaboration_mode = if let Some(collaboration_mode) = collaboration_mode { collaboration_mode } else { - self.codex - .session + self.session .collaboration_mode() .await .with_updates(model, effort, /*developer_instructions*/ None) @@ -411,30 +408,27 @@ impl CodexThread { /// Use sparingly: this is intended to be removed soon. pub async fn submit_with_id(&self, sub: Submission) -> CodexResult<()> { - self.codex.submit_with_id(sub).await + self.io.submit_with_id(sub).await } pub async fn next_event(&self) -> CodexResult { - self.codex.next_event().await + self.io.next_event().await } pub async fn agent_status(&self) -> AgentStatus { - self.codex.agent_status().await + self.io.agent_status().await } pub async fn list_background_terminals(&self) -> Vec { - self.codex.session.list_background_terminals().await + self.session.list_background_terminals().await } pub async fn terminate_background_terminal(&self, process_id: i32) -> bool { - self.codex - .session - .terminate_background_terminal(process_id) - .await + self.session.terminate_background_terminal(process_id).await } pub(crate) fn subscribe_status(&self) -> watch::Receiver { - self.codex.agent_status.clone() + self.io.agent_status.clone() } /// Returns the complete token usage snapshot currently cached for this thread. @@ -445,7 +439,7 @@ impl CodexThread { /// `total_token_usage` would drop last-turn usage and make the v2 /// `thread/tokenUsage/updated` payload incomplete. pub async fn token_usage_info(&self) -> Option { - self.codex.session.token_usage_info().await + self.session.token_usage_info().await } /// Records a user-role session-prefix message without creating a new user turn boundary. @@ -457,8 +451,7 @@ impl CodexThread { phase: None, internal_chat_message_metadata_passthrough: None, }; - self.codex - .session + self.session .inject_no_new_turn(vec![item], /*current_turn_context*/ None) .await; } @@ -471,24 +464,21 @@ impl CodexThread { )); } - let turn_context = self.codex.session.new_default_turn().await; - if self.codex.session.reference_context_item().await.is_none() { + let turn_context = self.session.new_default_turn().await; + if self.session.reference_context_item().await.is_none() { // This history-only API runs without run_turn, so it owns its initial step. let step_context = self - .codex .session .capture_step_context(Arc::clone(&turn_context)) .await; - self.codex - .session + self.session .record_context_updates_and_set_reference_context_item(step_context.as_ref()) .await; } - self.codex - .session + self.session .inject_no_new_turn(items, Some(turn_context.as_ref())) .await; - self.codex.session.flush_rollout().await?; + self.session.flush_rollout().await?; Ok(()) } @@ -501,12 +491,11 @@ impl CodexThread { } pub(crate) fn is_running(&self) -> bool { - !self.codex.tx_sub.is_closed() + !self.io.tx_sub.is_closed() } pub async fn guardian_trunk_rollout_path(&self) -> Option { - self.codex - .session + self.session .guardian_review_session .trunk_rollout_path() .await @@ -517,7 +506,6 @@ impl CodexThread { include_archived: bool, ) -> ThreadStoreResult { let live_thread = self - .codex .session .live_thread_for_persistence("load history") .map_err(|err| ThreadStoreError::Internal { @@ -532,7 +520,6 @@ impl CodexThread { include_history: bool, ) -> ThreadStoreResult { let live_thread = self - .codex .session .live_thread_for_persistence("read thread") .map_err(|err| ThreadStoreError::Internal { @@ -549,7 +536,6 @@ impl CodexThread { include_archived: bool, ) -> ThreadStoreResult { let live_thread = self - .codex .session .live_thread_for_persistence("update thread metadata") .map_err(|err| ThreadStoreError::Internal { @@ -561,7 +547,6 @@ impl CodexThread { /// Appends rollout items through the live thread so derived metadata stays in sync. pub async fn append_rollout_items(&self, items: &[RolloutItem]) -> ThreadStoreResult<()> { let live_thread = self - .codex .session .live_thread_for_persistence("append rollout items") .map_err(|err| ThreadStoreError::Internal { @@ -571,16 +556,16 @@ impl CodexThread { } pub fn state_db(&self) -> Option { - self.codex.state_db() + self.session.state_db() } pub async fn config_snapshot(&self) -> ThreadConfigSnapshot { - self.codex.thread_config_snapshot().await + self.session.thread_config_snapshot().await } /// Returns the files that supplied the thread's loaded model instructions. pub async fn instruction_sources(&self) -> Vec { - self.codex.instruction_sources().await + self.session.instruction_sources().await } /// Returns loaded instruction sources rendered as legacy app-server path strings. @@ -593,19 +578,18 @@ impl CodexThread { } pub async fn config(&self) -> Arc { - self.codex.session.get_config().await + self.session.get_config().await } /// Resolves the MCP runtime configuration using this thread's extension data. pub async fn runtime_mcp_config(&self, config: &crate::config::Config) -> codex_mcp::McpConfig { - self.codex.session.runtime_mcp_config(config).await + self.session.runtime_mcp_config(config).await } /// Returns the exact MCP config, environment bindings, and manager most recently published. pub async fn current_mcp_runtime(&self) -> Arc { - let turn_context = self.codex.session.new_default_turn().await; - self.codex - .session + let turn_context = self.session.new_default_turn().await; + self.session .capture_step_context(turn_context) .await .mcp @@ -613,30 +597,27 @@ impl CodexThread { } pub fn multi_agent_version(&self) -> Option { - self.codex.session.multi_agent_version() + self.session.multi_agent_version() } /// Refresh the thread's layer-backed user config state from a caller-supplied /// config snapshot. Thread-scoped layers and session-static settings remain /// unchanged. pub async fn refresh_runtime_config(&self, next_config: crate::config::Config) { - self.codex.session.refresh_runtime_config(next_config).await; + self.session.refresh_runtime_config(next_config).await; } pub async fn environment_selections(&self) -> Vec { - self.codex.thread_environment_selections().await + self.session.thread_environment_selections().await } /// Passively inspects the selected capability roots whose environments are ready now. pub fn inspect_selected_capability_roots(&self) -> SelectedCapabilityRootsStatus { - self.codex - .session + self.session .services .turn_environments .environment_manager() - .inspect_selected_capability_roots( - &self.codex.session.services.selected_capability_roots, - ) + .inspect_selected_capability_roots(&self.session.services.selected_capability_roots) } pub async fn read_mcp_resource( @@ -669,7 +650,7 @@ impl CodexThread { } pub fn enabled(&self, feature: Feature) -> bool { - self.codex.enabled(feature) + self.session.enabled(feature) } pub async fn increment_out_of_band_elicitation_count(&self) -> CodexResult { @@ -678,7 +659,7 @@ impl CodexThread { CodexErr::Fatal("out-of-band elicitation count overflowed".to_string()) })?; if elicitations.count == 0 { - elicitations.registration = Some(self.codex.session.services.elicitations.register()); + elicitations.registration = Some(self.session.services.elicitations.register()); } elicitations.count = incremented; Ok(incremented) diff --git a/codex-rs/core/src/guardian/review_session.rs b/codex-rs/core/src/guardian/review_session.rs index 565b09e8dccc..8acb839208be 100644 --- a/codex-rs/core/src/guardian/review_session.rs +++ b/codex-rs/core/src/guardian/review_session.rs @@ -43,7 +43,7 @@ use crate::config::NetworkProxySpec; use crate::config::Permissions; use crate::context::ContextualUserFragment; use crate::context::GuardianFollowupReviewReminder; -use crate::session::Codex; +use crate::session::SessionIo; use crate::session::session::Session; use crate::session::turn_context::TurnContext; use codex_config::types::McpServerConfig; @@ -106,7 +106,8 @@ struct GuardianReviewSessionState { } struct GuardianReviewSession { - codex: Codex, + session: Arc, + io: SessionIo, cancel_token: CancellationToken, reuse_key: GuardianReviewSessionReuseKey, review_lock: Semaphore, @@ -220,7 +221,7 @@ pub(crate) fn prompt_cache_key_override_for_review_session( impl GuardianReviewSession { async fn shutdown(&self) { self.cancel_token.cancel(); - let _ = self.codex.shutdown_and_wait().await; + let _ = self.io.shutdown_and_wait().await; } fn shutdown_in_background(self: &Arc) { @@ -235,7 +236,7 @@ impl GuardianReviewSession { } async fn refresh_last_committed_fork_snapshot(&self) { - match load_rollout_items_for_fork(&self.codex.session).await { + match load_rollout_items_for_fork(&self.session).await { Ok(Some(items)) if !items.is_empty() => { let mut state = self.state.lock().await; let prior_review_count = state.prior_review_count; @@ -332,8 +333,8 @@ impl GuardianReviewSessionManager { pub(crate) async fn trunk_rollout_path(&self) -> Option { let trunk = self.state.lock().await.trunk.clone()?; - trunk.codex.session.ensure_rollout_materialized().await; - match trunk.codex.session.current_rollout_path().await { + trunk.session.ensure_rollout_materialized().await; + match trunk.session.current_rollout_path().await { Ok(path) => path, Err(err) => { warn!("failed to resolve guardian trunk rollout path: {err}"); @@ -492,14 +493,15 @@ impl GuardianReviewSessionManager { } #[cfg(test)] - pub(crate) async fn cache_for_test(&self, codex: Codex) { + pub(crate) async fn cache_for_test(&self, session: Arc, io: SessionIo) { let reuse_key = GuardianReviewSessionReuseKey::from_spawn_config( - codex.session.get_config().await.as_ref(), - codex.session.user_instructions().await, + session.get_config().await.as_ref(), + session.user_instructions().await, ); self.state.lock().await.trunk = Some(Arc::new(GuardianReviewSession { reuse_key, - codex, + session, + io, cancel_token: CancellationToken::new(), review_lock: Semaphore::new(/*permits*/ 1), state: Mutex::new(GuardianReviewState { @@ -511,10 +513,10 @@ impl GuardianReviewSessionManager { } #[cfg(test)] - pub(crate) async fn register_ephemeral_for_test(&self, codex: Codex) { + pub(crate) async fn register_ephemeral_for_test(&self, session: Arc, io: SessionIo) { let reuse_key = GuardianReviewSessionReuseKey::from_spawn_config( - codex.session.get_config().await.as_ref(), - codex.session.user_instructions().await, + session.get_config().await.as_ref(), + session.user_instructions().await, ); self.state .lock() @@ -522,7 +524,8 @@ impl GuardianReviewSessionManager { .ephemeral_reviews .push(Arc::new(GuardianReviewSession { reuse_key, - codex, + session, + io, cancel_token: CancellationToken::new(), review_lock: Semaphore::new(/*permits*/ 1), state: Mutex::new(GuardianReviewState { @@ -553,7 +556,7 @@ impl GuardianReviewSessionManager { .trunk .clone() .expect("guardian trunk should exist"); - trunk.codex.session.send_event_raw(event).await; + trunk.session.send_event_raw(event).await; } async fn remove_trunk_if_current( @@ -664,7 +667,7 @@ async fn spawn_guardian_review_session( ), None => (None, 0, None), }; - let codex = Box::pin(run_codex_thread_interactive( + let (session, io) = Box::pin(run_codex_thread_interactive( spawn_config, parent_session.services.auth_manager.clone(), parent_session.services.models_manager.clone(), @@ -677,7 +680,8 @@ async fn spawn_guardian_review_session( .await?; Ok(GuardianReviewSession { - codex, + session, + io, cancel_token, reuse_key, review_lock: Semaphore::new(/*permits*/ 1), @@ -728,7 +732,7 @@ async fn run_review_on_session( .or_else(|| model_info.default_reasoning_level.clone()); let mut analytics_result = GuardianReviewAnalyticsResult::from_session(GuardianReviewSessionAnalyticsParams { - guardian_thread_id: review_session.codex.session.thread_id.to_string(), + guardian_thread_id: review_session.session.thread_id().to_string(), guardian_session_kind, guardian_model: params.model.clone(), guardian_reasoning_effort: guardian_reasoning_effort.map(|effort| effort.to_string()), @@ -751,9 +755,7 @@ async fn run_review_on_session( .parent_session .services .network_approval - .sync_session_approved_hosts_to( - &review_session.codex.session.services.network_approval, - ) + .sync_session_approved_hosts_to(&review_session.session.services.network_approval) .await; build_guardian_prompt_items_with_parent_turn( @@ -784,7 +786,6 @@ async fn run_review_on_session( let reviewed_action_truncated = prompt_items.reviewed_action_truncated; let transcript_cursor = prompt_items.transcript_cursor; let token_usage_at_review_start = review_session - .codex .session .total_token_usage() .await @@ -803,7 +804,7 @@ async fn run_review_on_session( let submit_result = run_before_review_deadline( deadline, params.external_cancel.as_ref(), - Box::pin(review_session.codex.submit(Op::UserInput { + Box::pin(review_session.io.submit(Op::UserInput { items: prompt_items.items, final_output_json_schema: Some(params.schema.clone()), responsesapi_client_metadata: None, @@ -857,7 +858,7 @@ async fn run_review_on_session( .await; if matches!(outcome.0, GuardianReviewSessionOutcome::Completed(_)) { if outcome.2 - && let Some(total_token_usage) = review_session.codex.session.total_token_usage().await + && let Some(total_token_usage) = review_session.session.total_token_usage().await { analytics_result.token_usage = Some(token_usage_delta( &token_usage_at_review_start, @@ -874,7 +875,6 @@ async fn run_review_on_session( async fn append_guardian_followup_reminder(review_session: &GuardianReviewSession) { let reminder: ResponseItem = ContextualUserFragment::into(GuardianFollowupReviewReminder); review_session - .codex .session .inject_no_new_turn(vec![reminder], /*current_turn_context*/ None) .await; @@ -905,7 +905,7 @@ async fn wait_for_guardian_review( tokio::select! { _ = &mut timeout => { let keep_review_session = interrupt_and_drain_turn( - &review_session.codex, + &review_session.io, expected_turn_id, ) .await @@ -920,14 +920,14 @@ async fn wait_for_guardian_review( } } => { let keep_review_session = interrupt_and_drain_turn( - &review_session.codex, + &review_session.io, expected_turn_id, ) .await .is_ok(); return (GuardianReviewSessionOutcome::Aborted, keep_review_session, false); } - event = review_session.codex.next_event() => { + event = review_session.io.next_event() => { match event { Ok(event) if !event_matches_turn(&event, expected_turn_id) => {} Ok(event) => match event.msg { @@ -1104,12 +1104,12 @@ async fn run_before_review_deadline_with_cancel( result } -async fn interrupt_and_drain_turn(codex: &Codex, expected_turn_id: &str) -> anyhow::Result<()> { - let _ = codex.submit(Op::Interrupt).await; +async fn interrupt_and_drain_turn(io: &SessionIo, expected_turn_id: &str) -> anyhow::Result<()> { + let _ = io.submit(Op::Interrupt).await; tokio::time::timeout(GUARDIAN_INTERRUPT_DRAIN_TIMEOUT, async { loop { - let event = codex.next_event().await?; + let event = io.next_event().await?; if event_matches_turn(&event, expected_turn_id) && matches!( event.msg, @@ -1154,11 +1154,11 @@ mod tests { ( GuardianReviewSession { - codex: Codex { + session, + io: SessionIo { tx_sub, rx_event, agent_status, - session, session_loop_termination: crate::session::completed_session_loop_termination(), }, cancel_token: CancellationToken::new(), @@ -1960,10 +1960,10 @@ mod tests { .await .expect("queue current turn abort"); - interrupt_and_drain_turn(&review_session.codex, "current-turn") + interrupt_and_drain_turn(&review_session.io, "current-turn") .await .expect("drain current turn"); - assert!(review_session.codex.rx_event.try_recv().is_err()); + assert!(review_session.io.rx_event.try_recv().is_err()); } } diff --git a/codex-rs/core/src/prompt_debug.rs b/codex-rs/core/src/prompt_debug.rs index bf424cc94169..d8174bd911c2 100644 --- a/codex-rs/core/src/prompt_debug.rs +++ b/codex-rs/core/src/prompt_debug.rs @@ -66,7 +66,7 @@ pub async fn build_prompt_input( ); let thread = thread_manager.start_thread(config).await?; - let output = build_prompt_input_from_session(&thread.thread.codex.session, input).await; + let output = build_prompt_input_from_session(&thread.thread.session, input).await; let shutdown = thread.thread.shutdown_and_wait().await; let _removed = thread_manager.remove_thread(&thread.thread_id).await; diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 8f7505523bfd..941be63427b2 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -383,14 +383,16 @@ use codex_utils_absolute_path::AbsolutePathBuf; #[cfg(test)] use codex_utils_stream_parser::ProposedPlanSegment; -/// The high-level interface to the Codex system. -/// It operates as a queue pair where you send submissions and receive events. -pub struct Codex { +/// Queue and lifecycle endpoints for a running [`Session`]. +/// +/// 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. +pub(crate) struct SessionIo { pub(crate) tx_sub: Sender, pub(crate) rx_event: Receiver, // Last known status of the agent. pub(crate) agent_status: watch::Receiver, - pub(crate) session: Arc, // Shared future for the background submission loop completion so multiple // callers can wait for shutdown. pub(crate) session_loop_termination: SessionLoopTermination, @@ -398,14 +400,7 @@ pub struct Codex { pub(crate) type SessionLoopTermination = Shared>; -/// Wrapper returned by [`Codex::spawn`] containing the spawned [`Codex`] and -/// the unique session id. -pub struct CodexSpawnOk { - pub codex: Codex, - pub thread_id: ThreadId, -} - -pub(crate) struct CodexSpawnArgs { +pub(crate) struct SessionSpawnArgs { pub(crate) config: Config, pub(crate) allow_provider_model_fallback: bool, pub(crate) user_instructions: LoadedUserInstructions, @@ -470,9 +465,9 @@ pub(crate) const SUBMISSION_CHANNEL_CAPACITY: usize = 512; const CYBER_VERIFY_URL: &str = "https://chatgpt.com/cyber"; const CYBER_SAFETY_URL: &str = "https://developers.openai.com/codex/concepts/cyber-safety"; -impl Codex { - /// Spawn a new [`Codex`] and initialize the session. - pub(crate) async fn spawn(args: CodexSpawnArgs) -> CodexResult { +impl Session { + /// Spawn and initialize a new session. + pub(crate) async fn spawn(args: SessionSpawnArgs) -> CodexResult<(Arc, SessionIo)> { let parent_trace = match args.parent_trace { Some(trace) => { if codex_otel::context_from_w3c_trace_context(&trace).is_some() { @@ -488,7 +483,7 @@ impl Codex { if let Some(trace) = parent_trace.as_ref() { let _ = set_parent_from_w3c_trace_context(&thread_spawn_span, trace); } - Self::spawn_internal(CodexSpawnArgs { + Self::spawn_internal(SessionSpawnArgs { parent_trace, ..args }) @@ -496,8 +491,8 @@ impl Codex { .await } - async fn spawn_internal(args: CodexSpawnArgs) -> CodexResult { - let CodexSpawnArgs { + async fn spawn_internal(args: SessionSpawnArgs) -> CodexResult<(Arc, SessionIo)> { + let SessionSpawnArgs { mut config, allow_provider_model_fallback, user_instructions, @@ -674,7 +669,7 @@ impl Codex { user_shell_override, }; - // Generate a unique ID for the lifetime of this Codex session. + // Generate a unique ID for the lifetime of this session. let session_source_clone = session_configuration.session_source.clone(); let (agent_status_tx, agent_status_rx) = watch::channel(AgentStatus::PendingInit); @@ -729,23 +724,24 @@ impl Codex { .instrument(info_span!("session_loop", thread_id = %thread_id)) .await; }); - let codex = Codex { + let io = SessionIo { tx_sub, rx_event, agent_status: agent_status_rx, - session, session_loop_termination: session_loop_termination_from_handle(session_loop_handle), }; - Ok(CodexSpawnOk { codex, thread_id }) + Ok((session, io)) } +} +impl SessionIo { /// Submit the `op` wrapped in a `Submission` with a unique ID. - pub async fn submit(&self, op: Op) -> CodexResult { + pub(crate) async fn submit(&self, op: Op) -> CodexResult { self.submit_with_trace(op, /*trace*/ None).await } - pub async fn submit_with_trace( + pub(crate) async fn submit_with_trace( &self, op: Op, trace: Option, @@ -761,7 +757,7 @@ impl Codex { Ok(id) } - pub async fn submit_user_input_with_client_user_message_id( + pub(crate) async fn submit_user_input_with_client_user_message_id( &self, op: Op, trace: Option, @@ -779,9 +775,8 @@ impl Codex { Ok(id) } - /// Use sparingly: prefer `submit()` so Codex is responsible for generating - /// unique IDs for each submission. - pub async fn submit_with_id(&self, mut sub: Submission) -> CodexResult<()> { + /// Use sparingly: prefer `submit()` so submission IDs are generated consistently. + pub(crate) async fn submit_with_id(&self, mut sub: Submission) -> CodexResult<()> { if sub.trace.is_none() { sub.trace = current_span_w3c_trace_context(); } @@ -792,18 +787,7 @@ impl Codex { Ok(()) } - /// Persist a thread-level memory mode update for the active session. - /// - /// This is a local-only operation that updates rollout metadata directly - /// and does not involve the model. - pub async fn set_thread_memory_mode( - &self, - mode: codex_protocol::protocol::ThreadMemoryMode, - ) -> anyhow::Result<()> { - handlers::persist_thread_memory_mode_update(&self.session, mode).await - } - - pub async fn shutdown_and_wait(&self) -> CodexResult<()> { + pub(crate) async fn shutdown_and_wait(&self) -> CodexResult<()> { let session_loop_termination = self.session_loop_termination.clone(); match self.submit(Op::Shutdown).await { Ok(_) => {} @@ -814,7 +798,7 @@ impl Codex { Ok(()) } - pub async fn next_event(&self) -> CodexResult { + pub(crate) async fn next_event(&self) -> CodexResult { let event = self .rx_event .recv() @@ -823,80 +807,9 @@ impl Codex { Ok(event) } - pub async fn steer_input( - &self, - input: Vec, - additional_context: BTreeMap, - expected_turn_id: Option<&str>, - client_user_message_id: Option, - responsesapi_client_metadata: Option>, - ) -> Result { - self.session - .steer_input( - input, - additional_context, - expected_turn_id, - client_user_message_id, - responsesapi_client_metadata, - ) - .await - } - - pub(crate) async fn set_app_server_client_info( - &self, - app_server_client_name: Option, - app_server_client_version: Option, - mcp_elicitations_auto_deny: bool, - ) -> ConstraintResult<()> { - self.session - .update_settings(SessionSettingsUpdate { - app_server_client_name, - app_server_client_version, - ..Default::default() - }) - .await?; - self.session - .services - .latest_mcp_runtime() - .manager() - .set_elicitations_auto_deny(mcp_elicitations_auto_deny); - Ok(()) - } - pub(crate) async fn agent_status(&self) -> AgentStatus { self.agent_status.borrow().clone() } - - pub(crate) async fn thread_config_snapshot(&self) -> ThreadConfigSnapshot { - let state = self.session.state.lock().await; - state.session_configuration.thread_config_snapshot() - } - - pub(crate) async fn instruction_sources(&self) -> Vec { - self.session - .services - .agents_md_manager - .get_loaded() - .await - .as_ref() - .map_or_else(Vec::new, |instructions| instructions.sources().collect()) - } - - pub(crate) async fn thread_environment_selections(&self) -> Vec { - let state = self.session.state.lock().await; - state - .session_configuration - .environment_selections() - .to_vec() - } - - pub(crate) fn state_db(&self) -> Option { - self.session.state_db() - } - - pub(crate) fn enabled(&self, feature: Feature) -> bool { - self.session.enabled(feature) - } } /// Generate a core submission ID. App-server exposes submission IDs that @@ -1197,6 +1110,13 @@ impl Session { self.services.live_thread.as_ref() } + pub(crate) async fn set_thread_memory_mode( + self: &Arc, + mode: ThreadMemoryMode, + ) -> anyhow::Result<()> { + handlers::persist_thread_memory_mode_update(self, mode).await + } + /// Flush rollout writes and return the final durability-barrier result. #[instrument(name = "session.flush_rollout", level = "trace", skip_all)] pub(crate) async fn flush_rollout(&self) -> std::io::Result<()> { @@ -1568,6 +1488,47 @@ impl Session { .map(|configuration| configuration.thread_config_snapshot()) } + pub(crate) async fn thread_config_snapshot(&self) -> ThreadConfigSnapshot { + let state = self.state.lock().await; + state.session_configuration.thread_config_snapshot() + } + + pub(crate) async fn set_app_server_client_info( + &self, + app_server_client_name: Option, + app_server_client_version: Option, + mcp_elicitations_auto_deny: bool, + ) -> ConstraintResult<()> { + self.update_settings(SessionSettingsUpdate { + app_server_client_name, + app_server_client_version, + ..Default::default() + }) + .await?; + self.services + .latest_mcp_runtime() + .manager() + .set_elicitations_auto_deny(mcp_elicitations_auto_deny); + Ok(()) + } + + pub(crate) async fn instruction_sources(&self) -> Vec { + self.services + .agents_md_manager + .get_loaded() + .await + .as_ref() + .map_or_else(Vec::new, |instructions| instructions.sources().collect()) + } + + pub(crate) async fn thread_environment_selections(&self) -> Vec { + let state = self.state.lock().await; + state + .session_configuration + .environment_selections() + .to_vec() + } + pub(crate) async fn set_session_startup_prewarm( &self, startup_prewarm: SessionStartupPrewarmHandle, diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 8c0aff826055..fbea8f551ed2 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -6387,15 +6387,13 @@ async fn request_permissions_is_auto_denied_when_granular_policy_blocks_tool_req #[tokio::test] async fn submit_with_id_captures_current_span_trace_context() { - let (session, _turn_context) = make_session_and_context().await; + let (_session, _turn_context) = make_session_and_context().await; let (tx_sub, rx_sub) = async_channel::bounded(1); let (_tx_event, rx_event) = async_channel::unbounded(); - let (_agent_status_tx, agent_status) = watch::channel(AgentStatus::PendingInit); - let codex = Codex { + let io = SessionIo { tx_sub, rx_event, - agent_status, - session: Arc::new(session), + agent_status: watch::channel(AgentStatus::PendingInit).1, session_loop_termination: completed_session_loop_termination(), }; @@ -6414,15 +6412,14 @@ async fn submit_with_id_captures_current_span_trace_context() { let expected_trace = async { let expected_trace = current_span_w3c_trace_context().expect("current span should have trace context"); - codex - .submit_with_id(Submission { - id: "sub-1".into(), - op: Op::Interrupt, - client_user_message_id: None, - trace: None, - }) - .await - .expect("submit should succeed"); + io.submit_with_id(Submission { + id: "sub-1".into(), + op: Op::Interrupt, + client_user_message_id: None, + trace: None, + }) + .await + .expect("submit should succeed"); expected_trace } .instrument(request_span) @@ -7115,30 +7112,28 @@ async fn submission_loop_channel_close_aborts_active_turn_before_thread_stop_lif #[tokio::test] async fn shutdown_and_wait_allows_multiple_waiters() { - let (session, _turn_context) = make_session_and_context().await; + let (_session, _turn_context) = make_session_and_context().await; let (tx_sub, rx_sub) = async_channel::bounded(4); let (_tx_event, rx_event) = async_channel::unbounded(); - let (_agent_status_tx, agent_status) = watch::channel(AgentStatus::PendingInit); let session_loop_handle = tokio::spawn(async move { let shutdown: Submission = rx_sub.recv().await.expect("shutdown submission"); assert_eq!(shutdown.op, Op::Shutdown); tokio::time::sleep(StdDuration::from_millis(50)).await; }); - let codex = Arc::new(Codex { + let io = Arc::new(SessionIo { tx_sub, rx_event, - agent_status, - session: Arc::new(session), + agent_status: watch::channel(AgentStatus::PendingInit).1, session_loop_termination: session_loop_termination_from_handle(session_loop_handle), }); let waiter_1 = { - let codex = Arc::clone(&codex); - tokio::spawn(async move { codex.shutdown_and_wait().await }) + let io = Arc::clone(&io); + tokio::spawn(async move { io.shutdown_and_wait().await }) }; let waiter_2 = { - let codex = Arc::clone(&codex); - tokio::spawn(async move { codex.shutdown_and_wait().await }) + let io = Arc::clone(&io); + tokio::spawn(async move { io.shutdown_and_wait().await }) }; waiter_1 @@ -7153,26 +7148,24 @@ async fn shutdown_and_wait_allows_multiple_waiters() { #[tokio::test] async fn shutdown_and_wait_waits_when_shutdown_is_already_in_progress() { - let (session, _turn_context) = make_session_and_context().await; + let (_session, _turn_context) = make_session_and_context().await; let (tx_sub, rx_sub) = async_channel::bounded(4); drop(rx_sub); let (_tx_event, rx_event) = async_channel::unbounded(); - let (_agent_status_tx, agent_status) = watch::channel(AgentStatus::PendingInit); let (shutdown_complete_tx, shutdown_complete_rx) = tokio::sync::oneshot::channel(); let session_loop_handle = tokio::spawn(async move { let _ = shutdown_complete_rx.await; }); - let codex = Arc::new(Codex { + let io = Arc::new(SessionIo { tx_sub, rx_event, - agent_status, - session: Arc::new(session), + agent_status: watch::channel(AgentStatus::PendingInit).1, session_loop_termination: session_loop_termination_from_handle(session_loop_handle), }); let waiter = { - let codex = Arc::clone(&codex); - tokio::spawn(async move { codex.shutdown_and_wait().await }) + let io = Arc::clone(&io); + tokio::spawn(async move { io.shutdown_and_wait().await }) }; tokio::time::sleep(StdDuration::from_millis(10)).await; @@ -7195,23 +7188,20 @@ async fn shutdown_and_wait_shuts_down_cached_guardian_subagent() { let parent_config = Arc::clone(&parent_turn_context.config); let (parent_tx_sub, parent_rx_sub) = async_channel::bounded(4); let (_parent_tx_event, parent_rx_event) = async_channel::unbounded(); - let (_parent_status_tx, parent_agent_status) = watch::channel(AgentStatus::PendingInit); let parent_session_for_loop = Arc::clone(&parent_session); let parent_session_loop_handle = tokio::spawn(async move { submission_loop(parent_session_for_loop, parent_config, parent_rx_sub).await; }); - let parent_codex = Codex { + let parent_io = SessionIo { tx_sub: parent_tx_sub, rx_event: parent_rx_event, - agent_status: parent_agent_status, - session: Arc::clone(&parent_session), + agent_status: watch::channel(AgentStatus::PendingInit).1, session_loop_termination: session_loop_termination_from_handle(parent_session_loop_handle), }; let (child_session, _child_turn_context) = make_session_and_context().await; let (child_tx_sub, child_rx_sub) = async_channel::bounded(4); let (_child_tx_event, child_rx_event) = async_channel::unbounded(); - let (_child_status_tx, child_agent_status) = watch::channel(AgentStatus::PendingInit); let (child_shutdown_tx, child_shutdown_rx) = tokio::sync::oneshot::channel(); let child_session_loop_handle = tokio::spawn(async move { let shutdown: Submission = child_rx_sub @@ -7223,19 +7213,19 @@ async fn shutdown_and_wait_shuts_down_cached_guardian_subagent() { .send(()) .expect("child shutdown signal should be delivered"); }); - let child_codex = Codex { + let child_session = Arc::new(child_session); + let child_io = SessionIo { tx_sub: child_tx_sub, rx_event: child_rx_event, - agent_status: child_agent_status, - session: Arc::new(child_session), + agent_status: watch::channel(AgentStatus::PendingInit).1, session_loop_termination: session_loop_termination_from_handle(child_session_loop_handle), }; parent_session .guardian_review_session - .cache_for_test(child_codex) + .cache_for_test(child_session, child_io) .await; - parent_codex + parent_io .shutdown_and_wait() .await .expect("parent shutdown should succeed"); @@ -7254,18 +7244,17 @@ async fn cached_guardian_subagent_exposes_its_rollout_path() { let child_rollout_path = attach_thread_persistence(&mut child_session).await; let (child_tx_sub, _child_rx_sub) = async_channel::bounded(4); let (_child_tx_event, child_rx_event) = async_channel::unbounded(); - let (_child_status_tx, child_agent_status) = watch::channel(AgentStatus::PendingInit); let child_session_loop_handle = tokio::spawn(async {}); - let child_codex = Codex { + let child_session = Arc::new(child_session); + let child_io = SessionIo { tx_sub: child_tx_sub, rx_event: child_rx_event, - agent_status: child_agent_status, - session: Arc::new(child_session), + agent_status: watch::channel(AgentStatus::PendingInit).1, session_loop_termination: session_loop_termination_from_handle(child_session_loop_handle), }; parent_session .guardian_review_session - .cache_for_test(child_codex) + .cache_for_test(child_session, child_io) .await; assert_eq!( @@ -7284,23 +7273,20 @@ async fn shutdown_and_wait_shuts_down_tracked_ephemeral_guardian_review() { let parent_config = Arc::clone(&parent_turn_context.config); let (parent_tx_sub, parent_rx_sub) = async_channel::bounded(4); let (_parent_tx_event, parent_rx_event) = async_channel::unbounded(); - let (_parent_status_tx, parent_agent_status) = watch::channel(AgentStatus::PendingInit); let parent_session_for_loop = Arc::clone(&parent_session); let parent_session_loop_handle = tokio::spawn(async move { submission_loop(parent_session_for_loop, parent_config, parent_rx_sub).await; }); - let parent_codex = Codex { + let parent_io = SessionIo { tx_sub: parent_tx_sub, rx_event: parent_rx_event, - agent_status: parent_agent_status, - session: Arc::clone(&parent_session), + agent_status: watch::channel(AgentStatus::PendingInit).1, session_loop_termination: session_loop_termination_from_handle(parent_session_loop_handle), }; let (child_session, _child_turn_context) = make_session_and_context().await; let (child_tx_sub, child_rx_sub) = async_channel::bounded(4); let (_child_tx_event, child_rx_event) = async_channel::unbounded(); - let (_child_status_tx, child_agent_status) = watch::channel(AgentStatus::PendingInit); let (child_shutdown_tx, child_shutdown_rx) = tokio::sync::oneshot::channel(); let child_session_loop_handle = tokio::spawn(async move { let shutdown: Submission = child_rx_sub @@ -7312,19 +7298,19 @@ async fn shutdown_and_wait_shuts_down_tracked_ephemeral_guardian_review() { .send(()) .expect("child shutdown signal should be delivered"); }); - let child_codex = Codex { + let child_session = Arc::new(child_session); + let child_io = SessionIo { tx_sub: child_tx_sub, rx_event: child_rx_event, - agent_status: child_agent_status, - session: Arc::new(child_session), + agent_status: watch::channel(AgentStatus::PendingInit).1, session_loop_termination: session_loop_termination_from_handle(child_session_loop_handle), }; parent_session .guardian_review_session - .register_ephemeral_for_test(child_codex) + .register_ephemeral_for_test(child_session, child_io) .await; - parent_codex + parent_io .shutdown_and_wait() .await .expect("parent shutdown should succeed"); diff --git a/codex-rs/core/src/session/tests/guardian_tests.rs b/codex-rs/core/src/session/tests/guardian_tests.rs index a2af3a271749..d2cee609f1b4 100644 --- a/codex-rs/core/src/session/tests/guardian_tests.rs +++ b/codex-rs/core/src/session/tests/guardian_tests.rs @@ -724,7 +724,7 @@ async fn guardian_subagent_does_not_inherit_parent_exec_policy_rules() { /*state_db*/ None, )); - let CodexSpawnOk { codex, .. } = Codex::spawn(CodexSpawnArgs { + let (session, io) = Session::spawn(SessionSpawnArgs { config, allow_provider_model_fallback: false, user_instructions: Default::default(), @@ -767,8 +767,7 @@ async fn guardian_subagent_does_not_inherit_parent_exec_policy_rules() { .expect("spawn guardian subagent"); assert_eq!( - codex - .session + session .services .exec_policy .current() @@ -781,5 +780,5 @@ async fn guardian_subagent_does_not_inherit_parent_exec_policy_rules() { }], } ); - drop(codex); + drop(io); } diff --git a/codex-rs/core/src/tasks/review.rs b/codex-rs/core/src/tasks/review.rs index 358333e02087..550eaa329ddf 100644 --- a/codex-rs/core/src/tasks/review.rs +++ b/codex-rs/core/src/tasks/review.rs @@ -137,7 +137,7 @@ async fn start_review_conversation( ) .await) .ok() - .map(|io| io.rx_event) + .map(|(_session, io)| io.rx_event) } async fn process_review_events( diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index f5ea09d60666..06e740233d3b 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -10,11 +10,11 @@ use crate::environment_selection::TurnEnvironmentSnapshot; use crate::environment_selection::default_thread_environment_selections; use crate::mcp::McpManager; use crate::rollout::truncation; -use crate::session::Codex; -use crate::session::CodexSpawnArgs; -use crate::session::CodexSpawnOk; use crate::session::INITIAL_SUBMIT_ID; +use crate::session::SessionIo; +use crate::session::SessionSpawnArgs; use crate::session::resolve_multi_agent_version; +use crate::session::session::Session; use crate::tasks::InterruptedTurnHistoryMarker; use crate::tasks::interrupted_turn_history_marker; use codex_agent_graph_store::AgentGraphStore; @@ -1281,7 +1281,7 @@ impl ThreadManagerState { // The spawn path retains only thread IDs, so look up the live // runtime again here to inherit its user instructions. Some(thread_id) => match self.get_thread(thread_id).await { - Ok(thread) => thread.codex.session.user_instructions().await, + Ok(thread) => thread.session.user_instructions().await, Err(_) => None, }, None => None, @@ -1617,9 +1617,7 @@ impl ThreadManagerState { forked_from_thread_id, ) .await; - let CodexSpawnOk { - codex, thread_id, .. - } = Box::pin(Codex::spawn(CodexSpawnArgs { + let (session, io) = Box::pin(Session::spawn(SessionSpawnArgs { config, allow_provider_model_fallback, user_instructions, @@ -1658,7 +1656,7 @@ impl ThreadManagerState { })) .await?; let new_thread = self - .finalize_thread_spawn(codex, thread_id, tracked_session_source) + .finalize_thread_spawn(session, io, tracked_session_source) .await?; if is_resumed_thread { new_thread.thread.emit_thread_resume_lifecycle().await; @@ -1668,11 +1666,12 @@ impl ThreadManagerState { async fn finalize_thread_spawn( &self, - codex: Codex, - thread_id: ThreadId, + session: Arc, + io: SessionIo, session_source: SessionSource, ) -> CodexResult { - let event = codex.next_event().await?; + let thread_id = session.thread_id(); + let event = io.next_event().await?; let session_configured = match event { Event { id, @@ -1687,7 +1686,8 @@ impl ThreadManagerState { let mut threads = self.threads.write().await; if let std::collections::hash_map::Entry::Vacant(e) = threads.entry(thread_id) { let thread = Arc::new(CodexThread::new( - codex, + session, + io, session_configured.clone(), session_configured.rollout_path.clone(), session_source, @@ -1701,7 +1701,7 @@ impl ThreadManagerState { } } - if let Err(err) = codex.shutdown_and_wait().await { + if let Err(err) = io.shutdown_and_wait().await { warn!("failed to shut down duplicate thread {thread_id}: {err}"); } Err(CodexErr::InvalidRequest(format!( @@ -1739,7 +1739,7 @@ impl ThreadManagerState { self.get_thread(*parent_thread_id) .await .ok() - .map(|thread| thread.codex.session.services.rollout_thread_trace.clone()) + .map(|thread| thread.session.services.rollout_thread_trace.clone()) .unwrap_or_else(codex_rollout_trace::ThreadTraceContext::disabled) } } diff --git a/codex-rs/core/src/thread_manager_tests.rs b/codex-rs/core/src/thread_manager_tests.rs index 9d722a94ad9c..09b8b4f99493 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/codex-rs/core/src/thread_manager_tests.rs @@ -426,14 +426,12 @@ async fn code_mode_session_provider_is_shared_across_threads() { let first_provider = first .thread - .codex .session .services .code_mode_service .session_provider(); let second_provider = second .thread - .codex .session .services .code_mode_service @@ -660,7 +658,7 @@ async fn start_thread_seeds_extension_data_for_mcp_and_lifecycle_contributors() }) .await .expect("start second thread"); - let first_session = &first_thread.thread.codex.session; + let first_session = &first_thread.thread.session; let first_originator = first_session.originator().await; let first_resolved = first_session .services @@ -673,7 +671,7 @@ async fn start_thread_seeds_extension_data_for_mcp_and_lifecycle_contributors() /*available_environment_ids*/ &[], ) .await; - let second_session = &second_thread.thread.codex.session; + let second_session = &second_thread.thread.session; let second_originator = second_session.originator().await; let second_resolved = second_session .services @@ -892,7 +890,6 @@ async fn resume_and_fork_do_not_restore_thread_environments_from_rollout() { .expect("resume source thread"); let resumed_turn = resumed .thread - .codex .session .new_turn_with_sub_id("resume-turn".to_string(), SessionSettingsUpdate::default()) .await @@ -919,7 +916,6 @@ async fn resume_and_fork_do_not_restore_thread_environments_from_rollout() { .expect("fork source thread"); let forked_turn = forked .thread - .codex .session .new_turn_with_sub_id("fork-turn".to_string(), SessionSettingsUpdate::default()) .await @@ -971,7 +967,7 @@ async fn explicit_installation_id_skips_codex_home_file() { .expect("start thread with explicit installation id"); assert!(!config.codex_home.join(INSTALLATION_ID_FILENAME).exists()); - assert_eq!(thread.thread.codex.session.installation_id, installation_id); + assert_eq!(thread.thread.session.installation_id, installation_id); thread .thread diff --git a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs b/codex-rs/core/src/tools/handlers/multi_agents_tests.rs index 6c963063f4f3..d7215f06a81d 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_tests.rs @@ -1584,9 +1584,8 @@ async fn multi_agent_v2_list_agents_returns_completed_status() { .get_thread(agent_id) .await .expect("child thread should exist"); - let child_turn = child_thread.codex.session.new_default_turn().await; + let child_turn = child_thread.session.new_default_turn().await; child_thread - .codex .session .send_event( child_turn.as_ref(), @@ -1989,7 +1988,7 @@ async fn multi_agent_v2_followup_task_completion_notifies_parent_on_every_turn() .expect("root thread should start"); // Production spawn_agent calls happen after the parent turn has resolved // and stored its runtime; mirror that before using the synthetic handler. - root.thread.codex.session.new_default_turn().await; + root.thread.session.new_default_turn().await; session.services.agent_control = manager.agent_control(); session.thread_id = root.thread_id; let session = Arc::new(session); @@ -2019,9 +2018,8 @@ async fn multi_agent_v2_followup_task_completion_notifies_parent_on_every_turn() .expect("worker thread should exist"); let worker_path = AgentPath::try_from("/root/worker").expect("worker path"); - let first_turn = thread.codex.session.new_default_turn().await; + let first_turn = thread.session.new_default_turn().await; thread - .codex .session .send_event( first_turn.as_ref(), @@ -2062,9 +2060,8 @@ async fn multi_agent_v2_followup_task_completion_notifies_parent_on_every_turn() ) })); - let second_turn = thread.codex.session.new_default_turn().await; + let second_turn = thread.session.new_default_turn().await; thread - .codex .session .send_event( second_turn.as_ref(), @@ -2226,9 +2223,8 @@ async fn multi_agent_v2_interrupted_turn_does_not_notify_parent() { .await .expect("worker thread should exist"); - let aborted_turn = thread.codex.session.new_default_turn().await; + let aborted_turn = thread.session.new_default_turn().await; thread - .codex .session .send_event( aborted_turn.as_ref(), @@ -2421,7 +2417,7 @@ async fn spawn_agent_reapplies_runtime_sandbox_after_role_config() { .get_thread(agent_id) .await .expect("spawned agent thread should exist"); - let child_turn = child_thread.codex.session.new_default_turn().await; + let child_turn = child_thread.session.new_default_turn().await; assert_eq!( child_turn.file_system_sandbox_policy(), expected_file_system_sandbox_policy @@ -3868,7 +3864,7 @@ async fn multi_agent_v2_interrupt_agent_accepts_task_name_target() { .get_thread(agent_id) .await .expect("worker thread should be resident"); - let worker_session = worker_thread.codex.session.clone(); + let worker_session = worker_thread.session.clone(); SpawnAgentHandlerV2::default() .handle(invocation( worker_session.clone(), @@ -4299,7 +4295,7 @@ async fn tool_handlers_cascade_close_and_resume_and_keep_explicitly_closed_subtr .await .expect("parent thread should start"); let parent_thread_id = parent.thread_id; - let parent_session = parent.thread.codex.session.clone(); + let parent_session = parent.thread.session.clone(); let child_turn = parent_session.new_default_turn().await; let child_spawn_output = SpawnAgentHandler::default() @@ -4326,7 +4322,7 @@ async fn tool_handlers_cascade_close_and_resume_and_keep_explicitly_closed_subtr .get_thread(child_thread_id) .await .expect("child thread should exist"); - let child_session = child_thread.codex.session.clone(); + let child_session = child_thread.session.clone(); let grandchild_spawn_output = SpawnAgentHandler::default() .handle(invocation( child_session.clone(), @@ -4430,7 +4426,7 @@ async fn tool_handlers_cascade_close_and_resume_and_keep_explicitly_closed_subtr .start_thread(config.clone()) .await .expect("operator thread should start"); - let operator_session = operator.thread.codex.session.clone(); + let operator_session = operator.thread.session.clone(); let _ = manager .agent_control() .shutdown_live_agent(parent_thread_id) @@ -4444,7 +4440,7 @@ async fn tool_handlers_cascade_close_and_resume_and_keep_explicitly_closed_subtr let parent_resume_output = ResumeAgentHandler .handle(invocation( operator_session, - operator.thread.codex.session.new_default_turn().await, + operator.thread.session.new_default_turn().await, "resume_agent", function_payload(json!({"id": parent_thread_id.to_string()})), ))