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
This commit is contained in:
rhan-oai
2026-11-02 21:09:37 +00:00
committed by copyberry
parent 6c984676b0
commit 40cf168567
24 changed files with 266 additions and 123 deletions
@@ -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))
+14
View File
@@ -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::<RestrictRootCapabilities>()
.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))
}
}
@@ -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(),
@@ -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();
+3 -1
View File
@@ -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,
),
);
+15 -4
View File
@@ -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;
@@ -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<ThreadId>,
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(),
+7 -8
View File
@@ -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(),