diff --git a/codex-rs/config/src/merge_tests.rs b/codex-rs/config/src/merge_tests.rs index f9da6e7ebc14..d9f856bce41f 100644 --- a/codex-rs/config/src/merge_tests.rs +++ b/codex-rs/config/src/merge_tests.rs @@ -2,6 +2,8 @@ use super::*; use crate::config_toml::ConfigToml; use crate::types::MemoriesToml; use pretty_assertions::assert_eq; +use std::collections::BTreeMap; +use std::collections::BTreeSet; fn parse_toml(value: &str) -> TomlValue { toml::from_str(value).expect("TOML should parse") @@ -124,3 +126,40 @@ fn merge_toml_values_normalizes_permission_network_domains_before_overlaying() { ); assert_eq!(base, expected); } + +#[test] +fn server_registered_mcp_tool_policy_replaces_lower_layer_policy() { + let mut base = parse_toml( + r#" +[features.server_registered_tools_only] +enabled = true +mcp_tools = [ + { server_name = "user_server", tool_name = "write" }, +] +"#, + ); + let overlay = parse_toml( + r#" +[features.server_registered_tools_only] +enabled = true +mcp_tools = [ + { server_name = "hoopa", tool_name = "lookup" }, + { server_name = "hoopa", tool_name = "search" }, +] +"#, + ); + + merge_toml_values(&mut base, &overlay); + + let config: ConfigToml = base.try_into().expect("merged config should deserialize"); + assert_eq!( + config + .features + .expect("features should be present") + .server_registered_mcp_tools(), + BTreeMap::from([( + "hoopa".to_string(), + BTreeSet::from(["lookup".to_string(), "search".to_string()]), + )]) + ); +} diff --git a/codex-rs/config/src/schema.rs b/codex-rs/config/src/schema.rs index 8547269ce2f4..6f168eafad69 100644 --- a/codex-rs/config/src/schema.rs +++ b/codex-rs/config/src/schema.rs @@ -44,6 +44,15 @@ pub fn features_schema(schema_gen: &mut SchemaGenerator) -> Schema { ); continue; } + if feature.id == codex_features::Feature::ServerRegisteredToolsOnly { + validation.properties.insert( + feature.key.to_string(), + schema_gen.subschema_for::>(), + ); + continue; + } if feature.id == codex_features::Feature::TokenBudget { validation.properties.insert( feature.key.to_string(), diff --git a/codex-rs/core/config.schema.json b/codex-rs/core/config.schema.json index fb4d8740ba77..4416a88953da 100644 --- a/codex-rs/core/config.schema.json +++ b/codex-rs/core/config.schema.json @@ -635,6 +635,9 @@ "secret_auth_storage": { "type": "boolean" }, + "server_registered_tools_only": { + "$ref": "#/definitions/FeatureToml_for_ServerRegisteredToolsOnlyConfigToml" + }, "shell_snapshot": { "type": "boolean" }, @@ -993,6 +996,16 @@ } ] }, + "FeatureToml_for_ServerRegisteredToolsOnlyConfigToml": { + "anyOf": [ + { + "type": "boolean" + }, + { + "$ref": "#/definitions/ServerRegisteredToolsOnlyConfigToml" + } + ] + }, "FeatureToml_for_TokenBudgetConfigToml": { "anyOf": [ { @@ -2783,6 +2796,38 @@ }, "type": "object" }, + "ServerRegisteredMcpToolToml": { + "additionalProperties": false, + "properties": { + "server_name": { + "type": "string" + }, + "tool_name": { + "type": "string" + } + }, + "required": [ + "server_name", + "tool_name" + ], + "type": "object" + }, + "ServerRegisteredToolsOnlyConfigToml": { + "additionalProperties": false, + "properties": { + "enabled": { + "type": "boolean" + }, + "mcp_tools": { + "description": "Raw MCP server and action identities authorized by the hosting server. An absent or empty list permits no MCP tools.", + "items": { + "$ref": "#/definitions/ServerRegisteredMcpToolToml" + }, + "type": "array" + } + }, + "type": "object" + }, "SessionPickerViewMode": { "description": "Preferred layout for the resume/fork session picker.", "enum": [ @@ -5007,6 +5052,9 @@ "secret_auth_storage": { "type": "boolean" }, + "server_registered_tools_only": { + "$ref": "#/definitions/FeatureToml_for_ServerRegisteredToolsOnlyConfigToml" + }, "shell_snapshot": { "type": "boolean" }, diff --git a/codex-rs/core/src/config/mod.rs b/codex-rs/core/src/config/mod.rs index 6348f14c0443..31ae82660cc1 100644 --- a/codex-rs/core/src/config/mod.rs +++ b/codex-rs/core/src/config/mod.rs @@ -120,6 +120,7 @@ use rmcp::model::UrlElicitationCapability; use serde::Deserialize; use serde::Serialize; use std::collections::BTreeMap; +use std::collections::BTreeSet; use std::collections::HashMap; use std::collections::HashSet; use std::io::ErrorKind; @@ -825,6 +826,10 @@ pub struct Config { /// Definition for MCP servers that Codex can reach out to for tool calls. pub mcp_servers: Constrained>, + /// Raw MCP identities authorized when only server-registered tools are enabled. + /// An absent server or action is denied at tool-registry construction time. + pub server_registered_mcp_tools: BTreeMap>, + /// Preferred store for MCP OAuth credentials. /// keyring: Use an OS-specific keyring service. /// Credentials stored in the keyring will only be readable by Codex unless the user explicitly grants access via OS-level keyring access. @@ -3068,6 +3073,11 @@ impl Config { }, feature_overrides, ); + let server_registered_mcp_tools = cfg + .features + .as_ref() + .map(FeaturesToml::server_registered_mcp_tools) + .unwrap_or_default(); let features = ManagedFeatures::from_configured_with_warnings( configured_features, feature_requirements, @@ -3832,6 +3842,7 @@ impl Config { env!("CARGO_PKG_VERSION"), ), mcp_servers, + server_registered_mcp_tools, // The config.toml omits "_mode" because it's a config file. However, "_mode" // is important in code to differentiate the mode from the store implementation. mcp_oauth_credentials_store_mode: resolve_mcp_oauth_credentials_store_mode( diff --git a/codex-rs/core/src/config/schema_tests.rs b/codex-rs/core/src/config/schema_tests.rs index dd67ead89867..be310be9c3df 100644 --- a/codex-rs/core/src/config/schema_tests.rs +++ b/codex-rs/core/src/config/schema_tests.rs @@ -73,3 +73,24 @@ fn config_schema_hides_unsupported_inline_mcp_bearer_token() { (false, true), ); } + +#[test] +fn config_schema_exposes_server_registered_tool_policy() { + let schema_json = config_schema_json().expect("serialize config schema"); + let schema_value: serde_json::Value = + serde_json::from_slice(&schema_json).expect("decode schema json"); + + assert_eq!( + schema_value.pointer("/properties/features/properties/server_registered_tools_only/$ref"), + Some(&serde_json::json!( + "#/definitions/FeatureToml_for_ServerRegisteredToolsOnlyConfigToml" + )) + ); + assert!( + schema_value + .pointer( + "/definitions/ServerRegisteredToolsOnlyConfigToml/properties/mcp_tools/items/$ref" + ) + .is_some() + ); +} diff --git a/codex-rs/core/src/tools/spec_plan.rs b/codex-rs/core/src/tools/spec_plan.rs index 57879aededfc..8523d5056754 100644 --- a/codex-rs/core/src/tools/spec_plan.rs +++ b/codex-rs/core/src/tools/spec_plan.rs @@ -89,6 +89,7 @@ use codex_tools::request_user_input_available_modes; use codex_tools::shell_command_backend_for_features; use codex_tools::shell_type_for_model_and_features; use std::collections::BTreeMap; +use std::collections::BTreeSet; use std::collections::HashSet; use std::sync::Arc; use tracing::instrument; @@ -194,13 +195,30 @@ fn build_tool_specs_and_registry( wait_agent_timeouts: wait_agent_timeout_options(turn_context), }; let mut planned_tools = PlannedTools::default(); - add_tool_sources(&context, &mut planned_tools); - apply_direct_model_only_namespace_overrides(turn_context, &mut planned_tools); - append_tool_search_executor(&context, &mut planned_tools); - prepend_code_mode_executors(&context, &mut planned_tools); + if server_registered_tools_only(turn_context) { + add_mcp_runtime_tools( + &context, + &mut planned_tools, + Some(&turn_context.config.server_registered_mcp_tools), + ); + add_dynamic_tools(&context, &mut planned_tools); + } else { + add_tool_sources(&context, &mut planned_tools); + apply_direct_model_only_namespace_overrides(turn_context, &mut planned_tools); + append_tool_search_executor(&context, &mut planned_tools); + prepend_code_mode_executors(&context, &mut planned_tools); + } build_model_visible_specs_and_registry(turn_context, planned_tools) } +fn server_registered_tools_only(turn_context: &TurnContext) -> bool { + turn_context + .config + .features + .get() + .enabled(Feature::ServerRegisteredToolsOnly) +} + fn apply_direct_model_only_namespace_overrides( turn_context: &TurnContext, planned_tools: &mut PlannedTools, @@ -239,6 +257,7 @@ fn build_model_visible_specs_and_registry( runtimes, hosted_specs, } = planned_tools; + let server_only = server_registered_tools_only(turn_context); let mut specs = Vec::new(); let mut seen_tool_names = HashSet::new(); for runtime in &runtimes { @@ -247,15 +266,15 @@ fn build_model_visible_specs_and_registry( continue; } let exposure = runtime.exposure(); - if exposure.is_direct() && !is_hidden_by_code_mode_only(turn_context, &tool_name, exposure) + if exposure.is_direct() + && (server_only || !is_hidden_by_code_mode_only(turn_context, &tool_name, exposure)) { let spec = runtime.spec(); - specs.push(spec_for_model_request( - turn_context, - exposure, - &tool_name, - spec, - )); + specs.push(if server_only { + spec + } else { + spec_for_model_request(turn_context, exposure, &tool_name, spec) + }); } } specs.extend(hosted_specs); @@ -322,12 +341,15 @@ fn hosted_model_tool_specs(context: &CoreToolPlanContext<'_>) -> Vec { } pub(crate) fn search_tool_enabled(turn_context: &TurnContext) -> bool { - turn_context.model_info.supports_search_tool && namespace_tools_enabled(turn_context) + !server_registered_tools_only(turn_context) + && turn_context.model_info.supports_search_tool + && namespace_tools_enabled(turn_context) } pub(crate) fn tool_suggest_enabled(turn_context: &TurnContext) -> bool { let features = turn_context.config.features.get(); - features.enabled(Feature::ToolSuggest) + !server_registered_tools_only(turn_context) + && features.enabled(Feature::ToolSuggest) && features.enabled(Feature::Apps) && features.enabled(Feature::Plugins) } @@ -610,7 +632,7 @@ fn add_tool_sources(context: &CoreToolPlanContext<'_>, planned_tools: &mut Plann add_mcp_resource_tools(context, planned_tools); add_core_utility_tools(context, planned_tools); add_collaboration_tools(context, planned_tools); - add_mcp_runtime_tools(context, planned_tools); + add_mcp_runtime_tools(context, planned_tools, /*allowed_raw_tools*/ None); add_extension_tools(context, planned_tools); add_dynamic_tools(context, planned_tools); for spec in hosted_model_tool_specs(context) { @@ -877,9 +899,16 @@ fn add_collaboration_tools(context: &CoreToolPlanContext<'_>, planned_tools: &mu deferred_mcp_tool_count = context.deferred_mcp_tools.map_or(0, <[ToolInfo]>::len) ) )] -fn add_mcp_runtime_tools(context: &CoreToolPlanContext<'_>, planned_tools: &mut PlannedTools) { +fn add_mcp_runtime_tools( + context: &CoreToolPlanContext<'_>, + planned_tools: &mut PlannedTools, + allowed_raw_tools: Option<&BTreeMap>>, +) { if let Some(mcp_tools) = context.mcp_tools { - for tool in mcp_tools { + for tool in mcp_tools + .iter() + .filter(|tool| raw_mcp_tool_is_allowed(tool, allowed_raw_tools)) + { match McpHandler::new(tool.clone()) { Ok(handler) => planned_tools.add(handler), Err(err) => warn!( @@ -891,7 +920,10 @@ fn add_mcp_runtime_tools(context: &CoreToolPlanContext<'_>, planned_tools: &mut } if let Some(deferred_mcp_tools) = context.deferred_mcp_tools { - for tool in deferred_mcp_tools { + for tool in deferred_mcp_tools + .iter() + .filter(|tool| raw_mcp_tool_is_allowed(tool, allowed_raw_tools)) + { match McpHandler::new(tool.clone()) { Ok(handler) => planned_tools.add_with_exposure(handler, ToolExposure::Deferred), Err(err) => warn!( @@ -903,6 +935,17 @@ fn add_mcp_runtime_tools(context: &CoreToolPlanContext<'_>, planned_tools: &mut } } +fn raw_mcp_tool_is_allowed( + tool: &ToolInfo, + allowed_raw_tools: Option<&BTreeMap>>, +) -> bool { + allowed_raw_tools.is_none_or(|allowed_raw_tools| { + allowed_raw_tools + .get(&tool.server_name) + .is_some_and(|allowed_tools| allowed_tools.contains(tool.tool.name.as_ref())) + }) +} + #[instrument( level = "trace", skip_all, diff --git a/codex-rs/core/src/tools/spec_plan_tests.rs b/codex-rs/core/src/tools/spec_plan_tests.rs index 7fd4c69cd5f5..14274d4e4afe 100644 --- a/codex-rs/core/src/tools/spec_plan_tests.rs +++ b/codex-rs/core/src/tools/spec_plan_tests.rs @@ -1,4 +1,5 @@ use std::collections::BTreeMap; +use std::collections::BTreeSet; use std::sync::Arc; use codex_features::Feature; @@ -899,6 +900,101 @@ async fn invalid_mcp_tools_are_not_registered() { plan.assert_registered_lacks(&[&ToolName::namespaced("mcp__invalid", "lookup").to_string()]); } +#[tokio::test] +async fn server_registered_tools_only_restricts_the_complete_tool_inventory() { + let plan = probe_with( + |turn| { + set_features( + turn, + &[ + Feature::Apps, + Feature::ServerRegisteredToolsOnly, + Feature::CodeMode, + Feature::CodeModeOnly, + Feature::ImageGeneration, + Feature::MultiAgentV2, + Feature::Plugins, + Feature::RequestPermissionsTool, + Feature::ShellTool, + Feature::ToolSuggest, + ], + ); + set_web_search_mode(turn, WebSearchMode::Live); + turn.model_info.supports_search_tool = true; + turn.model_info.web_search_tool_type = WebSearchToolType::TextAndImage; + turn.model_info.apply_patch_tool_type = Some(ApplyPatchToolType::Freeform); + turn.model_info.input_modalities = vec![InputModality::Image]; + update_config(turn, |config| { + config.server_registered_mcp_tools = BTreeMap::from([( + "hoopa".to_string(), + BTreeSet::from(["lookup".to_string(), "search".to_string()]), + )]); + }); + use_chatgpt_auth(turn); + assert!(!super::tool_suggest_enabled(turn)); + }, + ToolPlanInputs { + mcp_tools: Some(vec![ + mcp_tool("hoopa", "mcp__hoopa", "lookup"), + mcp_tool("hoopa", "mcp__hoopa", "search"), + mcp_tool("hoopa", "mcp__hoopa", "update"), + mcp_tool("user_server", "mcp__user_server", "lookup"), + ]), + deferred_mcp_tools: Some(vec![mcp_tool( + "user_deferred", + "mcp__user_deferred", + "lookup", + )]), + tool_suggest_candidates: Some(plugin_candidates(ToolSuggestPresentation::ListTool)), + extension_tool_executors: vec![ + Arc::new(DeferredExtensionTool), + Arc::new(TestNamespaceExtensionTool { + namespace: "image_gen", + tool_name: "imagegen", + }), + ], + dynamic_tools: vec![ + dynamic_tool( + /*namespace*/ None, + "server_callback", + /*defer_loading*/ false, + ), + dynamic_tool( + /*namespace*/ None, + "server_deferred", + /*defer_loading*/ true, + ), + ], + }, + ) + .await; + + assert_eq!(plan.visible_names, vec!["mcp__hoopa", "server_callback"]); + assert_eq!( + plan.namespace_function_names("mcp__hoopa"), + &["lookup".to_string(), "search".to_string()] + ); + assert_eq!( + plan.registered_names, + vec![ + ToolName::namespaced("mcp__hoopa", "lookup").to_string(), + ToolName::namespaced("mcp__hoopa", "search").to_string(), + "server_callback".to_string(), + "server_deferred".to_string(), + ] + ); + assert_eq!( + plan.exposure(&ToolName::namespaced("mcp__hoopa", "lookup").to_string()), + ToolExposure::Direct + ); + assert_eq!( + plan.exposure(&ToolName::namespaced("mcp__hoopa", "search").to_string()), + ToolExposure::Direct + ); + assert_eq!(plan.exposure("server_callback"), ToolExposure::Direct); + assert_eq!(plan.exposure("server_deferred"), ToolExposure::Deferred); +} + #[tokio::test] async fn request_plugin_install_requires_all_discovery_features() { for disabled_feature in [Feature::ToolSuggest, Feature::Apps, Feature::Plugins] { diff --git a/codex-rs/core/tests/suite/tools.rs b/codex-rs/core/tests/suite/tools.rs index 5e47eae52e2c..e390a3854b37 100644 --- a/codex-rs/core/tests/suite/tools.rs +++ b/codex-rs/core/tests/suite/tools.rs @@ -9,6 +9,10 @@ use anyhow::Context; use anyhow::Result; use codex_core::sandboxing::SandboxPermissions; use codex_features::Feature; +use codex_protocol::dynamic_tools::DynamicToolCallOutputContentItem; +use codex_protocol::dynamic_tools::DynamicToolFunctionSpec; +use codex_protocol::dynamic_tools::DynamicToolResponse; +use codex_protocol::dynamic_tools::DynamicToolSpec; use codex_protocol::models::PermissionProfile; use codex_protocol::permissions::FileSystemAccessMode; use codex_protocol::permissions::FileSystemPath; @@ -16,6 +20,9 @@ use codex_protocol::permissions::FileSystemSandboxEntry; use codex_protocol::permissions::FileSystemSandboxPolicy; use codex_protocol::permissions::NetworkSandboxPolicy; use codex_protocol::protocol::AskForApproval; +use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::Op; +use codex_protocol::user_input::UserInput; use core_test_support::assert_regex_match; use core_test_support::responses::ev_assistant_message; use core_test_support::responses::ev_completed; @@ -31,6 +38,8 @@ use core_test_support::skip_if_no_network; use core_test_support::skip_if_sandbox; use core_test_support::test_codex::local; use core_test_support::test_codex::test_codex; +use core_test_support::wait_for_event; +use pretty_assertions::assert_eq; use regex_lite::Regex; use serde_json::Value; use serde_json::json; @@ -93,6 +102,112 @@ async fn empty_turn_environments_omits_environment_backed_tools() -> Result<()> Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn server_registered_tools_only_is_enforced_in_model_requests() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + let call_id = "server-echo-call"; + let response_mock = mount_sse_sequence( + &server, + vec![ + sse(vec![ + ev_response_created("resp-1"), + ev_function_call(call_id, "server_echo", r#"{"message":"hello"}"#), + ev_completed("resp-1"), + ]), + sse(vec![ + ev_response_created("resp-2"), + ev_assistant_message("msg-1", "done"), + ev_completed("resp-2"), + ]), + ], + ) + .await; + + let mut builder = test_codex().with_config(|config| { + for feature in [ + Feature::ServerRegisteredToolsOnly, + Feature::CodeMode, + Feature::CodeModeOnly, + ] { + config + .features + .enable(feature) + .expect("test feature should enable"); + } + }); + let base_test = builder.build(&server).await?; + let new_thread = base_test + .thread_manager + .start_thread_with_tools( + base_test.config.clone(), + vec![DynamicToolSpec::Function(DynamicToolFunctionSpec { + name: "server_echo".to_string(), + description: "Echo a server-owned value.".to_string(), + input_schema: json!({ + "type": "object", + "properties": {"message": {"type": "string"}}, + "required": ["message"], + "additionalProperties": false, + }), + defer_loading: false, + })], + ) + .await?; + let mut test = base_test; + test.codex = new_thread.thread; + test.session_configured = new_thread.session_configured; + + test.codex + .submit(Op::UserInput { + items: vec![UserInput::Text { + text: "Use the server echo tool.".to_string(), + text_elements: Vec::new(), + }], + final_output_json_schema: None, + responsesapi_client_metadata: None, + additional_context: Default::default(), + thread_settings: Default::default(), + }) + .await?; + + let EventMsg::DynamicToolCallRequest(request) = wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::DynamicToolCallRequest(_)) + }) + .await + else { + unreachable!("event guard guarantees DynamicToolCallRequest"); + }; + assert_eq!(request.call_id, call_id); + assert_eq!(request.tool, "server_echo"); + assert_eq!(request.arguments, json!({"message": "hello"})); + + test.codex + .submit(Op::DynamicToolResponse { + id: request.call_id, + response: DynamicToolResponse { + content_items: vec![DynamicToolCallOutputContentItem::InputText { + text: "echo-ok".to_string(), + }], + success: true, + }, + }) + .await?; + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + + let requests = response_mock.requests(); + assert_eq!(requests.len(), 2); + for request in requests { + assert_eq!(tool_names(&request.body_json()), vec!["server_echo"]); + } + + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn turn_environment_selection_keeps_environment_backed_tools() -> Result<()> { skip_if_no_network!(Ok(())); diff --git a/codex-rs/features/src/feature_configs.rs b/codex-rs/features/src/feature_configs.rs index 0100dc1d8387..ad5ea0e9cc2b 100644 --- a/codex-rs/features/src/feature_configs.rs +++ b/codex-rs/features/src/feature_configs.rs @@ -29,6 +29,34 @@ impl FeatureConfig for CodeModeConfigToml { } } +#[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq, Eq, JsonSchema)] +#[serde(deny_unknown_fields)] +pub struct ServerRegisteredToolsOnlyConfigToml { + #[serde(skip_serializing_if = "Option::is_none")] + pub enabled: Option, + /// Raw MCP server and action identities authorized by the hosting server. + /// An absent or empty list permits no MCP tools. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub mcp_tools: Vec, +} + +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq, JsonSchema)] +#[serde(deny_unknown_fields)] +pub struct ServerRegisteredMcpToolToml { + pub server_name: String, + pub tool_name: String, +} + +impl FeatureConfig for ServerRegisteredToolsOnlyConfigToml { + fn enabled(&self) -> Option { + self.enabled + } + + fn set_enabled(&mut self, enabled: bool) { + self.enabled = Some(enabled); + } +} + #[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq, Eq, JsonSchema)] #[serde(deny_unknown_fields)] pub struct MultiAgentV2ConfigToml { diff --git a/codex-rs/features/src/lib.rs b/codex-rs/features/src/lib.rs index 716bb2e44132..befb2b43cc58 100644 --- a/codex-rs/features/src/lib.rs +++ b/codex-rs/features/src/lib.rs @@ -27,6 +27,8 @@ pub use feature_configs::NetworkProxyModeToml; pub use feature_configs::NetworkProxyUnixSocketPermissionToml; use feature_configs::RemovedAppsMcpPathOverrideConfigToml; pub use feature_configs::RolloutBudgetConfigToml; +pub use feature_configs::ServerRegisteredMcpToolToml; +pub use feature_configs::ServerRegisteredToolsOnlyConfigToml; pub use feature_configs::TokenBudgetConfigToml; use legacy::LegacyFeatureToggles; pub use legacy::legacy_feature_keys; @@ -164,6 +166,11 @@ pub enum Feature { ToolSearchAlwaysDeferMcpTools, /// Expose MCP model-visible namespaces without the legacy `mcp__` prefix. NonPrefixedMcpToolNames, + /// Restrict the model-visible tool surface to tools registered by the hosting server. + /// + /// This retains host-allowlisted MCP tools and dynamic tools while removing + /// Codex-owned core, extension, hosted, collaboration, and tool-discovery tools. + ServerRegisteredToolsOnly, /// Enable discoverable tool suggestions for apps. ToolSuggest, /// Enable plugins. @@ -647,6 +654,8 @@ pub struct FeaturesToml { pub rollout_budget: Option>, #[serde(default, skip_serializing_if = "Option::is_none")] pub current_time_reminder: Option>, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub server_registered_tools_only: Option>, #[serde(default, rename = "apps_mcp_path_override", skip_serializing)] #[schemars(skip)] removed_apps_mcp_path_override: Option>, @@ -692,6 +701,16 @@ impl FeaturesToml { { entries.insert(Feature::CurrentTimeReminder.key().to_string(), enabled); } + if let Some(enabled) = self + .server_registered_tools_only + .as_ref() + .and_then(FeatureToml::enabled) + { + entries.insert( + Feature::ServerRegisteredToolsOnly.key().to_string(), + enabled, + ); + } if let Some(enabled) = self.network_proxy.as_ref().and_then(FeatureToml::enabled) { entries.insert(Feature::NetworkProxy.key().to_string(), enabled); } @@ -706,6 +725,7 @@ impl FeaturesToml { token_budget, rollout_budget, current_time_reminder, + server_registered_tools_only, removed_apps_mcp_path_override: _, network_proxy, entries, @@ -725,6 +745,8 @@ impl FeaturesToml { materialize_resolved_feature_enabled(rollout_budget, enabled); } else if spec.id == Feature::CurrentTimeReminder { materialize_resolved_feature_enabled(current_time_reminder, enabled); + } else if spec.id == Feature::ServerRegisteredToolsOnly { + materialize_resolved_feature_enabled(server_registered_tools_only, enabled); } else if spec.id == Feature::NetworkProxy { materialize_resolved_feature_enabled(network_proxy, enabled); } else { @@ -732,6 +754,22 @@ impl FeaturesToml { } } } + + pub fn server_registered_mcp_tools(&self) -> BTreeMap> { + match self.server_registered_tools_only.as_ref() { + Some(FeatureToml::Config(config)) => { + let mut tools = BTreeMap::>::new(); + for tool in &config.mcp_tools { + tools + .entry(tool.server_name.clone()) + .or_default() + .insert(tool.tool_name.clone()); + } + tools + } + Some(FeatureToml::Enabled(_)) | None => BTreeMap::new(), + } + } } fn materialize_resolved_feature_enabled( @@ -1088,6 +1126,12 @@ pub const FEATURES: &[FeatureSpec] = &[ stage: Stage::UnderDevelopment, default_enabled: false, }, + FeatureSpec { + id: Feature::ServerRegisteredToolsOnly, + key: "server_registered_tools_only", + stage: Stage::UnderDevelopment, + default_enabled: false, + }, FeatureSpec { id: Feature::UnavailableDummyTools, key: "unavailable_dummy_tools", diff --git a/codex-rs/features/src/tests.rs b/codex-rs/features/src/tests.rs index 50da143fbb15..e3a1e923cedd 100644 --- a/codex-rs/features/src/tests.rs +++ b/codex-rs/features/src/tests.rs @@ -11,6 +11,7 @@ use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::WarningEvent; use pretty_assertions::assert_eq; use std::collections::BTreeMap; +use std::collections::BTreeSet; use toml::Table; use toml::Value as TomlValue; @@ -680,6 +681,37 @@ non_code_mode_only = true ); } +#[test] +fn server_registered_tools_only_deserializes_exact_raw_mcp_policy() { + let features: FeaturesToml = toml::from_str( + r#" +[server_registered_tools_only] +enabled = true +mcp_tools = [ + { server_name = "hoopa", tool_name = "lookup" }, + { server_name = "hoopa", tool_name = "search" }, +] +"#, + ) + .expect("features table should deserialize"); + + assert_eq!( + features.entries(), + BTreeMap::from([("server_registered_tools_only".to_string(), true)]) + ); + assert_eq!( + features.server_registered_mcp_tools(), + BTreeMap::from([( + "hoopa".to_string(), + BTreeSet::from(["lookup".to_string(), "search".to_string()]), + )]) + ); + + let boolean_only: FeaturesToml = toml::from_str("server_registered_tools_only = true") + .expect("boolean feature toggle should deserialize"); + assert_eq!(boolean_only.server_registered_mcp_tools(), BTreeMap::new()); +} + #[test] fn materialize_resolved_enabled_writes_all_features_and_preserves_custom_config() { let mut features = Features::with_defaults(); @@ -687,6 +719,7 @@ fn materialize_resolved_enabled_writes_all_features_and_preserves_custom_config( features.enable(Feature::MultiAgentV2); features.enable(Feature::NetworkProxy); features.enable(Feature::RespectSystemProxy); + features.enable(Feature::ServerRegisteredToolsOnly); let mut features_toml = FeaturesToml { multi_agent_v2: Some(FeatureToml::Config(crate::MultiAgentV2ConfigToml { @@ -699,6 +732,15 @@ fn materialize_resolved_enabled_writes_all_features_and_preserves_custom_config( proxy_url: Some("http://127.0.0.1:43128".to_string()), ..Default::default() })), + server_registered_tools_only: Some(FeatureToml::Config( + crate::ServerRegisteredToolsOnlyConfigToml { + enabled: Some(false), + mcp_tools: vec![crate::ServerRegisteredMcpToolToml { + server_name: "hoopa".to_string(), + tool_name: "search".to_string(), + }], + }, + )), entries: BTreeMap::new(), ..Default::default() }; @@ -730,6 +772,10 @@ fn materialize_resolved_enabled_writes_all_features_and_preserves_custom_config( ..Default::default() })) ); + assert_eq!( + features_toml.server_registered_mcp_tools(), + BTreeMap::from([("hoopa".to_string(), BTreeSet::from(["search".to_string()]),)]) + ); let replayed = Features::from_sources( FeatureConfigSource { features: Some(&features_toml), diff --git a/codex-rs/thread-manager-sample/src/main.rs b/codex-rs/thread-manager-sample/src/main.rs index c35daaa497fa..c684ff403219 100644 --- a/codex-rs/thread-manager-sample/src/main.rs +++ b/codex-rs/thread-manager-sample/src/main.rs @@ -229,6 +229,7 @@ fn new_config(model: Option, arg0_paths: Arg0DispatchPaths) -> anyhow::R workspace_roots_explicit: false, cli_auth_credentials_store_mode: AuthCredentialsStoreMode::File, mcp_servers: Constrained::allow_any(HashMap::new()), + server_registered_mcp_tools: BTreeMap::new(), mcp_oauth_credentials_store_mode: OAuthCredentialsStoreMode::File, mcp_oauth_callback_port: None, mcp_oauth_callback_url: None,