Enforce strict third-party tool deferral in Code Mode Only (#59071)

## What changed

Restore `code_mode_only_strict_3p_tools` as an opt-in feature under development. When the effective tool mode is Code Mode Only, keep eligible MCP/app and client-supplied dynamic tools deferred and callable through `exec`, regardless of `deferLoading`, `omit_tools_from`, or namespace exclusion and direct-only settings.

Preserve hidden tools, tool policy restrictions, and built-in tool classification. Log conflicting exposure settings without adding warnings to the model prompt. Leave other effective tool modes unaffected.

## Testing

Add coverage for exposure overrides, model mode overrides, built-in name collisions, and hidden or denied tools. Update integration coverage for app access restrictions and MCP catalog refreshes that preserve the cached prompt prefix while `exec` uses the current tools.

GitOrigin-RevId: 0cb085ba9d23bebc48dbc3fa298b982de0a66c47
This commit is contained in:
rishab
2026-11-03 06:05:55 +00:00
committed by copyberry
parent 9e23a4708a
commit 29eed28082
13 changed files with 356 additions and 102 deletions
@@ -177,6 +177,10 @@ impl CoreToolRuntime for DynamicToolHandler {
fn finishes_on_cancellation(&self) -> bool {
true
}
fn is_third_party_tool(&self) -> bool {
true
}
}
#[expect(
+4
View File
@@ -352,6 +352,10 @@ impl CoreToolRuntime for McpHandler {
}))
}
fn is_third_party_tool(&self) -> bool {
true
}
fn mcp_server_name(&self) -> Option<&str> {
Some(&self.definition.tool_info.server_name)
}
+6
View File
@@ -92,6 +92,12 @@ pub(crate) trait CoreToolRuntime: ToolExecutor<ToolInvocation> {
None
}
/// True for MCP/app and client-supplied dynamic tools.
/// A built-in stays false even if a client's tool has the same name.
fn is_third_party_tool(&self) -> bool {
false
}
/// Identify the environments selected by the original input.
fn environments_for_input(
&self,
+51 -1
View File
@@ -217,6 +217,15 @@ pub(crate) fn build_tool_router(
)
}
/// Use the effective mode because the model can override the thread's configured mode.
fn code_mode_only_strict_3p_tools(turn_context: &TurnContext, model_info: &ModelInfo) -> bool {
effective_tool_mode(turn_context, model_info) == ToolMode::CodeModeOnly
&& turn_context
.config
.features
.enabled(Feature::CodeModeOnlyStrictThirdPartyTools)
}
fn apply_mcp_tool_exposure_policy(
turn_context: &TurnContext,
model_info: &ModelInfo,
@@ -261,6 +270,17 @@ fn apply_mcp_tool_exposure_policy(
let Some(omitted_exposures) = omitted_exposures_by_tool.get(&tool_name) else {
continue;
};
if code_mode_only_strict_3p_tools(turn_context, model_info) {
// Ignore MCP/app omit_tools_from; eligible tools must stay deferred in exec.
// Log conflicts without adding a warning to the model's prompt.
if omitted_exposures.intersects(ToolExposures::DEFERRED | ToolExposures::CODE_MODE) {
tracing::warn!(
tool_name = %tool_name,
"Ignoring MCP/app omit_tools_from because code_mode_only_strict_3p_tools is enabled"
);
}
continue;
}
let tool_name = tool_name.with_default_namespace();
let mut exposures = ToolExposures::ALL.difference(*omitted_exposures);
@@ -388,6 +408,7 @@ pub(crate) fn finalize_tool_router(
let model_messages = ResolvedModelMessages::from_model(model_info);
hosted_specs.retain(|spec| registry.tool_policy.allows(&ToolName::plain(spec.name())));
apply_direct_model_only_namespace_overrides(turn_context, &mut registry);
enforce_strict_3p_tools(turn_context, model_info, &mut registry);
let tool_mode = effective_tool_mode(turn_context, model_info);
let code_mode_enabled = matches!(tool_mode, ToolMode::CodeMode | ToolMode::CodeModeOnly);
if code_mode_enabled {
@@ -610,6 +631,31 @@ fn tool_exposure_with_namespace_override(
}
}
/// Run after namespace overrides so eligible third-party tools always remain deferred.
fn enforce_strict_3p_tools(
turn_context: &TurnContext,
model_info: &ModelInfo,
registry: &mut ToolRegistry,
) {
if !code_mode_only_strict_3p_tools(turn_context, model_info) {
return;
}
for tool in registry.entries_mut() {
// Remaining Hidden entries failed eligibility or budget checks; keep them hidden.
if !tool.runtime.is_third_party_tool() || tool.exposure == ToolExposure::Hidden {
continue;
}
if tool.exposure != ToolExposure::Deferred {
tracing::warn!(
tool_name = %tool.runtime.tool_name().with_default_namespace(),
"Deferring third-party tool because code_mode_only_strict_3p_tools is enabled"
);
tool.exposure = ToolExposure::Deferred;
}
}
}
#[instrument(level = "trace", skip_all)]
fn build_model_visible_specs(
turn_context: &TurnContext,
@@ -956,7 +1002,11 @@ fn code_mode_tools<'a>(
}
let tool_name = tool.runtime.tool_name();
if is_excluded_from_code_mode(turn_context, &tool_name) {
// Deferred third-party tools still need an exec route to be discoverable.
if is_excluded_from_code_mode(turn_context, &tool_name)
&& !(code_mode_only_strict_3p_tools(turn_context, model_info)
&& tool.runtime.is_third_party_tool())
{
return None;
}
@@ -1,80 +0,0 @@
//! Code Mode finalization preserves tool policy and hidden-tool eligibility.
use super::*;
use crate::tools::registry::ToolRegistry;
use codex_extension_api::ToolPolicy;
use pretty_assertions::assert_eq;
#[tokio::test]
async fn code_mode_finalization_preserves_hidden_and_policy_denied_tools() {
let (_, mut turn) = make_session_and_context().await;
set_features(&mut turn, &[Feature::CodeMode, Feature::CodeModeOnly]);
update_config(&mut turn, |config| {
config.code_mode.direct_only_tool_namespaces = vec!["mcp__allowed".to_string()];
});
let allowed_mcp = ToolName::namespaced("mcp__allowed", "lookup");
let hidden_mcp = ToolName::namespaced("mcp__allowed", "over_budget");
let mut registry = ToolRegistry::with_tool_policy(Arc::new(ToolPolicy {
allowed_tools: Some(vec![
ToolName::plain("exec"),
ToolName::plain("wait"),
allowed_mcp.clone(),
hidden_mcp.clone(),
ToolName::plain("allowed_dynamic"),
]),
..Default::default()
}));
let hosted = append_source_tools(
&turn,
turn.model_info(),
&mut registry,
vec![
mcp_runtime("allowed", "mcp__allowed", "lookup", ToolExposure::Direct),
// MCP registration keeps budget-overflow tools in the registry as Hidden.
// This is different from a caller setting omit_tools_from.
mcp_runtime(
"allowed",
"mcp__allowed",
"over_budget",
ToolExposure::Hidden,
),
mcp_runtime("denied", "mcp__denied", "lookup", ToolExposure::Direct),
],
std::iter::empty::<Arc<dyn for<'call> ToolExecutor<ExtensionToolCall<'call>>>>(),
&[
dynamic_tool(
/*namespace*/ None,
"allowed_dynamic",
/*defer_loading*/ false,
),
dynamic_tool(
/*namespace*/ None,
"denied_dynamic",
/*defer_loading*/ false,
),
],
);
let plan = ToolPlanProbe::from_router(ToolRouter::from_registry(
&turn,
turn.model_info(),
registry,
hosted,
&Default::default(),
));
plan.assert_registered_lacks(&[
&ToolName::namespaced("mcp__denied", "lookup").to_string(),
"denied_dynamic",
]);
assert_eq!(
plan.exposure(&allowed_mcp.to_string()),
ToolExposure::DirectModelOnly
);
assert_eq!(plan.exposure(&hidden_mcp.to_string()), ToolExposure::Hidden);
assert_eq!(plan.exposure("allowed_dynamic"), ToolExposure::Direct);
assert_eq!(
plan.code_mode_tool_names
.values()
.cloned()
.collect::<std::collections::HashSet<_>>(),
[ToolName::plain("allowed_dynamic")].into_iter().collect(),
);
}
@@ -0,0 +1,248 @@
//! Strict Code Mode Only: deferral, mode overrides, and hidden/denied tools.
use super::*;
use crate::tools::registry::ToolRegistry;
use codex_extension_api::ToolPolicy;
use pretty_assertions::assert_eq;
use tracing_test::internal::MockWriter;
#[tokio::test(flavor = "current_thread")]
async fn strict_defers_third_party_tools_despite_caller_settings_and_logs_overrides() {
let buffer: &'static std::sync::Mutex<Vec<u8>> =
Box::leak(Box::new(std::sync::Mutex::new(Vec::new())));
let subscriber = tracing_subscriber::fmt()
.with_ansi(false)
.without_time()
.with_max_level(tracing::Level::WARN)
.with_writer(MockWriter::new(buffer))
.finish();
// Keep callsite interest independent of untraced parallel test threads.
let _untraced = tracing::Dispatch::new(tracing::subscriber::NoSubscriber::default());
// The test runner may already own the global subscriber.
let _subscriber = tracing::subscriber::set_default(subscriber);
for supports_search_tool in [false, true] {
let plan = probe_with(
|turn| {
set_features(
turn,
&[
Feature::CodeMode,
Feature::CodeModeOnly,
Feature::CodeModeOnlyStrictThirdPartyTools,
],
);
update_turn_settings_for_test(turn, |settings| {
Arc::make_mut(&mut settings.model_info).supports_search_tool =
supports_search_tool;
});
update_config(turn, |config| {
config.update_plan_enabled = true;
config.code_mode.direct_only_tool_namespaces =
vec!["direct_client".to_string()];
config.code_mode.excluded_tool_namespaces = vec!["excluded_client".to_string()];
});
},
ToolPlanInputs {
dynamic_tools: vec![
dynamic_tool(Some("client"), "eager", /*defer_loading*/ false),
dynamic_tool(Some("client"), "deferred", /*defer_loading*/ true),
dynamic_tool(Some("direct_client"), "eager", /*defer_loading*/ false),
dynamic_tool(
Some("excluded_client"),
"eager",
/*defer_loading*/ false,
),
// Legal legacy input: app-server forbids deferLoading=true on a plain tool.
dynamic_tool(
/*namespace*/ None,
"client_echo",
/*defer_loading*/ false,
),
// A collision must not make the built-in tool third-party.
dynamic_tool(
/*namespace*/ None,
"update_plan",
/*defer_loading*/ false,
),
],
..ToolPlanInputs::default()
},
)
.await;
plan.assert_visible_contains(&["exec", "wait"]);
plan.assert_visible_lacks(&[
"tool_search",
"client",
"direct_client",
"excluded_client",
"client_echo",
]);
assert_eq!(plan.exposure("update_plan"), ToolExposure::Direct);
let ToolSpec::Freeform(exec) = plan.visible_spec("exec") else {
panic!("expected exec");
};
for name in [
ToolName::namespaced("client", "eager"),
ToolName::namespaced("client", "deferred"),
ToolName::namespaced("direct_client", "eager"),
ToolName::namespaced("excluded_client", "eager"),
ToolName::plain("client_echo"),
] {
assert_eq!(plan.exposure(&name.to_string()), ToolExposure::Deferred);
assert!(
plan.code_mode_tool_names
.values()
.any(|nested| nested == &name),
"{name} must be callable via exec"
);
assert!(
!exec.description.contains(&format!("{}(args:", name.name)),
"{name} schema must not be eager"
);
}
assert!(
!exec.description.contains("Deferring third-party tool"),
"warnings must not enter the model prompt"
);
}
let logs = String::from_utf8(buffer.lock().unwrap().clone()).unwrap();
assert!(
logs.contains(
"Deferring third-party tool because code_mode_only_strict_3p_tools is enabled"
)
);
assert!(logs.contains("client_echo"));
}
#[tokio::test]
async fn strict_flag_does_nothing_outside_effective_code_mode_only() {
let (_, mut turn) = make_session_and_context().await;
set_features(&mut turn, &[Feature::CodeMode, Feature::CodeModeOnly]);
turn.code_mode_available = true;
update_config(&mut turn, |config| {
config.code_mode.direct_only_tool_namespaces = vec!["client".to_string()];
config.code_mode.excluded_tool_namespaces = vec!["client".to_string()];
});
for mode in [ToolMode::Direct, ToolMode::CodeMode] {
update_turn_settings_for_test(&mut turn, |settings| {
let model = Arc::make_mut(&mut settings.model_info);
model.tool_mode = Some(mode);
model.supports_search_tool = true;
});
let inputs = || ToolPlanInputs {
dynamic_tools: vec![
dynamic_tool(Some("client"), "eager", /*defer_loading*/ false),
dynamic_tool(Some("client"), "deferred", /*defer_loading*/ true),
dynamic_tool(
/*namespace*/ None,
"client_echo",
/*defer_loading*/ false,
),
],
tool_runtimes: vec![mcp_runtime(
"sample",
"mcp__sample",
"lookup",
ToolExposure::Direct,
)],
..ToolPlanInputs::default()
};
set_feature(
&mut turn,
Feature::CodeModeOnlyStrictThirdPartyTools,
/*enabled*/ false,
);
let off = ToolPlanProbe::from_router(plan_with_model(&turn, turn.model_info(), inputs()));
assert_eq!(off.tool_mode, mode);
set_feature(
&mut turn,
Feature::CodeModeOnlyStrictThirdPartyTools,
/*enabled*/ true,
);
let on = ToolPlanProbe::from_router(plan_with_model(&turn, turn.model_info(), inputs()));
assert_eq!(on, off, "strict flag changed effective {mode:?}");
}
}
#[tokio::test]
async fn strict_finalization_cannot_resurrect_hidden_or_policy_denied_tools() {
let (_, mut turn) = make_session_and_context().await;
set_features(
&mut turn,
&[
Feature::CodeMode,
Feature::CodeModeOnly,
Feature::CodeModeOnlyStrictThirdPartyTools,
],
);
update_config(&mut turn, |config| {
config.code_mode.direct_only_tool_namespaces = vec!["mcp__allowed".to_string()];
});
let allowed_mcp = ToolName::namespaced("mcp__allowed", "lookup");
let hidden_mcp = ToolName::namespaced("mcp__allowed", "over_budget");
let mut registry = ToolRegistry::with_tool_policy(Arc::new(ToolPolicy {
allowed_tools: Some(vec![
ToolName::plain("exec"),
ToolName::plain("wait"),
allowed_mcp.clone(),
hidden_mcp.clone(),
ToolName::plain("allowed_dynamic"),
]),
..Default::default()
}));
let hosted = append_source_tools(
&turn,
turn.model_info(),
&mut registry,
vec![
mcp_runtime("allowed", "mcp__allowed", "lookup", ToolExposure::Direct),
// MCP registration keeps budget-overflow tools in the registry as Hidden.
// This is different from a caller setting omit_tools_from.
mcp_runtime(
"allowed",
"mcp__allowed",
"over_budget",
ToolExposure::Hidden,
),
mcp_runtime("denied", "mcp__denied", "lookup", ToolExposure::Direct),
],
std::iter::empty::<Arc<dyn for<'call> ToolExecutor<ExtensionToolCall<'call>>>>(),
&[
dynamic_tool(
/*namespace*/ None,
"allowed_dynamic",
/*defer_loading*/ false,
),
dynamic_tool(
/*namespace*/ None,
"denied_dynamic",
/*defer_loading*/ false,
),
],
);
let plan = ToolPlanProbe::from_router(ToolRouter::from_registry(
&turn,
turn.model_info(),
registry,
hosted,
&Default::default(),
));
plan.assert_registered_lacks(&[
&ToolName::namespaced("mcp__denied", "lookup").to_string(),
"denied_dynamic",
]);
assert_eq!(
plan.exposure(&allowed_mcp.to_string()),
ToolExposure::Deferred
);
assert_eq!(plan.exposure(&hidden_mcp.to_string()), ToolExposure::Hidden);
assert_eq!(plan.exposure("allowed_dynamic"), ToolExposure::Deferred);
assert_eq!(
plan.code_mode_tool_names
.values()
.cloned()
.collect::<std::collections::HashSet<_>>(),
[allowed_mcp, ToolName::plain("allowed_dynamic")]
.into_iter()
.collect(),
);
}
+2 -2
View File
@@ -3428,5 +3428,5 @@ async fn hosted_web_search_and_standalone_image_generation_follow_runtime_gates(
bedrock_with_standalone_web_search.assert_visible_lacks(&["web"]);
}
#[path = "spec_plan_code_mode_policy_tests.rs"]
mod code_mode_policy;
#[path = "spec_plan_strict_third_party_tests.rs"]
mod strict_third_party;
@@ -231,5 +231,5 @@ pub(super) async fn connector_exposure_requests(
Ok(requests)
}
#[path = "app_tool_exposure_code_mode_tests.rs"]
mod code_mode;
#[path = "app_tool_exposure_strict_tests.rs"]
mod strict_third_party;
@@ -1,10 +1,11 @@
//! Code Mode honors app exposure settings while preserving app access restrictions.
//! App omissions cannot bypass strict deferral; app access restrictions still apply.
use super::*;
use pretty_assertions::assert_eq;
use std::sync::Mutex;
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn code_mode_honors_app_omissions_and_preserves_app_access() -> anyhow::Result<()> {
#[tracing_test::traced_test]
async fn strict_ignores_app_omissions_but_preserves_app_access() -> anyhow::Result<()> {
let server = responses::start_mock_server().await;
let mut tools = vec![];
for name in ["lookup", "disabled", "app_only"] {
@@ -62,6 +63,10 @@ enabled = false
.features
.enable(Feature::CodeModeHost)
.expect("code mode host");
config
.features
.enable(Feature::CodeModeOnlyStrictThirdPartyTools)
.expect("strict third-party tools");
config.analytics_enabled = Some(false);
});
let test = builder.build_with_auto_env(&server).await?;
@@ -73,7 +78,6 @@ enabled = false
.tool_by_name("mcp__codex_apps__calendar", "lookup")
.is_none()
);
assert!(requests[0].body_contains_text("mcp__codex_apps__calendar__lookup("));
let output =
super::super::code_mode::custom_tool_output_last_non_empty_text(&requests[1], "lookup")
.expect("exec reports actual discovery and dispatch");
@@ -83,8 +87,19 @@ enabled = false
"names": ["mcp__codex_apps__calendar__lookup"],
"result": "called lookup for at with "
}),
"disabled and app-only tools never enter ALL_TOOLS"
"app omissions are ignored; disabled and app-only tools never enter ALL_TOOLS"
);
// The spawned Core session does not inherit the test span.
let logs = String::from_utf8(
tracing_test::internal::global_buf()
.lock()
.expect("captured logs")
.clone(),
)?;
let thread_id = test.session_configured.thread_id.to_string();
assert!(logs.lines().any(|line| {
line.contains(&thread_id) && line.contains("Ignoring MCP/app omit_tools_from")
}));
test.codex.shutdown_and_wait().await?;
Ok(())
}
+2 -2
View File
@@ -143,8 +143,8 @@ mod context_management_tools;
#[path = "scenarios_code_mode_settled_helpers_tests.rs"]
mod code_mode_settled_helpers;
#[path = "scenarios_deferred_mcp_cache.rs"]
mod deferred_mcp_cache;
#[path = "scenarios_strict_3p_cache.rs"]
mod strict_3p_cache;
#[path = "scenarios_agent_message_board.rs"]
mod agent_message_board;
@@ -20,7 +20,7 @@ use std::sync::Arc;
use std::sync::Mutex;
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn deferred_mcp_refresh_preserves_cache_while_exec_tracks_available_tools() -> Result<()> {
async fn strict_3p_mcp_refresh_preserves_cache_while_exec_tracks_available_tools() -> Result<()> {
let server = responses::start_mock_server().await;
let first_mcp = responses::start_mock_server().await;
let second_mcp = responses::start_mock_server().await;
@@ -54,11 +54,17 @@ async fn deferred_mcp_refresh_preserves_cache_while_exec_tracks_available_tools(
super::configure_scenario_catalog(config);
config.code_mode.disable_in_process_fallback = true;
config.features.enable(Feature::CodeModeHost).unwrap();
config
.features
.enable(Feature::CodeModeOnlyStrictThirdPartyTools)
.unwrap();
config.code_mode.direct_only_tool_namespaces = vec!["mcp__environment".to_string()];
config.code_mode.excluded_tool_namespaces = vec!["mcp__environment".to_string()];
})
.with_model_info_override("gpt-6-astra", |model| {
model.tool_mode = Some(ToolMode::CodeModeOnly);
model.use_responses_lite = true;
model.supports_search_tool = true;
model.supports_search_tool = false;
})
.build_with_auto_env(&server)
.await?;
@@ -95,7 +101,7 @@ text(JSON.stringify(results));"#;
"environment".to_string(),
serde_json::from_value(json!({
"url": format!("{url}/api/codex/ps/mcp"),
"omit_tools_from": ["direct"],
"omit_tools_from": ["direct", "deferred", "code_mode"],
}))?,
);
}
@@ -164,7 +170,7 @@ text(JSON.stringify(results));"#;
let labels = [
"turn 1: no MCP server registered",
"turn 1: exec discovery and call results",
"turn 2: register deferred MCP with lookup",
"turn 2: register MCP with lookup; omit-all, direct-only and excluded",
"turn 2: exec discovery and call results",
"turn 3: same MCP name, new endpoint; replace lookup with resolve and search",
"turn 3: exec discovery and call results",
@@ -173,7 +179,7 @@ text(JSON.stringify(results));"#;
];
let sections = labels.into_iter().zip(&requests).collect::<Vec<_>>();
let snapshot = context_snapshot::format_labeled_requests_snapshot(
"Code Mode Only: deferred MCP config changes between turns. Exec calls the live catalog. One window means the captured prefix stayed stable. The execution environment remains running.",
"Strict Code Mode Only: MCP config changes between turns, with omit-all and namespace skips. Exec calls the live catalog. One window means the captured prefix stayed stable. The execution environment remains running.",
&sections,
&ContextSnapshotOptions::default()
.rewrite_known_segments()
@@ -185,7 +191,7 @@ text(JSON.stringify(results));"#;
r"(?m)(^00:additional_tools/developer \(\d+; hash=|^ +- namespace/functions; hash=)[0-9a-f]{16}",
)?
.replace_all(&snapshot, "${1}<INITIAL_BUILD_HASH>");
insta::assert_snapshot!("deferred_mcp_catalog_cache", snapshot);
insta::assert_snapshot!("strict_3p_mcp_catalog_cache", snapshot);
test.codex.shutdown_and_wait().await?;
Ok(())
}
@@ -1,8 +1,8 @@
---
source: core/tests/suite/scenarios_deferred_mcp_cache.rs
source: core/tests/suite/scenarios_strict_3p_cache.rs
expression: snapshot
---
Scenario: Code Mode Only: deferred MCP config changes between turns. Exec calls the live catalog. One window means the captured prefix stayed stable. The execution environment remains running.
Scenario: Strict Code Mode Only: MCP config changes between turns, with omit-all and namespace skips. Exec calls the live catalog. One window means the captured prefix stayed stable. The execution environment remains running.
## Window 1
Settings:
@@ -57,7 +57,7 @@ Settings:
Wall time <DURATION> seconds
Output:
| []
-- request 3 (turn; turn 2: register deferred MCP with lookup) --
-- request 3 (turn; turn 2: register MCP with lookup; omit-all, direct-only and excluded) --
10:message/developer:
<tools>
Deferred tool namespaces:
+3 -2
View File
@@ -164,7 +164,8 @@ pub enum Feature {
CodeModeToolDescriptionFirst,
/// Restrict model-visible tools to code mode entrypoints (`exec`, `wait`).
CodeModeOnly,
/// Removed compatibility flag for strict third-party tool deferral in Code Mode Only.
/// Keep eligible MCP/app and dynamic tools deferred in exec, only in Code Mode Only.
/// Ignores deferLoading, omit_tools_from, and namespace skip settings.
CodeModeOnlyStrictThirdPartyTools,
/// Use the single unified PTY-backed exec tool.
UnifiedExec,
@@ -1362,7 +1363,7 @@ pub const FEATURES: &[FeatureSpec] = &[
FeatureSpec {
id: Feature::CodeModeOnlyStrictThirdPartyTools,
key: "code_mode_only_strict_3p_tools",
stage: Stage::Removed,
stage: Stage::UnderDevelopment,
default_enabled: false,
},
FeatureSpec {