Refresh plugin recommendations when tooling state changes (#61387)

## Why

Cached recommendations can keep offering a plugin after installation. Refreshing the shared catalog during step preparation can also leave the recommendation prompt and install handler with different candidates.

## What changed

- Check `tooling_changed_at` before reusing recommendations, with a one-second reuse window and a 250 ms deadline that includes authentication. Refresh the catalog when the marker changes, preserving cached recommendations on failure and retrying on a later poll.
- Scope recommendation and featured-plugin caches by product, backend, and account identity, and share concurrent recommendation refreshes.
- Use the same recommendation snapshot for a step's prompt and install handler, while allowing later steps to observe updated recommendations.

## Testing

Add coverage for marker changes, missing markers, refresh failures, concurrent and cancelled requests, and the authentication deadline. Extend integration tests to verify consistent step candidates and stop offering an installed plugin while preserving the original recommendation block in conversation history.

GitOrigin-RevId: 44fa6723a623d9604f0c768118bc6889d2d3b693
This commit is contained in:
Alex Zamoshchin
2026-11-13 22:16:32 +00:00
committed by copyberry
parent 826b6e8199
commit 42f950cd11
7 changed files with 460 additions and 327 deletions
+27 -20
View File
@@ -309,6 +309,7 @@ use self::session::SessionSettingsCommit;
pub(crate) use self::session::SessionSettingsUpdate;
#[cfg(test)]
use self::turn::AssistantMessageStreamParsers;
use self::turn::PreparedToolRecommendations;
use self::turn::RealtimeEventText;
use self::turn::agent_message_text;
#[cfg(test)]
@@ -347,7 +348,9 @@ use crate::tools::network_approval::build_blocked_request_observer;
use crate::tools::network_approval::build_network_policy_decider;
#[cfg(test)]
use crate::tools::parallel::ToolCallRuntime;
use crate::tools::router::ToolSuggestCandidates;
use crate::tools::sandboxing::ApprovalStore;
use crate::tools::spec_plan::tool_suggest_enabled;
use crate::turn_timing::TurnTimingState;
use crate::turn_timing::record_turn_ttfm_metric;
use crate::unified_exec::UnifiedExecProcessManager;
@@ -356,7 +359,6 @@ use crate::windows_sandbox::local_binding_policy_for_sandbox;
use crate::windows_sandbox::managed_proxy_routing_for_windows_sandbox;
use codex_core_plugins::PluginCommandAttribution;
use codex_core_plugins::PluginsManager;
use codex_core_plugins::RecommendedPluginCandidatesInput;
use codex_git_utils::get_git_repo_root;
use codex_history::CodexHarnessMetadata;
use codex_history::CompactedItem;
@@ -4488,7 +4490,13 @@ impl Session {
&required_servers,
required_plugins,
)),
turn::prepare_tool_recommendations(self.as_ref(), turn_context.as_ref()),
async {
if tool_suggest_enabled(&turn_context) {
Some(PreparedToolRecommendations::load(self, &turn_context).await)
} else {
None
}
},
);
let search = crate::model_search::resolve(
turn_context.as_ref(),
@@ -5134,25 +5142,25 @@ impl Session {
developer_sections
.push(DeveloperInstructions::new(developer_instructions).render_fragment());
}
let recommended_plugin_candidates = if turn_context
.config
.features
.plugin_recommendations_enabled()
let step_candidates = step_context.extension_data.get::<ToolSuggestCandidates>();
// Prompt-only recommendations have no install handler to capture candidates for the step.
let prompt_only_candidates = if !tool_suggest_enabled(turn_context)
&& turn_context
.config
.features
.plugin_recommendations_enabled()
{
let auth = self.services.auth_manager.auth().await;
let plugins_config = turn_context.config.plugins_config_input();
self.services
.plugins_manager
.recommended_plugin_candidates_for_config(RecommendedPluginCandidatesInput {
plugins_config: &plugins_config,
auth: auth.as_ref(),
disabled_tools: &turn_context.config.tool_suggest.disabled_tools,
app_server_client_name: turn_context.app_server_client_name.as_deref(),
})
.await
match PreparedToolRecommendations::load(self, turn_context).await {
PreparedToolRecommendations::Endpoint(tools) => Some(tools),
PreparedToolRecommendations::Legacy(_) => None,
}
} else {
None
};
let recommended_plugin_candidates = step_candidates
.as_deref()
.map(|candidates| candidates.tools.as_slice())
.or(prompt_only_candidates.as_deref());
let context_contributors = self.services.extensions.context_contributors().to_vec();
for contributor in &context_contributors {
for fragment in contributor
@@ -5301,9 +5309,8 @@ impl Session {
}
if next_item.is_none()
&& let Some(recommended_plugins) = recommended_plugin_candidates
.as_deref()
.and_then(|plugins| {
&& let Some(recommended_plugins) =
recommended_plugin_candidates.and_then(|plugins| {
RecommendedPluginsInstructions::from_plugins(
plugins,
&ResolvedModelMessages::from_model(&step_context.settings.model_info)
+61 -83
View File
@@ -67,7 +67,6 @@ use crate::tools::router::ToolSuggestCandidates;
use crate::tools::router::ToolSuggestPresentation;
use crate::tools::router::record_registration_metrics;
use crate::tools::spec_plan::build_tool_router;
use crate::tools::spec_plan::tool_suggest_enabled;
use crate::turn_diff_tracker::TurnDiffTracker;
use crate::turn_timing::record_turn_ttft_metric;
use crate::util::error_or_panic;
@@ -2095,40 +2094,29 @@ async fn run_sampling_request(
}
}
pub(crate) struct PreparedToolRecommendations {
auth: Option<CodexAuth>,
endpoint_candidates: Option<Vec<DiscoverableTool>>,
pub(super) enum PreparedToolRecommendations {
Endpoint(Vec<DiscoverableTool>),
Legacy(Option<CodexAuth>),
}
#[instrument(level = "trace", skip_all)]
pub(crate) async fn prepare_tool_recommendations(
sess: &Session,
turn_context: &TurnContext,
) -> PreparedToolRecommendations {
let tool_suggest_is_enabled = tool_suggest_enabled(turn_context);
let auth = if tool_suggest_is_enabled {
sess.services.auth_manager.auth().await
} else {
None
};
let endpoint_candidates = if tool_suggest_is_enabled {
let plugins_config = turn_context.config.plugins_config_input();
sess.services
impl PreparedToolRecommendations {
#[instrument(level = "trace", skip_all)]
pub(super) async fn load(sess: &Session, turn_context: &TurnContext) -> Self {
let auth = sess.services.auth_manager.auth().await;
let endpoint_candidates = sess
.services
.plugins_manager
.recommended_plugin_candidates_for_config(RecommendedPluginCandidatesInput {
plugins_config: &plugins_config,
plugins_config: &turn_context.config.plugins_config_input(),
auth: auth.as_ref(),
disabled_tools: &turn_context.config.tool_suggest.disabled_tools,
app_server_client_name: turn_context.app_server_client_name.as_deref(),
})
.await
} else {
None
};
PreparedToolRecommendations {
auth,
endpoint_candidates,
.await;
match endpoint_candidates {
Some(tools) => Self::Endpoint(tools),
None => Self::Legacy(auth),
}
}
}
@@ -2141,77 +2129,61 @@ pub(crate) async fn prepare_tool_recommendations(
apps_enabled = turn_context.apps_enabled()
)
)]
pub(crate) async fn built_tools(
pub(super) async fn built_tools(
sess: &Session,
turn_context: &TurnContext,
model_info: &codex_protocol::openai_models::ModelInfo,
environments: &TurnEnvironmentSnapshot,
mcp: &Arc<codex_mcp::McpBinding>,
step_store: &ExtensionData,
prepared_recommendations: PreparedToolRecommendations,
prepared_recommendations: Option<PreparedToolRecommendations>,
dynamic_tools: &[codex_protocol::dynamic_tools::DynamicToolSpec],
) -> CodexResult<Arc<ToolRouter>> {
let all_mcp_tools = mcp.tools();
let connector_snapshot = mcp.config().connector_snapshot.clone();
let apps_enabled = turn_context.apps_enabled();
let accessible_connectors =
apps_enabled.then(|| connectors::accessible_connectors_from_mcp_tools(all_mcp_tools));
let tool_suggest_is_enabled = tool_suggest_enabled(turn_context);
let PreparedToolRecommendations {
auth,
endpoint_candidates: endpoint_recommended_plugin_candidates,
} = prepared_recommendations;
let tool_suggest_candidates =
if let Some(recommended_plugin_candidates) = endpoint_recommended_plugin_candidates {
Some(ToolSuggestCandidates {
tools: recommended_plugin_candidates,
presentation: ToolSuggestPresentation::RecommendationContext,
})
} else {
let loaded_plugin_app_connector_ids = connector_snapshot
let tool_suggest_candidates = match prepared_recommendations {
Some(PreparedToolRecommendations::Endpoint(tools)) => Some(ToolSuggestCandidates {
tools,
presentation: ToolSuggestPresentation::RecommendationContext,
}),
Some(PreparedToolRecommendations::Legacy(auth)) if apps_enabled => {
let accessible_connectors =
connectors::accessible_connectors_from_mcp_tools(mcp.tools());
let loaded_plugin_app_connector_ids = mcp
.config()
.connector_snapshot
.connector_ids()
.iter()
.map(|connector_id| connector_id.0.clone())
.collect::<Vec<_>>();
async {
if apps_enabled && tool_suggest_is_enabled {
if let Some(accessible_connectors) = accessible_connectors.as_ref() {
match connectors::list_tool_suggest_discoverable_tools_with_auth(
&turn_context.config,
sess.services.plugins_manager.as_ref(),
auth.as_ref(),
accessible_connectors.as_slice(),
&loaded_plugin_app_connector_ids,
)
.await
.map(|discoverable_tools| {
filter_request_plugin_install_discoverable_tools_for_client(
discoverable_tools,
turn_context.app_server_client_name.as_deref(),
)
}) {
Ok(discoverable_tools) if discoverable_tools.is_empty() => None,
Ok(discoverable_tools) => Some(ToolSuggestCandidates {
tools: discoverable_tools,
presentation: ToolSuggestPresentation::ListTool,
}),
Err(err) => {
warn!("failed to load discoverable tool suggestions: {err:#}");
None
}
}
} else {
None
}
} else {
match connectors::list_tool_suggest_discoverable_tools_with_auth(
&turn_context.config,
sess.services.plugins_manager.as_ref(),
auth.as_ref(),
&accessible_connectors,
&loaded_plugin_app_connector_ids,
)
.instrument(trace_span!("built_tools.load_discoverable_tools"))
.await
{
Ok(tools) => {
let tools = filter_request_plugin_install_discoverable_tools_for_client(
tools,
turn_context.app_server_client_name.as_deref(),
);
(!tools.is_empty()).then_some(ToolSuggestCandidates {
tools,
presentation: ToolSuggestPresentation::ListTool,
})
}
Err(err) => {
warn!("failed to load discoverable tool suggestions: {err:#}");
None
}
}
.instrument(trace_span!("built_tools.load_discoverable_tools"))
.await
};
Ok(Arc::new(build_tool_router(
}
Some(PreparedToolRecommendations::Legacy(_)) | None => None,
};
let router = build_tool_router(
sess,
turn_context,
model_info,
@@ -2221,7 +2193,13 @@ pub(crate) async fn built_tools(
step_store,
tool_suggest_candidates.as_ref(),
dynamic_tools,
)?))
)?;
if let Some(candidates) = tool_suggest_candidates
&& candidates.presentation == ToolSuggestPresentation::RecommendationContext
{
step_store.insert(candidates);
}
Ok(Arc::new(router))
}
#[derive(Debug)]
@@ -382,19 +382,38 @@ async fn resolve_install_elicitation(
Ok(())
}
async fn mount_remote_calendar_recommendation(server: &wiremock::MockServer) {
async fn mount_remote_calendar_recommendation(
server: &wiremock::MockServer,
installed: Arc<AtomicBool>,
) {
let installed_for_marker = Arc::clone(&installed);
Mock::given(method("GET"))
.and(path("/ps/plugins/tooling-state/turtle"))
.respond_with(move |_: &wiremock::Request| {
let marker = installed_for_marker
.load(Ordering::SeqCst)
.then_some("1700000000000");
ResponseTemplate::new(200).set_body_json(json!({ "tooling_changed_at": marker }))
})
.mount(server)
.await;
Mock::given(method("GET"))
.and(path("/ps/plugins/suggested/codex"))
.and(query_param("scope", "GLOBAL"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
"enabled": true,
"plugins": [{
"id": REMOTE_CALENDAR_PLUGIN_ID,
"name": "calendar",
"display_name": "Calendar"
}]
})))
.expect(1)
.and(query_param("includeInstalled", "false"))
.respond_with(move |_: &wiremock::Request| {
let plugins = if installed.load(Ordering::SeqCst) {
json!([])
} else {
json!([{
"id": REMOTE_CALENDAR_PLUGIN_ID,
"name": "calendar",
"display_name": "Calendar"
}])
};
ResponseTemplate::new(200).set_body_json(json!({ "enabled": true, "plugins": plugins }))
})
.expect(2)
.mount(server)
.await;
Mock::given(method("GET"))
@@ -788,6 +807,18 @@ async fn mcp_discovery_overlaps_endpoint_plugin_recommendations() -> Result<()>
response.requests().is_empty(),
"sampling should wait for the complete MCP catalog"
);
// Change the shared catalog after the step's recommendation request, while MCP is blocked.
Mock::given(method("GET"))
.and(path("/ps/plugins/suggested/codex"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
"enabled": true, "plugins": []
})))
.with_priority(1)
.mount(&server)
.await;
test.thread_manager
.plugins_manager()
.clear_recommended_plugins_cache();
test.fs()
.write_file(
&barrier,
@@ -809,6 +840,11 @@ async fn mcp_discovery_overlaps_endpoint_plugin_recommendations() -> Result<()>
.contains("github@openai-curated-remote"),
"the completed request should preserve endpoint recommendations"
);
assert!(
tool_names(&request.body_json())
.iter()
.any(|name| name == REQUEST_PLUGIN_INSTALL_TOOL_NAME)
);
assert!(
request
.tool_by_name("mcp__step_preparation", "echo")
@@ -1541,7 +1577,8 @@ async fn run_remote_plugin_install_refresh_case(refreshed_tools: RefreshedAppsTo
}
RefreshedAppsTools::Missing => AppsTestServer::mount_without_tools(&server).await?,
};
mount_remote_calendar_recommendation(&server).await;
let installed = Arc::new(AtomicBool::new(false));
mount_remote_calendar_recommendation(&server, Arc::clone(&installed)).await;
let initial_remote_installed_plugins = mount_empty_remote_installed_plugins(&server).await;
let install_call_id = "install-calendar";
@@ -1579,6 +1616,7 @@ async fn run_remote_plugin_install_refresh_case(refreshed_tools: RefreshedAppsTo
let elicitation = start_install_turn(&test, "use Calendar").await?;
mount_remote_calendar_installed_plugins(&server).await;
drop(initial_remote_installed_plugins);
installed.store(true, Ordering::SeqCst);
tools_available.store(
matches!(refreshed_tools, RefreshedAppsTools::Available),
Ordering::SeqCst,
@@ -1623,10 +1661,32 @@ async fn run_remote_plugin_install_refresh_case(refreshed_tools: RefreshedAppsTo
.codex
.refresh_runtime_config(current_config, test.config.clone())
.await;
// Let the one-second marker poll window expire before the next turn.
tokio::time::sleep(Duration::from_secs(1)).await;
test.submit_turn("check whether Calendar is still installed")
.await?;
let requests = mock.requests();
assert_eq!(requests.len(), 3);
let recommendation_blocks = |index: usize| {
requests[index]
.message_input_texts("developer")
.into_iter()
.filter(|text| text.starts_with("<recommended_plugins>"))
.collect::<Vec<_>>()
};
let initial = recommendation_blocks(0);
assert_eq!(initial.len(), 1);
assert_eq!(
recommendation_blocks(2),
initial,
"refreshing recommendations must preserve the original block"
);
assert!(
!tool_names(&requests[2].body_json())
.iter()
.any(|name| name == REQUEST_PLUGIN_INSTALL_TOOL_NAME),
"the updated recommendations must stop offering to install Calendar"
);
assert_eq!(
requests[2]
.tool_by_name(CALENDAR_NAMESPACE, CALENDAR_CREATE_EVENT_TOOL)