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(()));