github/codex
4
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-28 08:04:55 +00:00
Code Issues Packages Projects Releases Wiki Activity

Separate session state from session I/O (#33209)

Browse Source 
## What changed

- Replace the combined `Codex` wrapper with an `Arc<Session>` and a dedicated
  `SessionIo` handle for submissions, events, status, and loop termination.
- Make `Session::spawn` return the session and I/O handle separately, and have
  `CodexThread` own both.
- Update thread management, delegated sessions, guardian reviews, and agent
  control to use the appropriate session or I/O interface.

GitOrigin-RevId: 8523bad7e7366e2969fd0cc3849b4247b141427a
This commit is contained in:
pakrym-oai
2026-07-15 02:16:53 +00:00
committed by copyberry
parent 9cddda7556
commit e4711f2a3b
19 changed files with 360 additions and 527 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/agent/control.rs
+2 -5
View File
@@ -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(
codex-rs/core/src/agent/control/execution.rs
+2 -2
View File
@@ -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());
codex-rs/core/src/agent/control/legacy.rs
+2 -2
View File
@@ -6,8 +6,8 @@ impl AgentControl {
pub(crate) async fn shutdown_live_agent(&self, agent_id: ThreadId) -> CodexResult<String> {
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 {
codex-rs/core/src/agent/control/residency.rs
+2 -7
View File
@@ -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)]
codex-rs/core/src/agent/control/residency_tests.rs
+3 -5
View File
@@ -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;
}
codex-rs/core/src/agent/control/spawn.rs
+5 -16
View File
@@ -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<String> =
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
codex-rs/core/src/agent/control_tests.rs
+29 -91
View File
@@ -242,7 +242,6 @@ async fn wait_for_subagent_notification(parent_thread: &Arc<CodexThread>) -> 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<CodexThread>, 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
codex-rs/core/src/codex_delegate.rs
+64 -63
View File
@@ -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<McpToolApprovalMetadata>,
}
/// 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<InitialHistory>,
) -> Result<Codex, CodexErr> {
) -> Result<(Arc<Session>, 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::<String, PendingMcpInvocation>::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<Value>,
initial_history: Option<InitialHistory>,
) -> Result<Codex, CodexErr> {
) -> Result<(Arc<Session>, 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<Codex>,
io: Arc<SessionIo>,
session: Arc<Session>,
tx_sub: Sender<Event>,
parent_session: Arc<Session>,
parent_ctx: Arc<TurnContext>,
@@ -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<Event>,
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<Codex>,
io: Arc<SessionIo>,
rx_ops: Receiver<Submission>,
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<Session>,
parent_ctx: &Arc<TurnContext>,
@@ -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<Session>,
parent_ctx: &Arc<TurnContext>,
@@ -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<Session>,
parent_ctx: &Arc<TurnContext>,
@@ -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<Session>,
parent_ctx: &Arc<TurnContext>,
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,
codex-rs/core/src/codex_delegate_tests.rs
+11 -15
View File
@@ -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,
codex-rs/core/src/codex_thread.rs
+58 -77
View File
@@ -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<Session>,
pub(crate) io: SessionIo,
pub(crate) session_source: SessionSource,
session_configured: SessionConfiguredEvent,
rollout_path: Option<PathBuf>,
@@ -184,13 +186,15 @@ pub struct BackgroundTerminalInfo {
/// (formerly called a conversation) in Codex.
impl CodexThread {
pub(crate) fn new(
codex: Codex,
session: Arc<Session>,
io: SessionIo,
session_configured: SessionConfiguredEvent,
rollout_path: Option<PathBuf>,
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<String> {
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<W3cTraceContext>,
) -> CodexResult<String> {
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<W3cTraceContext>,
client_user_message_id: Option<String>,
) -> CodexResult<String> {
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<String>,
responsesapi_client_metadata: Option<HashMap<String, String>>,
) -> Result<String, SteerInputError> {
self.codex
self.session
.steer_input(
input,
additional_context,
@@ -308,7 +307,7 @@ impl CodexThread {
&self,
items: Vec<ResponseItem>,
) -> Result<(), Vec<ResponseItem>> {
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<ResponseItem>,
) -> 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<String>,
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<ThreadConfigSnapshot> {
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<Event> {
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<BackgroundTerminalInfo> {
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<AgentStatus> {
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<TokenUsageInfo> {
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<PathBuf> {
self.codex
.session
self.session
.guardian_review_session
.trunk_rollout_path()
.await
@@ -517,7 +506,6 @@ impl CodexThread {
include_archived: bool,
) -> ThreadStoreResult<StoredThreadHistory> {
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<StoredThread> {
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<StoredThread> {
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<StateDbHandle> {
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<PathUri> {
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<crate::config::Config> {
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<crate::session::McpRuntimeSnapshot> {
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<MultiAgentVersion> {
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<TurnEnvironmentSelection> {
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<i64> {
@@ -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)
codex-rs/core/src/guardian/review_session.rs
+35 -35
View File
@@ -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<Session>,
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<Self>) {
@@ -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<PathBuf> {
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<Session>, 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<Session>, 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<T>(
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());
}
}
codex-rs/core/src/prompt_debug.rs
+1 -1
View File
@@ -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;
codex-rs/core/src/session/mod.rs
+73 -112
View File
@@ -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<Submission>,
pub(crate) rx_event: Receiver<Event>,
// Last known status of the agent.
pub(crate) agent_status: watch::Receiver<AgentStatus>,
pub(crate) session: Arc<Session>,
// 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<BoxFuture<'static, ()>>;
/// 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<CodexSpawnOk> {
impl Session {
/// Spawn and initialize a new session.
pub(crate) async fn spawn(args: SessionSpawnArgs) -> CodexResult<(Arc<Self>, 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<CodexSpawnOk> {
let CodexSpawnArgs {
async fn spawn_internal(args: SessionSpawnArgs) -> CodexResult<(Arc<Self>, 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<String> {
pub(crate) async fn submit(&self, op: Op) -> CodexResult<String> {
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<W3cTraceContext>,
@@ -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<W3cTraceContext>,
@@ -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<Event> {
pub(crate) async fn next_event(&self) -> CodexResult<Event> {
let event = self
.rx_event
.recv()
@@ -823,80 +807,9 @@ impl Codex {
Ok(event)
}
pub async fn steer_input(
&self,
input: Vec<UserInput>,
additional_context: BTreeMap<String, AdditionalContextEntry>,
expected_turn_id: Option<&str>,
client_user_message_id: Option<String>,
responsesapi_client_metadata: Option<HashMap<String, String>>,
) -> Result<String, SteerInputError> {
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<String>,
app_server_client_version: Option<String>,
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<PathUri> {
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<TurnEnvironmentSelection> {
let state = self.session.state.lock().await;
state
.session_configuration
.environment_selections()
.to_vec()
}
pub(crate) fn state_db(&self) -> Option<state_db::StateDbHandle> {
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<Self>,
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<String>,
app_server_client_version: Option<String>,
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<PathUri> {
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<TurnEnvironmentSelection> {
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,
codex-rs/core/src/session/tests.rs
+41 -55
View File
@@ -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");
codex-rs/core/src/session/tests/guardian_tests.rs
+3 -4
View File
@@ -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);
}
codex-rs/core/src/tasks/review.rs
+1 -1
View File
@@ -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(
codex-rs/core/src/thread_manager.rs
+14 -14
View File
@@ -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<Session>,
io: SessionIo,
session_source: SessionSource,
) -> CodexResult<NewThread> {
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)
}
}
codex-rs/core/src/thread_manager_tests.rs
+3 -7
View File
@@ -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
codex-rs/core/src/tools/handlers/multi_agents_tests.rs
+11 -15
View File
@@ -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()})),
))