From 52b45262754bbe8712c9d385e7a5705b406fdec9 Mon Sep 17 00:00:00 2001 From: rhan-oai Date: Mon, 2 Nov 2026 22:34:04 +0000 Subject: [PATCH] Allow hosted child threads to inherit a root session ID (#59000) ## What changed Expose `InheritedSessionId` so a host can supply a root session ID through `ExtensionDataInit` when starting a fresh child thread. The child retains its distinct thread ID, and the inherited session ID is persisted in rollout metadata and restored on resume. Reject inherited session IDs for resumed, forked, or native agent threads, and reject an inherited session ID that equals the child thread ID. ## Testing Add an integration test covering distinct child identity, inherited session metadata, and restoration after shutdown and resume. GitOrigin-RevId: e3462433f9d8a0bba24251ee2b6d5329557ddb31 --- codex-rs/core/src/agent/host/mod.rs | 1 + codex-rs/core/src/agent/host/types.rs | 10 +++ codex-rs/core/src/lib.rs | 1 + codex-rs/core/src/session/session.rs | 26 ++++++-- codex-rs/core/src/thread_manager.rs | 10 +++ .../core/tests/suite/agent_control_tests.rs | 65 +++++++++++++++++++ 6 files changed, 106 insertions(+), 7 deletions(-) diff --git a/codex-rs/core/src/agent/host/mod.rs b/codex-rs/core/src/agent/host/mod.rs index b256ca4830..aba89129e3 100644 --- a/codex-rs/core/src/agent/host/mod.rs +++ b/codex-rs/core/src/agent/host/mod.rs @@ -18,6 +18,7 @@ pub use mcp_metadata::AgentMcpMetadataRecipient; pub use ownership::AgentRuntimeOwnership; pub(crate) use ownership::AgentRuntimeServices; pub use types::HostThreadSpawnReservation; +pub use types::InheritedSessionId; pub use types::PreparedChildContext; pub(crate) use types::PreparedChildInheritance; pub use types::RestrictRootCapabilities; diff --git a/codex-rs/core/src/agent/host/types.rs b/codex-rs/core/src/agent/host/types.rs index 7ae6353ed9..6c1301d9f2 100644 --- a/codex-rs/core/src/agent/host/types.rs +++ b/codex-rs/core/src/agent/host/types.rs @@ -7,6 +7,7 @@ use crate::agents_md_manager::SessionInstructions; use crate::exec_policy::ExecPolicyManager; use codex_execpolicy::Policy; use codex_execpolicy::RequirementsExecPolicy; +use codex_protocol::SessionId; use codex_protocol::ThreadId; use codex_protocol::capabilities::SelectedCapabilityRoot; use codex_protocol::config_types::ApprovalsReviewer; @@ -31,6 +32,15 @@ use std::sync::Arc; #[derive(Clone, Copy, Debug)] pub struct RestrictRootCapabilities; +/// The root's session ID, supplied by a trusted host when starting a fresh child. +/// +/// Insert this into the thread's startup [`codex_extension_api::ExtensionDataInit`]. +/// The child keeps its own distinct thread ID. Resumed threads take their session +/// ID from persisted metadata, and native children take it from their controller. +/// To restrict root capabilities, also supply [`RestrictRootCapabilities`]. +#[derive(Clone, Copy, Debug)] +pub struct InheritedSessionId(pub SessionId); + /// A fresh child ID reserved in the host's tree. /// The host commits membership, serializes startup attempts and deduplicates retries. /// Core derives the parent, session and path from the invocation and child request. diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index 82a938d170..884f61fefb 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -164,6 +164,7 @@ pub use agent::host::AgentUsageAccounting; pub use agent::host::AgentUsageReport; pub use agent::host::HostThreadSpawnReservation; pub use agent::host::HostToolInvocationLookup; +pub use agent::host::InheritedSessionId; pub use agent::host::PreparedChildContext; pub use agent::host::RestrictRootCapabilities; pub use agent::host::ToolInvocationContext; diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 2836e9d21b..ae17b9a009 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -9,6 +9,7 @@ use super::step_settings::StepSettingsUpdate; use super::*; use crate::agent::api::AgentControl; use crate::agent::host::AgentRuntimeServices; +use crate::agent::host::InheritedSessionId; use crate::agents_md_manager::AgentsMdManager; use crate::agents_md_manager::SessionInstructions; use crate::config::ConstraintError; @@ -954,6 +955,15 @@ impl Session { )); } }; + let inherited_session_id = thread_extension_init + .remove::() + .map(|inherited| inherited.0); + if inherited_session_id == Some(SessionId::from(thread_id)) { + return Err(CodexErr::InvalidRequest( + "an inherited agent session ID must differ from the child thread ID".to_string(), + ) + .into()); + } let isolation = thread_extension_init .get::() .map(|policy| *policy) @@ -996,13 +1006,15 @@ impl Session { || *session_id != SessionId::from(thread_id) }); // session_id is equal to the root thread's ID. - let session_id = resumed_session_id.unwrap_or_else(|| { - if session_configuration.session_source.is_non_root_agent() { - agent_control.control().identity() - } else { - SessionId::from(thread_id) - } - }); + let session_id = resumed_session_id + .or(inherited_session_id) + .unwrap_or_else(|| { + if session_configuration.session_source.is_non_root_agent() { + agent_control.control().identity() + } else { + SessionId::from(thread_id) + } + }); let initial_auto_compact_window_ids = AutoCompactWindowIds::new_initial(); let restore_child_window = matches!(&initial_history, InitialHistory::Forked(_)) && session_configuration.session_source.is_non_root_agent() diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index 966e666b97..f07c7d18a0 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -15,6 +15,7 @@ use crate::agent::LocalAgentControl; use crate::agent::api::AgentConfigUpdate; use crate::agent::api::AgentControl; use crate::agent::control::AgentControlInit; +use crate::agent::host::InheritedSessionId; use crate::agents_md_manager::SessionInstructions; use crate::attestation::AttestationProvider; use crate::codex_thread::CodexThread; @@ -2287,6 +2288,15 @@ impl ThreadManagerState { } let inherited_environments = captured_environments.or(inherited_environments); let session_source = session_source.unwrap_or_else(|| self.session_source.clone()); + if thread_extension_init.get::().is_some() + && (!matches!(&initial_history, InitialHistory::New) + || session_source.is_non_root_agent()) + { + return Err(CodexErr::InvalidRequest( + "an inherited agent session ID requires a fresh thread started directly by the host" + .to_string(), + )); + } // Older callers and saved reviewers identify isolation through their source. // New internal callers supply an explicit runtime policy before startup. let isolation = thread_extension_init diff --git a/codex-rs/core/tests/suite/agent_control_tests.rs b/codex-rs/core/tests/suite/agent_control_tests.rs index c6da8fa8f0..e7da4dd829 100644 --- a/codex-rs/core/tests/suite/agent_control_tests.rs +++ b/codex-rs/core/tests/suite/agent_control_tests.rs @@ -13,6 +13,7 @@ use codex_core::AgentTurnOutcome; use codex_core::DeliveryReceipt; use codex_core::ForkSnapshot; use codex_core::GuardianRootSnapshot; +use codex_core::InheritedSessionId; use codex_core::LiveAgent; use codex_core::PendingAgentMessage; use codex_core::RolloutBudgetReminder; @@ -526,6 +527,70 @@ async fn host_factory_follows_thread_lifecycle() -> anyhow::Result<()> { Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn fresh_hosted_child_persists_inherited_session_id() -> anyhow::Result<()> { + skip_if_no_network!(Ok(())); + let server = responses::start_mock_server().await; + let root = test_codex().build_with_auto_env(&server).await?; + let root_session_id = root.session_configured.session_id; + let manager = &root.thread_manager; + let options = || StartThreadOptions { + session_source: Some(SessionSource::Custom("host".to_string())), + environments: Some(vec![root.executor_environment().request()]), + ..StartThreadOptions::new(root.config.clone()) + }; + let mut child_options = options(); + child_options + .thread_extension_init + .insert(InheritedSessionId(root_session_id)); + let child = manager.start_thread(child_options).await?; + assert_ne!(SessionId::from(child.thread_id), root_session_id); + assert_eq!( + ( + child.session_configured.session_id, + child.session_configured.parent_thread_id + ), + (root_session_id, None) + ); + child.thread.ensure_rollout_materialized().await; + child.thread.flush_rollout().await?; + let history = child + .thread + .load_history(/*include_archived*/ false) + .await?; + let Some(RolloutItem::SessionMeta(meta)) = history.items.first() else { + anyhow::bail!("expected child metadata as the first rollout item"); + }; + assert_eq!( + (meta.meta.id, meta.meta.session_id), + (child.thread_id, root_session_id) + ); + child.thread.shutdown_and_wait().await?; + let resumed = manager + .start_thread(StartThreadOptions { + initial_history: InitialHistory::Resumed(ResumedHistory { + history_revision: None, + conversation_id: child.thread_id, + history: Arc::new(history.items), + rollout_path: child.thread.rollout_path(), + }), + ..options() + }) + .await?; + assert!(!Arc::ptr_eq(&resumed.thread, &child.thread)); + assert_eq!( + ( + resumed.thread_id, + resumed.session_configured.session_id, + resumed.session_configured.parent_thread_id + ), + (child.thread_id, root_session_id, None) + ); + resumed.thread.shutdown_and_wait().await?; + root.codex.shutdown_and_wait().await?; + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn collaboration_tools_dispatch_to_the_host_controller() -> anyhow::Result<()> { skip_if_no_network!(Ok(()));