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
This commit is contained in:
Keyan Zhang
2026-10-26 22:05:23 +00:00
committed by copyberry
parent 136d891a94
commit b46a3801e4
9 changed files with 174 additions and 150 deletions
@@ -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::<Value>(&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::<Value>(&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<_>>(),
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(())
}
@@ -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"],
+6
View File
@@ -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"
},
+22
View File
@@ -1802,6 +1802,18 @@ impl Session {
}
}
self.persist_rollout_items(&rollout_items).await;
if self.preserve_parent_tool_declarations() {
let notice = DeveloperInstructions::new(
"<mcp_tools>\nMCP tools are unavailable for this request.\n</mcp_tools>",
)
.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::<crate::McpServersDisabled>()
.is_some()
}
pub fn enabled(&self, feature: Feature) -> bool {
self.features.enabled(feature)
}
+3 -1
View File
@@ -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(
@@ -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()
@@ -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: "<PROMPT_CACHE_KEY 1>"
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:
<MODEL_INSTRUCTIONS>
02:message/developer:
<PERMISSIONS_INSTRUCTIONS>
03:message/user:
<ENVIRONMENT_CONTEXT>
04:message/user:
Start with the echo service.
-- request 2 (turn) --
05:message/developer:
<mcp_tools>
MCP tools are unavailable for this request.
</mcp_tools>
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:
<MODEL_INSTRUCTIONS>
02:message/developer:
<PERMISSIONS_INSTRUCTIONS>
03:message/user:
<ENVIRONMENT_CONTEXT>
04:message/user:
Start with the echo service.
05:message/user:
Continue using the echo service.
@@ -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: "<PROMPT_CACHE_KEY 1>"
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:
<MODEL_INSTRUCTIONS>
02:message/developer:
<PERMISSIONS_INSTRUCTIONS>
03:message/user:
<ENVIRONMENT_CONTEXT>
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:
<MODEL_INSTRUCTIONS>
02:message/developer:
<PERMISSIONS_INSTRUCTIONS>
03:message/user:
<ENVIRONMENT_CONTEXT>
04:message/user:
Start with the echo service.
05:message/user:
Continue using the echo service.
+8
View File
@@ -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",