agent/codex
2
0
Fork 0
mirror of https://github.com/pchuan98/codex.git synced 2026-07-01 00:31:56 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: persist and restore codex app's tools after search (#11780)

Browse Source 
### What changed
1. Removed per-turn MCP selection reset in `core/src/tasks/mod.rs`.
2. Added `SessionState::set_mcp_tool_selection(Vec<String>)` in
`core/src/state/session.rs` for authoritative restore behavior (deduped,
order-preserving, empty clears).
3. Added rollout parsing in `core/src/codex.rs` to recover
`active_selected_tools` from prior `search_tool_bm25` outputs:
   - tracks matching `call_id`s
   - parses function output text JSON
   - extracts `active_selected_tools`
   - latest valid payload wins
   - malformed/non-matching payloads are ignored
4. Applied restore logic to resumed and forked startup paths in
`core/src/codex.rs`.
5. Updated instruction text to session/thread scope in
`core/templates/search_tool/tool_description.md`.
6. Expanded tests in `core/tests/suite/search_tool.rs`, plus unit
coverage in:
   - `core/src/codex.rs`
   - `core/src/state/session.rs`

### Behavior after change
1. Search activates matched tools.
2. Additional searches union into active selection.
3. Selection survives new turns in the same thread.
4. Resume/fork restores selection from rollout history.
5. Separate threads do not inherit selection unless forked.
This commit is contained in:
Anton Panasenko
2026-02-15 19:18:41 -08:00
committed by GitHub
Unverified
parent 060a320e7d
commit 02abd9a8ea
12 changed files with 1201 additions and 201 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/codex.rs
+235 -13
View File
@@ -150,6 +150,7 @@ use crate::mcp::with_codex_apps_mcp;
use crate::mcp_connection_manager::McpConnectionManager;
use crate::mcp_connection_manager::filter_codex_apps_mcp_tools_only;
use crate::mcp_connection_manager::filter_mcp_tools_by_name;
use crate::mcp_connection_manager::filter_non_codex_apps_mcp_tools_only;
use crate::memories;
use crate::mentions::build_connector_slug_counts;
use crate::mentions::build_skill_name_counts;
@@ -223,6 +224,7 @@ use crate::tasks::SessionTask;
use crate::tasks::SessionTaskContext;
use crate::tools::ToolRouter;
use crate::tools::context::SharedTurnDiffTracker;
use crate::tools::handlers::SEARCH_TOOL_BM25_TOOL_NAME;
use crate::tools::js_repl::JsReplHandle;
use crate::tools::network_approval::NetworkApprovalService;
use crate::tools::network_approval::build_blocked_request_observer;
@@ -1489,6 +1491,11 @@ impl Session {
state.merge_mcp_tool_selection(tool_names)
}
pub(crate) async fn set_mcp_tool_selection(&self, tool_names: Vec<String>) {
let mut state = self.state.lock().await;
state.set_mcp_tool_selection(tool_names);
}
pub(crate) async fn get_mcp_tool_selection(&self) -> Option<Vec<String>> {
let state = self.state.lock().await;
state.get_mcp_tool_selection()
@@ -1522,6 +1529,7 @@ impl Session {
async fn record_initial_history(&self, conversation_history: InitialHistory) {
let turn_context = self.new_default_turn().await;
self.clear_mcp_tool_selection().await;
match conversation_history {
InitialHistory::New => {
// Build and record initial items (user instructions + environment context)
@@ -1537,6 +1545,8 @@ impl Session {
}
InitialHistory::Resumed(resumed_history) => {
let rollout_items = resumed_history.history;
let restored_tool_selection =
Self::extract_mcp_tool_selection_from_rollout(&rollout_items);
let previous_model = Self::last_rollout_model_name(&rollout_items)
.map(std::string::ToString::to_string);
{
@@ -1578,12 +1588,17 @@ impl Session {
let mut state = self.state.lock().await;
state.set_token_info(Some(info));
}
if let Some(selected_tools) = restored_tool_selection {
self.set_mcp_tool_selection(selected_tools).await;
}
// Defer seeding the session's initial context until the first turn starts so
// turn/start overrides can be merged before we write to the rollout.
self.flush_rollout().await;
}
InitialHistory::Forked(rollout_items) => {
let restored_tool_selection =
Self::extract_mcp_tool_selection_from_rollout(&rollout_items);
let previous_model = Self::last_rollout_model_name(&rollout_items)
.map(std::string::ToString::to_string);
self.set_previous_model(previous_model).await;
@@ -1603,6 +1618,9 @@ impl Session {
let mut state = self.state.lock().await;
state.set_token_info(Some(info));
}
if let Some(selected_tools) = restored_tool_selection {
self.set_mcp_tool_selection(selected_tools).await;
}
// If persisting, persist all rollout items as-is (recorder filters)
if !rollout_items.is_empty() {
@@ -1644,6 +1662,54 @@ impl Session {
})
}
fn extract_mcp_tool_selection_from_rollout(
rollout_items: &[RolloutItem],
) -> Option<Vec<String>> {
let mut search_call_ids = HashSet::new();
let mut active_selected_tools: Option<Vec<String>> = None;
for item in rollout_items {
let RolloutItem::ResponseItem(response_item) = item else {
continue;
};
match response_item {
ResponseItem::FunctionCall { name, call_id, .. } => {
if name == SEARCH_TOOL_BM25_TOOL_NAME {
search_call_ids.insert(call_id.clone());
}
}
ResponseItem::FunctionCallOutput { call_id, output } => {
if !search_call_ids.contains(call_id) {
continue;
}
let Some(content) = output.body.to_text() else {
continue;
};
let Ok(payload) = serde_json::from_str::<Value>(&content) else {
continue;
};
let Some(selected_tools) = payload
.get("active_selected_tools")
.and_then(Value::as_array)
else {
continue;
};
let Some(selected_tools) = selected_tools
.iter()
.map(|value| value.as_str().map(str::to_string))
.collect::<Option<Vec<_>>>()
else {
continue;
};
active_selected_tools = Some(selected_tools);
}
_ => {}
}
}
active_selected_tools
}
async fn previous_model(&self) -> Option<String> {
let state = self.state.lock().await;
state.previous_model()
@@ -4899,6 +4965,10 @@ async fn built_tools(
None
};
let app_tools = connectors
.as_ref()
.map(|connectors| filter_codex_apps_mcp_tools(&mcp_tools, connectors));
if let Some(connectors) = connectors.as_ref() {
let skill_name_counts_lower = skills_outcome.map_or_else(HashMap::new, |outcome| {
build_skill_name_counts(&outcome.skills, &outcome.disabled_paths).1
@@ -4911,24 +4981,20 @@ async fn built_tools(
&skill_name_counts_lower,
);
let mut selected_mcp_tools =
if let Some(selected_tools) = sess.get_mcp_tool_selection().await {
filter_mcp_tools_by_name(&mcp_tools, &selected_tools)
} else {
HashMap::new()
};
let mut selected_mcp_tools = filter_non_codex_apps_mcp_tools_only(&mcp_tools);
let apps_mcp_tools =
filter_codex_apps_mcp_tools_only(&mcp_tools, explicitly_enabled.as_ref());
selected_mcp_tools.extend(apps_mcp_tools);
if let Some(selected_tools) = sess.get_mcp_tool_selection().await {
selected_mcp_tools.extend(filter_mcp_tools_by_name(&mcp_tools, &selected_tools));
}
selected_mcp_tools.extend(filter_codex_apps_mcp_tools_only(
&mcp_tools,
explicitly_enabled.as_ref(),
));
mcp_tools = selected_mcp_tools;
}
let app_tools = connectors
.as_ref()
.map(|connectors| filter_codex_apps_mcp_tools(&mcp_tools, connectors));
Ok(Arc::new(ToolRouter::from_config(
&turn_context.tools_config,
Some(
@@ -5789,6 +5855,22 @@ mod tests {
}
}
fn function_call_rollout_item(name: &str, call_id: &str) -> RolloutItem {
RolloutItem::ResponseItem(ResponseItem::FunctionCall {
id: None,
name: name.to_string(),
arguments: "{}".to_string(),
call_id: call_id.to_string(),
})
}
fn function_call_output_rollout_item(call_id: &str, output: &str) -> RolloutItem {
RolloutItem::ResponseItem(ResponseItem::FunctionCallOutput {
call_id: call_id.to_string(),
output: FunctionCallOutputPayload::from_text(output.to_string()),
})
}
#[tokio::test]
async fn get_base_instructions_no_user_content() {
let prompt_with_apply_patch_instructions =
@@ -5978,6 +6060,46 @@ mod tests {
assert_eq!(connector_ids, HashSet::<String>::new());
}
#[test]
fn non_app_mcp_tools_remain_visible_without_search_selection() {
let mcp_tools = HashMap::from([
(
"mcp__codex_apps__calendar_create_event".to_string(),
make_mcp_tool(
CODEX_APPS_MCP_SERVER_NAME,
"calendar_create_event",
Some("calendar"),
Some("Calendar"),
),
),
(
"mcp__rmcp__echo".to_string(),
make_mcp_tool("rmcp", "echo", None, None),
),
]);
let mut selected_mcp_tools = mcp_tools
.iter()
.filter(|(_, tool)| tool.server_name != CODEX_APPS_MCP_SERVER_NAME)
.map(|(name, tool)| (name.clone(), tool.clone()))
.collect::<HashMap<_, _>>();
let connectors = connectors::accessible_connectors_from_mcp_tools(&mcp_tools);
let explicitly_enabled_connectors = HashSet::new();
let connectors = filter_connectors_for_input(
&connectors,
&[user_message("run echo")],
&explicitly_enabled_connectors,
&HashMap::new(),
);
let apps_mcp_tools = filter_codex_apps_mcp_tools_only(&mcp_tools, &connectors);
selected_mcp_tools.extend(apps_mcp_tools);
let mut tool_names: Vec<String> = selected_mcp_tools.into_keys().collect();
tool_names.sort();
assert_eq!(tool_names, vec!["mcp__rmcp__echo".to_string()]);
}
#[test]
fn search_tool_selection_keeps_codex_apps_tools_without_mentions() {
let selected_tool_names = vec![
@@ -6065,6 +6187,106 @@ mod tests {
);
}
#[test]
fn extract_mcp_tool_selection_from_rollout_reads_search_tool_output() {
let rollout_items = vec![
function_call_rollout_item(SEARCH_TOOL_BM25_TOOL_NAME, "search-1"),
function_call_output_rollout_item(
"search-1",
&json!({
"active_selected_tools": [
"mcp__codex_apps__calendar_create_event",
"mcp__codex_apps__calendar_list_events",
],
})
.to_string(),
),
];
let selected = Session::extract_mcp_tool_selection_from_rollout(&rollout_items);
assert_eq!(
selected,
Some(vec![
"mcp__codex_apps__calendar_create_event".to_string(),
"mcp__codex_apps__calendar_list_events".to_string(),
])
);
}
#[test]
fn extract_mcp_tool_selection_from_rollout_latest_valid_payload_wins() {
let rollout_items = vec![
function_call_rollout_item(SEARCH_TOOL_BM25_TOOL_NAME, "search-1"),
function_call_output_rollout_item(
"search-1",
&json!({
"active_selected_tools": ["mcp__codex_apps__calendar_create_event"],
})
.to_string(),
),
function_call_rollout_item(SEARCH_TOOL_BM25_TOOL_NAME, "search-2"),
function_call_output_rollout_item(
"search-2",
&json!({
"active_selected_tools": ["mcp__codex_apps__calendar_delete_event"],
})
.to_string(),
),
];
let selected = Session::extract_mcp_tool_selection_from_rollout(&rollout_items);
assert_eq!(
selected,
Some(vec!["mcp__codex_apps__calendar_delete_event".to_string(),])
);
}
#[test]
fn extract_mcp_tool_selection_from_rollout_ignores_non_search_and_malformed_payloads() {
let rollout_items = vec![
function_call_rollout_item("shell", "shell-1"),
function_call_output_rollout_item(
"shell-1",
&json!({
"active_selected_tools": ["mcp__codex_apps__should_be_ignored"],
})
.to_string(),
),
function_call_rollout_item(SEARCH_TOOL_BM25_TOOL_NAME, "search-1"),
function_call_output_rollout_item("search-1", "{not-json"),
function_call_output_rollout_item(
"unknown-search-call",
&json!({
"active_selected_tools": ["mcp__codex_apps__also_ignored"],
})
.to_string(),
),
function_call_output_rollout_item(
"search-1",
&json!({
"active_selected_tools": ["mcp__codex_apps__calendar_list_events"],
})
.to_string(),
),
];
let selected = Session::extract_mcp_tool_selection_from_rollout(&rollout_items);
assert_eq!(
selected,
Some(vec!["mcp__codex_apps__calendar_list_events".to_string(),])
);
}
#[test]
fn extract_mcp_tool_selection_from_rollout_returns_none_without_valid_search_output() {
let rollout_items = vec![function_call_rollout_item(
SEARCH_TOOL_BM25_TOOL_NAME,
"search-1",
)];
let selected = Session::extract_mcp_tool_selection_from_rollout(&rollout_items);
assert_eq!(selected, None);
}
#[tokio::test]
async fn reconstruct_history_matches_live_compactions() {
let (session, turn_context) = make_session_and_context().await;