From b46a3801e41414cc5a73040d2ecb282dfbc527ba Mon Sep 17 00:00:00 2001 From: Keyan Zhang <2268452+keyz@users.noreply.github.com> Date: Mon, 26 Oct 2026 21:11:10 +0000 Subject: [PATCH] Preserve parent tool declarations in prediction forks (#57287) ## What changed Keep inherited incremental tool declarations in MCP-disabled prediction forks and `thread/sample` requests, and append a developer notice that MCP tools are unavailable. Skip tool declaration updates within the inherited context window; new windows still declare the current tools. Add the default-enabled `prediction_preserve_parent_tools` feature flag and its configuration schema entries. ## Testing Update core snapshots and app-server tests to cover inherited declarations, the MCP-unavailable notice, preservation across MCP refresh, and continued parent tool availability. GitOrigin-RevId: 3adf96c7d0e0093986dbdbd540bd47674a2c105c --- .../app-server/tests/suite/v2/thread_fork.rs | 103 ++++++++++-------- .../tests/suite/v2/thread_sample.rs | 11 +- codex-rs/core/config.schema.json | 6 + codex-rs/core/src/session/mod.rs | 22 ++++ codex-rs/core/src/session/world_state.rs | 4 +- .../suite/scenarios_incremental_tools.rs | 23 +++- ...ools__prediction_fork_inherited_tools.snap | 57 ++++++++++ ...al_tools__prediction_fork_mcp_removal.snap | 90 --------------- codex-rs/features/src/lib.rs | 8 ++ 9 files changed, 174 insertions(+), 150 deletions(-) create mode 100644 codex-rs/core/tests/suite/snapshots/all__suite__scenarios__incremental_tools__prediction_fork_inherited_tools.snap delete mode 100644 codex-rs/core/tests/suite/snapshots/all__suite__scenarios__incremental_tools__prediction_fork_mcp_removal.snap diff --git a/codex-rs/app-server/tests/suite/v2/thread_fork.rs b/codex-rs/app-server/tests/suite/v2/thread_fork.rs index 355f25ba84..6c902ce83b 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_fork.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_fork.rs @@ -2972,7 +2972,7 @@ async fn fork_option_disables_mcp_servers( let server = responses::start_mock_server().await; let mock = responses::mount_sse_sequence( &server, - ["parent", "child", "refreshed-child", "continued-parent"] + ["parent", "prediction", "continued-parent"] .into_iter() .map(|id| responses::sse(vec![responses::ev_completed(id)])) .collect(), @@ -3014,7 +3014,7 @@ async fn fork_option_disables_mcp_servers( let mut model = construct_model_info_offline("gpt-5.4", &config); model.use_responses_lite = responses_lite; model.supports_search_tool = false; - model.tool_mode = Some(codex_protocol::openai_models::ToolMode::Direct); + model.tool_mode = Some(codex_protocol::openai_models::ToolMode::CodeMode); let catalog = codex_home.path().join("catalog.json"); std::fs::write( &catalog, @@ -3063,38 +3063,36 @@ async fn fork_option_disables_mcp_servers( }, }) .await?; - let removes_mcp = disable_mcp && incremental_tools && responses_lite; - for (index, thread_id) in [&child.thread.id, &child.thread.id, &parent.thread.id] - .into_iter() - .enumerate() - { - if index == 1 { - let initializations = mcp_server - .received_requests() - .await - .expect("MCP request recording must be enabled") - .iter() - .filter(|request| request.method.as_str() == "POST") - .filter(|request| { - serde_json::from_slice::(&request.body) - .expect("MCP request must be valid JSON")["method"] - == "initialize" - }) - .count(); - if removes_mcp { - assert_eq!(initializations, 1, "disabled fork must not connect MCPs"); - } - let _: McpServerRefreshResponse = app - .request(|request_id| ClientRequest::McpServerRefresh { - request_id, - params: None, - }) - .await?; - } + let disables_mcp = disable_mcp && incremental_tools && responses_lite; + if disables_mcp { + let initializations = mcp_server + .received_requests() + .await + .expect("MCP request recording must be enabled") + .iter() + .filter(|request| request.method.as_str() == "POST") + .filter(|request| { + serde_json::from_slice::(&request.body) + .expect("MCP request must be valid JSON")["method"] + == "initialize" + }) + .count(); + assert_eq!(initializations, 1, "disabled fork must not connect MCPs"); + } + let _: McpServerRefreshResponse = app + .request(|request_id| ClientRequest::McpServerRefresh { + request_id, + params: None, + }) + .await?; + for (thread_id, text) in [ + (&child.thread.id, "Predict the next message."), + (&parent.thread.id, "Continue"), + ] { app.start_turn_and_wait_for_completion(TurnStartParams { thread_id: thread_id.clone(), input: vec![UserInput::Text { - text: "Continue".into(), + text: text.into(), text_elements: vec![], }], ..Default::default() @@ -3108,21 +3106,36 @@ async fn fork_option_disables_mcp_servers( .to_string() .contains("prediction_mcp_lookup") ); - assert_eq!( - requests - .iter() - .map(|request| request.body_json()["input"] - .as_array() - .expect("input") + let declarations = requests[0].inputs_of_type("additional_tools"); + if disables_mcp { + assert!( + declarations .iter() - .any(|item| { - let text = item.to_string(); - text.contains("namespaces are no longer available") - && text.contains("mcp__fixture") - })) - .collect::>(), - vec![false, removes_mcp, removes_mcp, false], - ); + .any(|item| responses::namespace_child_tool(item, "functions", "exec").is_some()) + ); + } + let notice_counts: Vec<_> = requests + .iter() + .map(|request| { + request + .message_input_texts("developer") + .iter() + .filter(|text| text.contains("MCP tools are unavailable for this request.")) + .count() + }) + .collect(); + assert_eq!(notice_counts, [0, usize::from(disables_mcp), 0]); + for request in &requests[1..] { + if disables_mcp { + assert_eq!(request.inputs_of_type("additional_tools"), declarations); + } + assert!( + !request + .input() + .iter() + .any(|item| item.to_string().contains("no longer available")) + ); + } Ok(()) } diff --git a/codex-rs/app-server/tests/suite/v2/thread_sample.rs b/codex-rs/app-server/tests/suite/v2/thread_sample.rs index efe380839a..b01dd5d5f0 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_sample.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_sample.rs @@ -778,14 +778,9 @@ async fn sample_blocks_mcp_execution_without_changing_parent_tools( let prefix = parent["input"].as_array().expect("parent input"); let sample_input = sample["input"].as_array().expect("sample input"); assert_eq!(&sample_input[..prefix.len()], prefix); - let removes_mcp = disable_mcp && incremental_tools; - assert_eq!( - sample_input[prefix.len()..].iter().any(|item| { - let text = item.to_string(); - text.contains("namespaces are no longer available") && text.contains("mcp__sample_test") - }), - removes_mcp - ); + assert!(!sample_input[prefix.len()..].iter().any(|item| { + item["type"] == "additional_tools" || item.to_string().contains("no longer available") + })); assert_eq!(sample, comparable_request(requests[3].body_json())); assert_eq!( requests[2].body_json()["tools"], diff --git a/codex-rs/core/config.schema.json b/codex-rs/core/config.schema.json index 795b6590bd..a1b54e6bcc 100644 --- a/codex-rs/core/config.schema.json +++ b/codex-rs/core/config.schema.json @@ -1212,6 +1212,9 @@ "powershell_shell_version": { "type": "boolean" }, + "prediction_preserve_parent_tools": { + "type": "boolean" + }, "prefer_mxc": { "type": "boolean" }, @@ -7839,6 +7842,9 @@ "powershell_shell_version": { "type": "boolean" }, + "prediction_preserve_parent_tools": { + "type": "boolean" + }, "prefer_mxc": { "type": "boolean" }, diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 9c5f50d1ac..d156184403 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -1802,6 +1802,18 @@ impl Session { } } self.persist_rollout_items(&rollout_items).await; + if self.preserve_parent_tool_declarations() { + let notice = DeveloperInstructions::new( + "\nMCP tools are unavailable for this request.\n", + ) + .with_metadata(ContentItemMetadata::tool(/*namespace*/ None)); + self.record_conversation_items( + &turn_context, + turn_context.model_info(), + &[ContextualUserFragment::into(notice)], + ) + .await; + } // Agent spawn owns its final durability barrier so it can overlap the edge write. let persist_context = match self.fork_persistence { @@ -4572,6 +4584,16 @@ impl Session { true } + fn preserve_parent_tool_declarations(&self) -> bool { + // Only prediction forks and thread/sample set this marker when they disable MCP. + self.enabled(Feature::PredictionPreserveParentTools) + && self + .services + .thread_extension_data + .get::() + .is_some() + } + pub fn enabled(&self, feature: Feature) -> bool { self.features.enabled(feature) } diff --git a/codex-rs/core/src/session/world_state.rs b/codex-rs/core/src/session/world_state.rs index 41f9faf607..63a05f3fef 100644 --- a/codex-rs/core/src/session/world_state.rs +++ b/codex-rs/core/src/session/world_state.rs @@ -165,7 +165,9 @@ impl Session { let extension_metrics = super::extension_metrics::from_session_telemetry( step_context.session_telemetry.clone(), ); - if incremental_tools { + let preserve_parent_tool_declarations = + !new_window && self.preserve_parent_tool_declarations(); + if incremental_tools && !preserve_parent_tool_declarations { let specs = step_context.tool_router.model_visible_specs(); let definitions = codex_tools::create_tools_json_for_responses_lite(&specs)?; world_state.add_section(TopLevelToolsState::new( diff --git a/codex-rs/core/tests/suite/scenarios_incremental_tools.rs b/codex-rs/core/tests/suite/scenarios_incremental_tools.rs index 27173e088b..3164cf537a 100644 --- a/codex-rs/core/tests/suite/scenarios_incremental_tools.rs +++ b/codex-rs/core/tests/suite/scenarios_incremental_tools.rs @@ -38,7 +38,7 @@ use super::super::rmcp_client::remote_aware_environment_id; use super::super::rmcp_client::remote_aware_stdio_server_bin; #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn prediction_fork_appends_mcp_removal_without_rewriting_history() -> Result<()> { +async fn prediction_fork_keeps_inherited_tools() -> Result<()> { skip_if_wine_exec!( Ok(()), "requires a Windows test_stdio_server in the Wine-exec environment" @@ -55,7 +55,7 @@ async fn prediction_fork_appends_mcp_removal_without_rewriting_history() -> Resu .with_model_info_override("gpt-5.4", |model| { model.use_responses_lite = true; model.supports_search_tool = false; - model.tool_mode = Some(ToolMode::Direct); + model.tool_mode = Some(ToolMode::CodeModeOnly); }) .with_config(move |config| { config.workspace_roots = vec![config.cwd.clone()]; @@ -77,7 +77,7 @@ async fn prediction_fork_appends_mcp_removal_without_rewriting_history() -> Resu wait_for_mcp_server(&test.codex, "echo_service").await?; let mock = responses::mount_sse_sequence( &server, - ["parent", "child", "child-again", "parent-again"] + ["parent", "prediction", "continued-parent"] .into_iter() .map(|id| responses::sse(vec![responses::ev_completed(id)])) .collect(), @@ -111,7 +111,6 @@ async fn prediction_fork_appends_mcp_removal_without_rewriting_history() -> Resu child.thread.restore_thread_settings(settings).await?; test.codex = child.thread; test.submit_text_turn("Predict the next message.").await?; - test.submit_text_turn("Predict again.").await?; test.codex = parent; test.submit_text_turn("Continue using the echo service.") .await?; @@ -120,14 +119,26 @@ async fn prediction_fork_appends_mcp_removal_without_rewriting_history() -> Resu let parent = requests[0].body_json(); let child = requests[1].body_json(); let prefix = parent["input"].as_array().expect("parent input"); + let exec = prefix + .iter() + .filter(|item| item["type"] == "additional_tools") + .find_map(|item| responses::namespace_child_tool(item, "functions", "exec")) + .expect("parent declares functions.exec"); + assert!( + exec["description"] + .as_str() + .expect("exec description") + .contains("mcp__echo_service__echo"), + "the inherited functions.exec must describe the MCP helper", + ); assert_eq!( &child["input"].as_array().expect("child input")[..prefix.len()], prefix ); insta::assert_snapshot!( - "prediction_fork_mcp_removal", + "prediction_fork_inherited_tools", context_snapshot::format_request_history_snapshot( - "The prediction fork preserves the incremental parent's complete prefix, then appends MCP removal. Its second turn retains that update; continuing the parent keeps its original tools.", + "Request 1 starts the parent thread. Request 2 is the prediction child's only turn. Request 3 continues the original parent and verifies that it still has its MCP tools.", &requests, &ContextSnapshotOptions::default() .rewrite_known_segments() diff --git a/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__incremental_tools__prediction_fork_inherited_tools.snap b/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__incremental_tools__prediction_fork_inherited_tools.snap new file mode 100644 index 0000000000..c2e8f9c579 --- /dev/null +++ b/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__incremental_tools__prediction_fork_inherited_tools.snap @@ -0,0 +1,57 @@ +--- +source: core/tests/suite/scenarios_incremental_tools.rs +expression: "context_snapshot::format_request_history_snapshot(\"Request 1 starts the parent thread. Request 2 is the prediction child's only turn. Request 3 continues the original parent and verifies that it still has its MCP tools.\",\n&requests,\n&ContextSnapshotOptions::default().rewrite_known_segments().include_request_settings(),)" +--- +Scenario: Request 1 starts the parent thread. Request 2 is the prediction child's only turn. Request 3 continues the original parent and verifies that it still has its MCP tools. + +## Window 1 +Settings: + include: ["reasoning.encrypted_content"] + model: "gpt-5.4" + parallel_tool_calls: false + prompt_cache_key: "" + reasoning: {"context":"all_turns","effort":"medium"} + store: false + stream: true + text: {"verbosity":"low"} + tool_choice: "auto" +-- request 1 (turn) -- +00:additional_tools/developer (1; hash=a523a0cf594329a1): + - namespace/functions; hash=d3e1dde6718dc7a3 + - custom/exec + - function/wait + - function/request_user_input +01:message/developer: + +02:message/developer: + +03:message/user: + +04:message/user: + Start with the echo service. +-- request 2 (turn) -- +05:message/developer: + + MCP tools are unavailable for this request. + +06:message/user: + Predict the next message. + +## Window 2 (after request 2: input diverged at item 05) +Settings: same as window 1 +-- request 3 (turn) -- +00:additional_tools/developer (1; hash=a523a0cf594329a1): + - namespace/functions; hash=d3e1dde6718dc7a3 + - custom/exec + - function/wait + - function/request_user_input +01:message/developer: + +02:message/developer: + +03:message/user: + +04:message/user: + Start with the echo service. +05:message/user: + Continue using the echo service. diff --git a/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__incremental_tools__prediction_fork_mcp_removal.snap b/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__incremental_tools__prediction_fork_mcp_removal.snap deleted file mode 100644 index 655d2bcbdf..0000000000 --- a/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__incremental_tools__prediction_fork_mcp_removal.snap +++ /dev/null @@ -1,90 +0,0 @@ ---- -source: core/tests/suite/scenarios_incremental_tools.rs -expression: "context_snapshot::format_request_history_snapshot(\"The prediction fork preserves the incremental parent's complete prefix, then appends MCP removal. Its second turn retains that update; continuing the parent keeps its original tools.\",\n&requests,\n&ContextSnapshotOptions::default().rewrite_known_segments().include_request_settings(),)" ---- -Scenario: The prediction fork preserves the incremental parent's complete prefix, then appends MCP removal. Its second turn retains that update; continuing the parent keeps its original tools. - -## Window 1 -Settings: - include: ["reasoning.encrypted_content"] - model: "gpt-5.4" - parallel_tool_calls: false - prompt_cache_key: "" - reasoning: {"context":"all_turns","effort":"medium"} - store: false - stream: true - text: {"verbosity":"low"} - tool_choice: "auto" --- request 1 (turn) -- -00:additional_tools/developer (3; hash=94348ab8527021ff): - - namespace/functions; hash=124fee69a69ed6b2 - - function/exec_command - - function/write_stdin - - function/list_mcp_resources - - function/list_mcp_resource_templates - - function/read_mcp_resource - - function/request_user_input - - custom/apply_patch - - function/view_image - - namespace/multi_agent_v1: Tools for spawning and managing sub-agents.; hash=2327142bda177c50 - - function/close_agent - - function/resume_agent - - function/send_input - - function/spawn_agent - - function/wait_agent - - namespace/mcp__echo_service: Use these tools to exercise the rmcp test server.; hash=6650953285437e87 - - function/echo -01:message/developer: - -02:message/developer: - -03:message/user: - -04:message/user: - Start with the echo service. --- request 2 (turn) -- -05:message/developer: - The following namespaces are no longer available. Do not call tools in them unless those tools are declared in a later update: - - mcp__echo_service - - The following tools are no longer available. Do not call them: - - functions.list_mcp_resource_templates - - functions.list_mcp_resources - - functions.read_mcp_resource -06:message/user: - Predict the next message. --- request 3 (turn) -- -07:message/user: - Predict again. - -## Window 2 (after request 3: input diverged at item 05) -Settings: same as window 1 --- request 4 (turn) -- -00:additional_tools/developer (3; hash=94348ab8527021ff): - - namespace/functions; hash=124fee69a69ed6b2 - - function/exec_command - - function/write_stdin - - function/list_mcp_resources - - function/list_mcp_resource_templates - - function/read_mcp_resource - - function/request_user_input - - custom/apply_patch - - function/view_image - - namespace/multi_agent_v1: Tools for spawning and managing sub-agents.; hash=2327142bda177c50 - - function/close_agent - - function/resume_agent - - function/send_input - - function/spawn_agent - - function/wait_agent - - namespace/mcp__echo_service: Use these tools to exercise the rmcp test server.; hash=6650953285437e87 - - function/echo -01:message/developer: - -02:message/developer: - -03:message/user: - -04:message/user: - Start with the echo service. -05:message/user: - Continue using the echo service. diff --git a/codex-rs/features/src/lib.rs b/codex-rs/features/src/lib.rs index c6bfe00c75..29234609c8 100644 --- a/codex-rs/features/src/lib.rs +++ b/codex-rs/features/src/lib.rs @@ -287,6 +287,8 @@ pub enum Feature { DeferredToolWorldState, /// Track top-level tool definitions in world state and emit incremental context updates. IncrementalTools, + /// Preserve inherited tool declarations in MCP-disabled prediction forks. + PredictionPreserveParentTools, /// Expose MCP model-visible namespaces without the legacy `mcp__` prefix. NonPrefixedMcpToolNames, /// Enable discoverable tool suggestions for apps. @@ -1695,6 +1697,12 @@ pub const FEATURES: &[FeatureSpec] = &[ stage: Stage::Stable, default_enabled: true, }, + FeatureSpec { + id: Feature::PredictionPreserveParentTools, + key: "prediction_preserve_parent_tools", + stage: Stage::Stable, + default_enabled: true, + }, FeatureSpec { id: Feature::NonPrefixedMcpToolNames, key: "non_prefixed_mcp_tool_names",