From 40cf168567aa3cb3e0cf6eadef82e6eca2cd16c1 Mon Sep 17 00:00:00 2001 From: rhan-oai Date: Mon, 2 Nov 2026 20:54:47 +0000 Subject: [PATCH] Allow hosts to restrict a thread's root-only capabilities (#58984) ## Why Hosts need to apply subagent-style capability restrictions without changing a thread's session attribution, agent lineage, or permission profile. ## What changed Add the public `RestrictRootCapabilities` marker to thread startup extension data. Capture the restriction once at startup and apply it alongside the restrictions native subagents already receive. Hosts must supply the marker again when cold-resuming a restricted thread. Use the restriction to gate user-input and plugin-install requests, async user messaging, legacy MCP tool approvals, and Full Access MCP form input. Exclude inherited user messages from Guardian's retained authorization context for restricted threads, including after cold resume. ## Testing Add coverage for host-restricted tool registration, MCP approval rejection, Full Access form-input handling, and inherited instruction authorization after cold resume. Update handler tests to exercise the shared restriction flag. GitOrigin-RevId: 4d6d7102a20dc747431c7511a74a204c773e8420 --- codex-rs/core/src/agent/host/mod.rs | 1 + codex-rs/core/src/agent/host/types.rs | 16 ++++- codex-rs/core/src/context_manager/history.rs | 8 ++- .../core/src/context_manager/history_tests.rs | 2 +- codex-rs/core/src/lib.rs | 1 + codex-rs/core/src/mcp_tool_call.rs | 2 +- .../src/session/guardian_checkpoint_tests.rs | 5 +- codex-rs/core/src/session/mod.rs | 14 ++++ .../src/session/rollout_reconstruction.rs | 2 +- .../session/rollout_reconstruction_tests.rs | 2 +- codex-rs/core/src/session/session.rs | 4 +- codex-rs/core/src/session/tests.rs | 19 ++++-- codex-rs/core/src/session/turn_context.rs | 4 ++ codex-rs/core/src/session/world_state.rs | 15 ++--- codex-rs/core/src/thread_manager.rs | 14 +--- .../tools/handlers/request_plugin_install.rs | 2 +- .../handlers/request_plugin_install_tests.rs | 15 +---- .../src/tools/handlers/request_user_input.rs | 2 +- .../handlers/request_user_input_tests.rs | 15 +---- codex-rs/core/src/tools/spec_plan.rs | 4 +- .../tests/suite/guardian_retained_context.rs | 67 +++++++++++++++++++ .../core/tests/suite/mcp_auth_elicitation.rs | 61 +++++++++++------ .../tests/suite/mcp_subagent_elicitation.rs | 66 ++++++++++++------ .../tests/suite/request_user_input_async.rs | 48 ++++++++----- 24 files changed, 266 insertions(+), 123 deletions(-) diff --git a/codex-rs/core/src/agent/host/mod.rs b/codex-rs/core/src/agent/host/mod.rs index 9a88c60c77..b256ca4830 100644 --- a/codex-rs/core/src/agent/host/mod.rs +++ b/codex-rs/core/src/agent/host/mod.rs @@ -20,6 +20,7 @@ pub(crate) use ownership::AgentRuntimeServices; pub use types::HostThreadSpawnReservation; pub use types::PreparedChildContext; pub(crate) use types::PreparedChildInheritance; +pub use types::RestrictRootCapabilities; pub use usage::AgentUsageAccounting; pub use usage::AgentUsageReport; pub use usage::UsageAccountingDecision; diff --git a/codex-rs/core/src/agent/host/types.rs b/codex-rs/core/src/agent/host/types.rs index f96d94a7b1..7ae6353ed9 100644 --- a/codex-rs/core/src/agent/host/types.rs +++ b/codex-rs/core/src/agent/host/types.rs @@ -1,4 +1,4 @@ -//! Prepared inputs and host reservations for native child startup. +//! Host inputs for child startup and root-only capability restrictions. use super::admission::AgentTurnAdmission; use super::usage::AgentUsageAccounting; @@ -17,6 +17,20 @@ use codex_protocol::protocol::TurnEnvironmentRequest; use codex_utils_absolute_path::AbsolutePathBuf; use std::sync::Arc; +/// Restrict a thread's root-only capabilities as for a native subagent. +/// +/// A trusted host supplies this marker in the thread's startup [`codex_extension_api::ExtensionDataInit`], +/// including when cold-resuming the thread. Core captures it once at startup. +/// Embeddings can read it during extension startup and MCP projection to apply their own +/// root-only restrictions. Session attribution, agent lineage and the thread's permission +/// profile come from their usual startup inputs. Native subagents are always restricted. +/// +/// For agent forks, hosts must also supply history with native agent-fork provenance. +/// This marks copied messages as inherited and excludes the source's Guardian checkpoints. +/// Obtaining current root evidence for Guardian is a separate host responsibility. +#[derive(Clone, Copy, Debug)] +pub struct RestrictRootCapabilities; + /// 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/context_manager/history.rs b/codex-rs/core/src/context_manager/history.rs index 621d3b7ee1..523730fd86 100644 --- a/codex-rs/core/src/context_manager/history.rs +++ b/codex-rs/core/src/context_manager/history.rs @@ -60,7 +60,6 @@ use codex_protocol::models::ImageReference; use codex_protocol::models::ResponseItem; use codex_protocol::openai_models::InputModality; use codex_protocol::protocol::InterAgentCommunication; -use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::TokenUsage; use codex_protocol::protocol::TokenUsageInfo; use codex_protocol::protocol::TurnContextItem; @@ -282,10 +281,13 @@ impl ContextManager { &self.retained_context } - pub(crate) fn for_session(source: &SessionSource, features: &ManagedFeatures) -> Self { + pub(crate) fn for_session( + root_capabilities_restricted: bool, + features: &ManagedFeatures, + ) -> Self { let independent = !features.enabled(Feature::GuardianReuseParentCompaction); Self { - retain_inherited_user_messages: !source.is_non_root_agent(), + retain_inherited_user_messages: !root_capabilities_restricted, guardian_review_mode: if independent { GuardianContextMode::Independent } else { diff --git a/codex-rs/core/src/context_manager/history_tests.rs b/codex-rs/core/src/context_manager/history_tests.rs index 42895a8db2..5c3f1330ad 100644 --- a/codex-rs/core/src/context_manager/history_tests.rs +++ b/codex-rs/core/src/context_manager/history_tests.rs @@ -133,7 +133,7 @@ fn inter_agent_assistant_msg(text: &str) -> ResponseItem { fn root_history() -> ContextManager { ContextManager::for_session( - &SessionSource::Cli, + /*root_capabilities_restricted*/ false, &ManagedFeatures::from(codex_features::Features::with_defaults()), ) } diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index 7e8fb8763b..82a938d170 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -165,6 +165,7 @@ pub use agent::host::AgentUsageReport; pub use agent::host::HostThreadSpawnReservation; pub use agent::host::HostToolInvocationLookup; pub use agent::host::PreparedChildContext; +pub use agent::host::RestrictRootCapabilities; pub use agent::host::ToolInvocationContext; pub use agent::host::UsageAccountingDecision; pub use thread_spawn::ThreadCreateRequest; diff --git a/codex-rs/core/src/mcp_tool_call.rs b/codex-rs/core/src/mcp_tool_call.rs index 5abbbce2a3..453bc96c4f 100644 --- a/codex-rs/core/src/mcp_tool_call.rs +++ b/codex-rs/core/src/mcp_tool_call.rs @@ -1779,7 +1779,7 @@ pub(crate) async fn request_mcp_tool_user_approval( cancellation_reason, ) } else { - if turn_context.session_source.is_non_root_agent() { + if turn_context.root_capabilities_restricted { return ApprovalResolution { decision: ReviewDecision::denied(concat!( "MCP tool approval via request_user_input requires the root thread. ", diff --git a/codex-rs/core/src/session/guardian_checkpoint_tests.rs b/codex-rs/core/src/session/guardian_checkpoint_tests.rs index edac477801..0f7061ff6e 100644 --- a/codex-rs/core/src/session/guardian_checkpoint_tests.rs +++ b/codex-rs/core/src/session/guardian_checkpoint_tests.rs @@ -4,7 +4,6 @@ use crate::context_manager::ContextManager; use crate::session::tests::make_session_and_context; use codex_history::InitialHistory; use codex_protocol::models::ResponseItem; -use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::ThreadHistoryMode; use codex_protocol::protocol::TokenUsage; use pretty_assertions::assert_eq; @@ -36,7 +35,7 @@ async fn guardian_checkpoint_preserves_live_context_without_storage( { let mut state = session.state.lock().await; state.history = ContextManager::for_session( - &SessionSource::default(), + /*root_capabilities_restricted*/ false, &crate::config::ManagedFeatures::from(codex_features::Features::with_defaults()), ); state @@ -107,7 +106,7 @@ async fn guardian_checkpoint_preserves_live_context_without_storage( let (fork, _) = make_session_and_context().await; fork.state.lock().await.session_configuration.history_mode = history_mode; fork.state.lock().await.history = ContextManager::for_session( - &SessionSource::default(), + /*root_capabilities_restricted*/ false, &crate::config::ManagedFeatures::from(codex_features::Features::with_defaults()), ); fork.record_initial_history(InitialHistory::Forked(items)) diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index ff0872efa3..9ebd91ad9e 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -12,6 +12,7 @@ use std::sync::atomic::AtomicU64; use std::time::SystemTime; use std::time::UNIX_EPOCH; +use crate::RestrictRootCapabilities; use crate::agent::AgentStatus; use crate::agent::agent_status_from_event; use crate::agent::api::AgentConfigUpdate; @@ -372,6 +373,7 @@ use codex_protocol::config_types::ReasoningSummary as ReasoningSummaryConfig; use codex_protocol::config_types::Settings; use codex_protocol::config_types::WindowsSandboxLevel; use codex_protocol::mcp::ClientMcpExtensions; +use codex_protocol::mcp::OPENAI_STANDARD_FORM_INPUT_EXTENSION_ID; use codex_protocol::models::LocalImagePreparation; use codex_protocol::models::ResponseInputItem; use codex_protocol::models::ResponseItem; @@ -941,6 +943,13 @@ impl Session { &model_info, ); let storage_originator = AuthStorageOriginator::from_client_name(&originator); + let root_capabilities_restricted = session_source.is_non_root_agent() + || thread_extension_init + .get::() + .is_some(); + let enable_full_access_form_input = !root_capabilities_restricted + && matches!(thread_source.as_ref(), Some(ThreadSource::User)) + && client_mcp_extensions.contains(OPENAI_STANDARD_FORM_INPUT_EXTENSION_ID); let session_configuration = SessionConfiguration { turn_extension_init, provider: create_model_provider( @@ -976,6 +985,7 @@ impl Session { app_server_client_version: None, trusted_guardian_reviewer, session_source, + root_capabilities_restricted, history_mode, forked_from_thread_id, parent_thread_id, @@ -1077,6 +1087,10 @@ impl Session { let _ = startup.io.set(io.clone()); startup.persistence.lock().await.commit(); } + // A required MCP server must not block startup waiting for form input. + if enable_full_access_form_input { + session.services.mcp_runtime.enable_full_access_form_input(); + } Ok((session, io)) } } diff --git a/codex-rs/core/src/session/rollout_reconstruction.rs b/codex-rs/core/src/session/rollout_reconstruction.rs index 1205827911..130d5e9ddc 100644 --- a/codex-rs/core/src/session/rollout_reconstruction.rs +++ b/codex-rs/core/src/session/rollout_reconstruction.rs @@ -195,7 +195,7 @@ impl Session { rollout_items, turn_context.history_mode, ContextManager::for_session( - &turn_context.session_source, + turn_context.root_capabilities_restricted, &turn_context.config.features, ), turn_context.model_info().truncation_policy.into(), diff --git a/codex-rs/core/src/session/rollout_reconstruction_tests.rs b/codex-rs/core/src/session/rollout_reconstruction_tests.rs index c102a49743..ebc312eafa 100644 --- a/codex-rs/core/src/session/rollout_reconstruction_tests.rs +++ b/codex-rs/core/src/session/rollout_reconstruction_tests.rs @@ -126,7 +126,7 @@ async fn sender_context_follows_its_delivery_through_checkpoint_and_rollback() { let (session, turn_context) = make_session_and_context().await; let mut live = ContextManager::for_session( - &SessionSource::default(), + /*root_capabilities_restricted*/ false, &crate::config::ManagedFeatures::from(codex_features::Features::with_defaults()), ); let mut items = Vec::new(); diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 9115427c9d..2836e9d21b 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -159,6 +159,8 @@ pub(crate) struct SessionConfiguration { pub(super) trusted_guardian_reviewer: bool, /// Source of the session (cli, vscode, exec, mcp, ...) pub(super) session_source: SessionSource, + /// Whether native lineage or the trusted host restricts this session's root-only capabilities. + pub(super) root_capabilities_restricted: bool, /// Persisted thread history contract selected when this thread was created. pub(super) history_mode: ThreadHistoryMode, /// Immediate history source copied into this thread, when this thread was forked. @@ -1574,7 +1576,7 @@ impl Session { session_configuration.clone(), initial_auto_compact_window_ids, ContextManager::for_session( - &session_configuration.session_source, + session_configuration.root_capabilities_restricted, &config.features, ), ); diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index d3190c789e..5711ec9207 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -4661,6 +4661,7 @@ async fn set_rate_limits_retains_previous_credits() { app_server_client_version: None, trusted_guardian_reviewer: false, session_source: SessionSource::Exec, + root_capabilities_restricted: false, history_mode: Default::default(), forked_from_thread_id: None, parent_thread_id: None, @@ -4784,6 +4785,7 @@ async fn set_rate_limits_updates_plan_type_when_present() { app_server_client_version: None, trusted_guardian_reviewer: false, session_source: SessionSource::Exec, + root_capabilities_restricted: false, history_mode: Default::default(), forked_from_thread_id: None, parent_thread_id: None, @@ -5409,6 +5411,7 @@ pub(crate) async fn make_session_configuration_for_tests() -> SessionConfigurati app_server_client_version: None, trusted_guardian_reviewer: false, session_source: SessionSource::Exec, + root_capabilities_restricted: false, history_mode: Default::default(), forked_from_thread_id: None, parent_thread_id: None, @@ -6987,6 +6990,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { app_server_client_version: None, trusted_guardian_reviewer: false, session_source: SessionSource::Exec, + root_capabilities_restricted: false, history_mode: Default::default(), forked_from_thread_id: None, parent_thread_id: None, @@ -7003,8 +7007,10 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { ); let mut state = SessionState::new(session_configuration.clone()); - state.history = - ContextManager::for_session(&session_configuration.session_source, &config.features); + state.history = ContextManager::for_session( + session_configuration.root_capabilities_restricted, + &config.features, + ); let (environment_manager, resolved_environments) = resolved_environments_for_configuration(&session_configuration, &default_environments) .await; @@ -7309,6 +7315,7 @@ async fn make_session_with_config_and_rx( app_server_client_version: None, trusted_guardian_reviewer: false, session_source: SessionSource::Exec, + root_capabilities_restricted: false, history_mode: Default::default(), forked_from_thread_id: None, parent_thread_id: None, @@ -7446,6 +7453,7 @@ async fn make_session_with_history_source_and_agent_control_and_rx( app_server_client_version: None, trusted_guardian_reviewer: false, session_source: session_source.clone(), + root_capabilities_restricted: session_source.is_non_root_agent(), history_mode: Default::default(), forked_from_thread_id: None, parent_thread_id: None, @@ -9312,6 +9320,7 @@ where app_server_client_version: None, trusted_guardian_reviewer: false, session_source: SessionSource::Exec, + root_capabilities_restricted: false, history_mode: Default::default(), forked_from_thread_id: None, parent_thread_id: None, @@ -9328,8 +9337,10 @@ where ); let mut state = SessionState::new(session_configuration.clone()); - state.history = - ContextManager::for_session(&session_configuration.session_source, &config.features); + state.history = ContextManager::for_session( + session_configuration.root_capabilities_restricted, + &config.features, + ); let (environment_manager, resolved_turn_environments) = resolved_environments_for_configuration(&session_configuration, &default_environments) .await; diff --git a/codex-rs/core/src/session/turn_context.rs b/codex-rs/core/src/session/turn_context.rs index c08333d350..aabc769ab7 100644 --- a/codex-rs/core/src/session/turn_context.rs +++ b/codex-rs/core/src/session/turn_context.rs @@ -347,6 +347,7 @@ pub struct TurnContext { pub(crate) session_telemetry: SessionTelemetry, pub(crate) provider: SharedModelProvider, pub(crate) session_source: SessionSource, + pub(crate) root_capabilities_restricted: bool, pub(crate) history_mode: ThreadHistoryMode, pub(crate) parent_thread_id: Option, pub(crate) originator: String, @@ -437,6 +438,7 @@ impl TurnContext { session_telemetry, provider: self.provider.clone(), session_source: self.session_source.clone(), + root_capabilities_restricted: self.root_capabilities_restricted, history_mode: self.history_mode, parent_thread_id: self.parent_thread_id, originator: self.originator.clone(), @@ -848,6 +850,7 @@ impl TurnContext { session_telemetry, provider: self.provider.clone(), session_source: self.session_source.clone(), + root_capabilities_restricted: self.root_capabilities_restricted, history_mode: self.history_mode, parent_thread_id: self.parent_thread_id, originator: self.originator.clone(), @@ -1132,6 +1135,7 @@ impl Session { session_telemetry: session_telemetry_for_context, provider, session_source, + root_capabilities_restricted: session_configuration.root_capabilities_restricted, history_mode: session_configuration.history_mode, parent_thread_id: session_configuration.parent_thread_id, originator: session_configuration.originator.clone(), diff --git a/codex-rs/core/src/session/world_state.rs b/codex-rs/core/src/session/world_state.rs index a3eff58391..cbf840b611 100644 --- a/codex-rs/core/src/session/world_state.rs +++ b/codex-rs/core/src/session/world_state.rs @@ -282,14 +282,13 @@ impl Session { )); } if !crate::guardian::is_basic_session_source(&turn_context.session_source) { - let send_user_message_async_available = - !turn_context.session_source.is_non_root_agent() - && step_context - .settings - .model_info - .experimental_supported_tools - .iter() - .any(|tool| tool == "send_user_message_async"); + let send_user_message_async_available = !turn_context.root_capabilities_restricted + && step_context + .settings + .model_info + .experimental_supported_tools + .iter() + .any(|tool| tool == "send_user_message_async"); world_state.add_section(PersistentModeState::new( turn_context.config.features.persistent_execution_enabled( step_context.settings.effective_reasoning_effort().as_ref(), diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index a30a2b8ab7..966e666b97 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -81,7 +81,6 @@ use codex_protocol::config_types::CollaborationModeMask; use codex_protocol::error::CodexErr; use codex_protocol::error::Result as CodexResult; use codex_protocol::mcp::ClientMcpExtensions; -use codex_protocol::mcp::OPENAI_STANDARD_FORM_INPUT_EXTENSION_ID; use codex_protocol::openai_models::ModelPreset; use codex_protocol::protocol::Event; use codex_protocol::protocol::EventMsg; @@ -2562,7 +2561,7 @@ impl ThreadManagerState { }, forked_from_thread_id, parent_thread_id, - thread_source: thread_source.clone(), + thread_source, originator: startup_context.originator, agent_control, dynamic_tools, @@ -2618,17 +2617,6 @@ impl ThreadManagerState { warn!(%thread_id, %source_thread_id, %error, "failed to copy fork attachments"); } } - // Enable Full Access form input only after session startup so a required MCP server cannot - // block startup while waiting for form input. - if session - .services - .client_mcp_extensions - .contains(OPENAI_STANDARD_FORM_INPUT_EXTENSION_ID) - && matches!(thread_source.as_ref(), Some(ThreadSource::User)) - && !tracked_session_source.is_non_root_agent() - { - session.services.mcp_runtime.enable_full_access_form_input(); - } let finalize_result = self .finalize_thread_spawn(session, io, tracked_session_source, initial_host_config) .await; diff --git a/codex-rs/core/src/tools/handlers/request_plugin_install.rs b/codex-rs/core/src/tools/handlers/request_plugin_install.rs index 271c12f30d..281f938b5c 100644 --- a/codex-rs/core/src/tools/handlers/request_plugin_install.rs +++ b/codex-rs/core/src/tools/handlers/request_plugin_install.rs @@ -114,7 +114,7 @@ impl RequestPluginInstallHandler { } }; - if turn.session_source.is_non_root_agent() { + if turn.root_capabilities_restricted { return Err(FunctionCallError::RespondToModel( "request_plugin_install can only be used by the root thread".to_string(), )); diff --git a/codex-rs/core/src/tools/handlers/request_plugin_install_tests.rs b/codex-rs/core/src/tools/handlers/request_plugin_install_tests.rs index f6af475f90..133a5c2ccd 100644 --- a/codex-rs/core/src/tools/handlers/request_plugin_install_tests.rs +++ b/codex-rs/core/src/tools/handlers/request_plugin_install_tests.rs @@ -16,9 +16,6 @@ use codex_config::types::ToolSuggestDiscoverableType; use codex_core_plugins::PluginInstallRequest; use codex_core_plugins::startup_sync::curated_plugins_repo_path; use codex_login::test_support::auth_manager_from_optional_auth; -use codex_protocol::ThreadId; -use codex_protocol::protocol::SessionSource; -use codex_protocol::protocol::SubAgentSource; use codex_rmcp_client::ElicitationResponse; use codex_tools::DiscoverablePluginInfo; use codex_utils_absolute_path::AbsolutePathBuf; @@ -33,15 +30,9 @@ use tokio::sync::Mutex; #[test_case(ToolSuggestPresentation::ListTool; "legacy install")] #[test_case(ToolSuggestPresentation::RecommendationContext; "recommended plugin")] #[tokio::test] -async fn request_plugin_install_rejects_subagent_threads(presentation: ToolSuggestPresentation) { +async fn request_plugin_install_rejects_restricted_threads(presentation: ToolSuggestPresentation) { let (session, mut turn) = make_session_and_context().await; - turn.session_source = SessionSource::SubAgent(SubAgentSource::ThreadSpawn { - parent_thread_id: ThreadId::new(), - depth: 1, - agent_path: None, - agent_nickname: None, - agent_role: None, - }); + turn.root_capabilities_restricted = true; let turn = Arc::new(turn); let arguments = match presentation { ToolSuggestPresentation::ListTool => json!({ @@ -75,7 +66,7 @@ async fn request_plugin_install_rejects_subagent_threads(presentation: ToolSugge .await; let Err(err) = result else { - panic!("subagent request_plugin_install should fail"); + panic!("restricted request_plugin_install should fail"); }; assert_eq!( err, diff --git a/codex-rs/core/src/tools/handlers/request_user_input.rs b/codex-rs/core/src/tools/handlers/request_user_input.rs index d55eca0297..0bac1b16fa 100644 --- a/codex-rs/core/src/tools/handlers/request_user_input.rs +++ b/codex-rs/core/src/tools/handlers/request_user_input.rs @@ -64,7 +64,7 @@ impl RequestUserInputHandler { } }; - if turn.session_source.is_non_root_agent() { + if turn.root_capabilities_restricted { return Err(FunctionCallError::RespondToModel( "request_user_input can only be used by the root thread".to_string(), )); diff --git a/codex-rs/core/src/tools/handlers/request_user_input_tests.rs b/codex-rs/core/src/tools/handlers/request_user_input_tests.rs index ecb5ed51bd..30ec27e68f 100644 --- a/codex-rs/core/src/tools/handlers/request_user_input_tests.rs +++ b/codex-rs/core/src/tools/handlers/request_user_input_tests.rs @@ -11,10 +11,7 @@ use crate::tools::context::ToolInvocation; use crate::tools::context::ToolPayload; use crate::turn_diff_tracker::TurnDiffTracker; use codex_guardian_context::RenderedVerifiedAnswers; -use codex_protocol::ThreadId; use codex_protocol::protocol::EventMsg; -use codex_protocol::protocol::SessionSource; -use codex_protocol::protocol::SubAgentSource; use codex_protocol::request_user_input::RequestUserInputAnswer; use codex_protocol::request_user_input::RequestUserInputResponse; use pretty_assertions::assert_eq; @@ -25,15 +22,9 @@ use test_case::test_case; use tokio::sync::Mutex; #[tokio::test] -async fn multi_agent_v2_request_user_input_rejects_subagent_threads() { +async fn request_user_input_rejects_restricted_threads() { let (session, mut turn) = make_session_and_context().await; - turn.session_source = SessionSource::SubAgent(SubAgentSource::ThreadSpawn { - parent_thread_id: ThreadId::new(), - depth: 1, - agent_path: None, - agent_nickname: None, - agent_role: None, - }); + turn.root_capabilities_restricted = true; let turn = Arc::new(turn); let result = RequestUserInputHandler { @@ -74,7 +65,7 @@ async fn multi_agent_v2_request_user_input_rejects_subagent_threads() { .await; let Err(err) = result else { - panic!("sub-agent request_user_input should fail"); + panic!("restricted request_user_input should fail"); }; assert_eq!( err, diff --git a/codex-rs/core/src/tools/spec_plan.rs b/codex-rs/core/src/tools/spec_plan.rs index 80ffe7b138..eb8d9a628a 100644 --- a/codex-rs/core/src/tools/spec_plan.rs +++ b/codex-rs/core/src/tools/spec_plan.rs @@ -1324,7 +1324,7 @@ fn add_core_utility_tools(context: &CoreToolPlanContext<'_>, registry: &mut Tool } if turn_context.config.experimental_request_user_input_enabled - && !turn_context.session_source.is_non_root_agent() + && !turn_context.root_capabilities_restricted && context .model_info .experimental_supported_tools @@ -1351,7 +1351,7 @@ fn add_core_utility_tools(context: &CoreToolPlanContext<'_>, registry: &mut Tool ); } - if !turn_context.session_source.is_non_root_agent() + if !turn_context.root_capabilities_restricted && (features.enabled(Feature::SendMessageToUserAsync) || context .model_info diff --git a/codex-rs/core/tests/suite/guardian_retained_context.rs b/codex-rs/core/tests/suite/guardian_retained_context.rs index 7fea5d2ee3..d6904fa094 100644 --- a/codex-rs/core/tests/suite/guardian_retained_context.rs +++ b/codex-rs/core/tests/suite/guardian_retained_context.rs @@ -9,6 +9,8 @@ use anyhow::Context; use anyhow::Result; use codex_core::CodexThread; use codex_core::ForkSnapshot; +use codex_core::RestrictRootCapabilities; +use codex_core::StartThreadOptions; use codex_core::TurnInputRequest; use codex_core::compact::content_items_to_text; use codex_core::config::Config; @@ -21,8 +23,10 @@ use codex_extension_api::ToolLifecycleContributor; use codex_extension_api::ToolLifecycleFuture; use codex_extension_api::ToolStartInput; use codex_features::Feature; +use codex_history::CodexHarnessMetadata; use codex_history::GuardianHistoryCheckpoint; use codex_history::InitialHistory; +use codex_history::ResponseItemEnvelope; use codex_history::ResumedHistory; use codex_history::RetainedContext; use codex_history::RetainedContextEntry; @@ -40,6 +44,7 @@ use codex_protocol::models::ContentItem; use codex_protocol::models::ResponseItem; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::Op; +use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::ThreadHistoryMode; use codex_protocol::protocol::ThreadSettingsOverrides; use codex_protocol::request_user_input::RequestUserInputAnswer; @@ -59,6 +64,7 @@ use core_test_support::responses::mount_sse_once_match; use core_test_support::responses::mount_sse_sequence; use core_test_support::responses::sse; use core_test_support::responses::start_mock_server; +use core_test_support::responses::user_message_item; use core_test_support::skip_if_no_network; use core_test_support::streaming_sse::StreamingSseChunk; use core_test_support::streaming_sse::start_streaming_sse_server; @@ -782,6 +788,67 @@ async fn legacy_checkpoint_recovers_root_excerpt_before_discarding_backup( Ok(()) } +#[test_case(false; "root")] +#[test_case(true; "host restricted")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn inherited_instruction_authorization_after_cold_resume(restricted: bool) -> Result<()> { + skip_if_no_network!(Ok(())); + let server = start_mock_server().await; + let test = test_codex().build_with_auto_env(&server).await?; + let options = |initial_history| { + let mut options = StartThreadOptions::new(test.config.clone()); + options.initial_history = initial_history; + options.session_source = Some(SessionSource::Custom("host".into())); + if restricted { + options + .thread_extension_init + .insert(RestrictRootCapabilities); + } + options + }; + let instruction = "Publish the report."; + let inherited = RolloutItem::ResponseItem(ResponseItemEnvelope { + item: user_message_item(instruction), + metadata: Some(CodexHarnessMetadata { + inherited_user_message: true, + ..Default::default() + }), + }); + let fork = test + .thread_manager + .start_thread(options(InitialHistory::Forked(vec![inherited]))) + .await?; + fork.thread.shutdown_and_wait().await?; + test.thread_manager.remove_thread(&fork.thread_id).await; + let saved = load_context(&test, &fork.thread).await?; + let resumed = test + .thread_manager + .start_thread(options(InitialHistory::Resumed(ResumedHistory { + conversation_id: fork.thread_id, + history: Arc::new(saved), + history_revision: None, + rollout_path: None, + }))) + .await?; + + let snapshot = resumed.thread.conversation_history_snapshot().await; + assert!(snapshot.items().any(|item| matches!( + item, ResponseItem::Message { role, content, .. } + if role == "user" && content_items_to_text(content).as_deref() == Some(instruction) + ))); + let retained = snapshot + .retained_context() + .context("resumed authorization")?; + assert_eq!( + retained.ordered_entries().any(|(_, entry)| matches!( + entry, RetainedContextEntry::UserMessage(user) if user.text == instruction + )), + !restricted + ); + resumed.thread.shutdown_and_wait().await?; + Ok(()) +} + // Cover live copied history and a truncated referenced checkpoint. Parent-answer // omission is already exercised by retained_answers_cross_real_session_boundaries. #[test_case(ThreadHistoryMode::Legacy; "copied worker history")] diff --git a/codex-rs/core/tests/suite/mcp_auth_elicitation.rs b/codex-rs/core/tests/suite/mcp_auth_elicitation.rs index 732162376b..5247ac1ab5 100644 --- a/codex-rs/core/tests/suite/mcp_auth_elicitation.rs +++ b/codex-rs/core/tests/suite/mcp_auth_elicitation.rs @@ -4,6 +4,7 @@ use anyhow::Result; use codex_analytics::AnalyticsEventsClient; use codex_analytics::AppServerRpcTransport; use codex_app_server_protocol as app; +use codex_core::RestrictRootCapabilities; use codex_core::StartThreadOptions; use codex_core::TurnInputRequest; use codex_core::TurnInputSubmission; @@ -15,6 +16,7 @@ use codex_extension_api::ToolLifecycleFuture; use codex_features::Feature; use codex_login::CodexAuth; use codex_mcp::CODEX_APPS_MCP_SERVER_NAME; +use codex_protocol::SessionId; use codex_protocol::approvals::ElicitationRequest; use codex_protocol::config_types::ApprovalsReviewer; use codex_protocol::items::McpToolCallStatus; @@ -522,15 +524,17 @@ enum SubagentRequestKind { ConnectorAuth(AuthFailureFormat), } -#[test_case(SubagentRequestKind::ToolApproval, ElicitationAction::Accept; "child_tool_approval_waits_for_user_before_execution")] -#[test_case(SubagentRequestKind::LegacyToolApproval, ElicitationAction::Accept; "child_legacy_tool_approval_is_rejected_before_execution")] -#[test_case(SubagentRequestKind::ConnectorAuth(AuthFailureFormat::Text), ElicitationAction::Accept; "child_connector_auth_accepts_text_failure")] -#[test_case(SubagentRequestKind::ConnectorAuth(AuthFailureFormat::Structured), ElicitationAction::Accept; "child_connector_auth_accepts_structured_failure")] -#[test_case(SubagentRequestKind::ConnectorAuth(AuthFailureFormat::Structured), ElicitationAction::Decline; "child_declined_connector_auth_preserves_original_failure")] +#[test_case(SubagentRequestKind::ToolApproval, ElicitationAction::Accept, false; "child_tool_approval_waits_for_user_before_execution")] +#[test_case(SubagentRequestKind::LegacyToolApproval, ElicitationAction::Accept, false; "child_legacy_tool_approval_is_rejected_before_execution")] +#[test_case(SubagentRequestKind::ConnectorAuth(AuthFailureFormat::Text), ElicitationAction::Accept, false; "child_connector_auth_accepts_text_failure")] +#[test_case(SubagentRequestKind::ConnectorAuth(AuthFailureFormat::Structured), ElicitationAction::Accept, false; "child_connector_auth_accepts_structured_failure")] +#[test_case(SubagentRequestKind::ConnectorAuth(AuthFailureFormat::Structured), ElicitationAction::Decline, false; "child_declined_connector_auth_preserves_original_failure")] +#[test_case(SubagentRequestKind::LegacyToolApproval, ElicitationAction::Accept, true; "restricted_standalone_legacy_tool_approval_is_rejected_before_execution")] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn core_generated_mcp_elicitations_support_subagents( +async fn core_generated_mcp_elicitations_support_restricted_threads( kind: SubagentRequestKind, decision: ElicitationAction, + restrict_root_capabilities: bool, ) -> Result<()> { skip_if_no_network!(Ok(())); let server = responses::start_mock_server().await; @@ -585,21 +589,38 @@ approvals_reviewer = "user" }) .build_with_auto_env(&server) .await?; - let child = test - .thread_manager - .start_thread(StartThreadOptions { - session_source: Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { - parent_thread_id: test.session_configured.thread_id, - depth: 1, - agent_path: None, - agent_nickname: None, - agent_role: None, - })), - environments: Some(vec![test.executor_environment().request()]), - ..StartThreadOptions::new(test.config.clone()) + let session_source = if restrict_root_capabilities { + SessionSource::Custom("host".to_string()) + } else { + SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id: test.session_configured.thread_id, + depth: 1, + agent_path: None, + agent_nickname: None, + agent_role: None, }) - .await? - .thread; + }; + let mut options = StartThreadOptions { + session_source: Some(session_source), + environments: Some(vec![test.executor_environment().request()]), + ..StartThreadOptions::new(test.config.clone()) + }; + if restrict_root_capabilities { + options + .thread_extension_init + .insert(RestrictRootCapabilities); + } + let child = test.thread_manager.start_thread(options).await?; + if restrict_root_capabilities { + assert_eq!( + ( + child.session_configured.parent_thread_id, + child.session_configured.session_id + ), + (None, SessionId::from(child.thread_id)), + ); + } + let child = child.thread; wait_for_mcp_server(&child, CODEX_APPS_MCP_SERVER_NAME).await?; responses::mount_sse_once( &server, diff --git a/codex-rs/core/tests/suite/mcp_subagent_elicitation.rs b/codex-rs/core/tests/suite/mcp_subagent_elicitation.rs index 51ac8aa8a0..06a0bb3310 100644 --- a/codex-rs/core/tests/suite/mcp_subagent_elicitation.rs +++ b/codex-rs/core/tests/suite/mcp_subagent_elicitation.rs @@ -1,10 +1,13 @@ //! MCP server elicitations preserve approval policy and signal abandonment before shutdown. use anyhow::Result; +use codex_core::RestrictRootCapabilities; use codex_core::StartThreadOptions; use codex_core::TurnInputRequest; use codex_protocol::approvals::ElicitationAction; use codex_protocol::config_types::ApprovalsReviewer; +use codex_protocol::mcp::ClientMcpExtensions; +use codex_protocol::mcp::OPENAI_STANDARD_FORM_INPUT_EXTENSION_ID; use codex_protocol::models::PermissionProfile; use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::EventMsg; @@ -12,6 +15,7 @@ use codex_protocol::protocol::Op; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::SubAgentSource; use codex_protocol::protocol::ThreadSettingsOverrides; +use codex_protocol::protocol::ThreadSource; use codex_protocol::user_input::UserInput; use core_test_support::responses; use core_test_support::responses::ResponsesRequest; @@ -69,6 +73,7 @@ pub(super) enum Caller { Root, Subagent, FullAccessSubagent, + FullAccessHost(Option), } #[derive(Clone, Copy)] @@ -84,6 +89,8 @@ pub(super) enum RequestKind { #[test_case(Caller::Root, RequestKind::BrowserAuth; "root_browser_auth_remains_interactive")] #[test_case(Caller::Root, RequestKind::ShutdownFormInput; "pending_root_form_is_abandoned_before_shutdown")] #[test_case(Caller::Subagent, RequestKind::FormInput; "subagent_receives_form_input")] +#[test_case(Caller::FullAccessHost(None), RequestKind::FormInput; "full_access_host_receives_form_input")] +#[test_case(Caller::FullAccessHost(Some(RestrictRootCapabilities)), RequestKind::FormInput; "restricted_full_access_host_declines_form_input")] #[test_case(Caller::Subagent, RequestKind::Permission; "subagent_permission_remains_interactive")] #[test_case(Caller::FullAccessSubagent, RequestKind::Permission; "full_access_subagent_automatically_approves_permission")] #[test_case(Caller::Subagent, RequestKind::StrictReview; "subagent_keeps_strict_automatic_safety_review")] @@ -143,28 +150,43 @@ pub(super) async fn mcp_server_elicitation_scenario( "default_tools_approval_mode": "approve" } }))?)?; - let full_access = matches!(caller, Caller::FullAccessSubagent); - let expects_prompt = !full_access && !matches!(request_kind, RequestKind::StrictReview); - let session_source = if matches!(caller, Caller::Root) { - SessionSource::Exec - } else { - SessionSource::SubAgent(SubAgentSource::ThreadSpawn { - parent_thread_id: test.session_configured.thread_id, - depth: 1, - agent_path: None, - agent_nickname: None, - agent_role: None, - }) + let full_access = matches!( + caller, + Caller::FullAccessSubagent | Caller::FullAccessHost(_) + ); + let expects_prompt = match caller { + Caller::FullAccessHost(restriction) => restriction.is_none(), + Caller::FullAccessSubagent => false, + Caller::Root | Caller::Subagent => !matches!(request_kind, RequestKind::StrictReview), }; - let thread = test - .thread_manager - .start_thread(StartThreadOptions { - session_source: Some(session_source), - environments: Some(vec![test.executor_environment().request()]), - ..StartThreadOptions::new(config) - }) - .await? - .thread; + let mut options = StartThreadOptions { + environments: Some(vec![test.executor_environment().request()]), + ..StartThreadOptions::new(config) + }; + options.session_source = Some(match caller { + Caller::Root => SessionSource::Exec, + Caller::Subagent | Caller::FullAccessSubagent => { + SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id: test.session_configured.thread_id, + depth: 1, + agent_path: None, + agent_nickname: None, + agent_role: None, + }) + } + Caller::FullAccessHost(restriction) => { + options.thread_source = Some(ThreadSource::User); + options.client_mcp_extensions = ClientMcpExtensions::new([( + OPENAI_STANDARD_FORM_INPUT_EXTENSION_ID.to_string(), + json!({}), + )]); + if let Some(restriction) = restriction { + options.thread_extension_init.insert(restriction); + } + SessionSource::Custom("host".to_string()) + } + }); + let thread = test.thread_manager.start_thread(options).await?.thread; wait_for_mcp_server(&thread, "elicitation").await?; let tool_call = responses::mount_sse_once( &server, @@ -314,6 +336,8 @@ pub(super) async fn mcp_server_elicitation_scenario( "action": "accept", "content": {}, "_meta": {"approvals_reviewer": "auto_review"} }) + } else if matches!(caller, Caller::FullAccessHost(Some(_))) { + json!({"action": "decline"}) } else { json!({"action": "accept", "content": response_content}) }; diff --git a/codex-rs/core/tests/suite/request_user_input_async.rs b/codex-rs/core/tests/suite/request_user_input_async.rs index 48fa1668a0..ab9596c2a0 100644 --- a/codex-rs/core/tests/suite/request_user_input_async.rs +++ b/codex-rs/core/tests/suite/request_user_input_async.rs @@ -1,4 +1,5 @@ use anyhow::Result; +use codex_core::RestrictRootCapabilities; use codex_core::StartThreadOptions; use codex_core::TurnInputRequest; use codex_features::Feature; @@ -33,15 +34,22 @@ use pretty_assertions::assert_eq; use serde_json::json; use test_case::test_case; -#[test_case(SessionSource::Exec, None, true, false; "root_without_capability")] -#[test_case(SessionSource::Exec, Some("send_user_message_async"), true, true; "legacy_capability_enabled")] -#[test_case(SessionSource::Exec, Some("send_user_message_async"), false, false; "legacy_capability_disabled")] -#[test_case(SessionSource::Exec, Some("request_user_input_async"), true, true; "current_capability_enabled")] -#[test_case(SessionSource::Exec, Some("request_user_input_async"), false, false; "current_capability_disabled")] -#[test_case(SessionSource::SubAgent(SubAgentSource::Other("test".to_string())), Some("request_user_input_async"), true, false; "subagent")] +enum ThreadKind { + Root, + NativeChild, + HostRestricted, +} + +#[test_case(ThreadKind::Root, None, true, false; "root_without_capability")] +#[test_case(ThreadKind::Root, Some("send_user_message_async"), true, true; "legacy_capability_enabled")] +#[test_case(ThreadKind::Root, Some("send_user_message_async"), false, false; "legacy_capability_disabled")] +#[test_case(ThreadKind::Root, Some("request_user_input_async"), true, true; "current_capability_enabled")] +#[test_case(ThreadKind::Root, Some("request_user_input_async"), false, false; "current_capability_disabled")] +#[test_case(ThreadKind::NativeChild, Some("request_user_input_async"), true, false; "subagent")] +#[test_case(ThreadKind::HostRestricted, Some("request_user_input_async"), true, false; "host_restricted")] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn async_question_registration_respects_user_input_setting( - session_source: SessionSource, + thread_kind: ThreadKind, catalog_tool: Option<&'static str>, enabled: bool, expect_question_tool: bool, @@ -68,15 +76,21 @@ async fn async_question_registration_respects_user_input_setting( }) .build_with_auto_env(&server) .await?; - let is_root = !session_source.is_non_root_agent(); - let thread = test - .thread_manager - .start_thread(StartThreadOptions { - session_source: Some(session_source), - ..StartThreadOptions::new(test.config.clone()) - }) - .await? - .thread; + let session_source = match thread_kind { + ThreadKind::Root => SessionSource::Exec, + ThreadKind::NativeChild => SessionSource::SubAgent(SubAgentSource::Other("test".into())), + ThreadKind::HostRestricted => SessionSource::Custom("host".into()), + }; + let mut options = StartThreadOptions { + session_source: Some(session_source), + ..StartThreadOptions::new(test.config.clone()) + }; + if matches!(thread_kind, ThreadKind::HostRestricted) { + options + .thread_extension_init + .insert(RestrictRootCapabilities); + } + let thread = test.thread_manager.start_thread(options).await?.thread; thread .start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text { text: "Report progress.".to_string(), @@ -99,7 +113,7 @@ async fn async_question_registration_respects_user_input_setting( .iter() .filter(|tool| tool["name"] == "send_message_to_user_async") .count(), - usize::from(is_root), + usize::from(matches!(thread_kind, ThreadKind::Root)), ); Ok(()) }