mirror of
https://github.com/openai/codex.git
synced 2026-11-14 14:49:31 +00:00
Reuse MCP bindings while cached servers remain dormant (#44121)
## Why Dormant MCP servers with usable cached tool catalogs prevented binding reuse, causing each model step to capture a new binding even when the catalog was unchanged. ## What changed Track dormant catalog revisions alongside ready client revisions so bindings can be reused until the catalog changes or the server starts. Explicit server and plugin requirements still trigger binding capture and startup. Check cached catalog availability without cloning tool definitions. ## Testing Extend the lazy-startup integration test to verify that two model steps share one binding while the server stays dormant, and that starting the server invalidates that binding. GitOrigin-RevId: ef0469a85000f84a928946c24350e3bda1522fb5
This commit is contained in:
1 parent
9e868bd9dc
commit
20f109eadb
6 files changed
+156
-16
No files matched your search
@@ -21,6 +21,7 @@ use startup::emit_update;
|
|||||||
use startup::mcp_init_error_display;
|
use startup::mcp_init_error_display;
|
||||||
use startup::mcp_startup_failure_reason;
|
use startup::mcp_startup_failure_reason;
|
||||||
use startup::should_share_codex_apps_tools_cache;
|
use startup::should_share_codex_apps_tools_cache;
|
||||||
|
pub(crate) use tool_catalog::BindingCatalogRevision;
|
||||||
pub use tool_catalog::tool_is_model_visible;
|
pub use tool_catalog::tool_is_model_visible;
|
||||||
|
|
||||||
use std::collections::HashMap;
|
use std::collections::HashMap;
|
||||||
|
|||||||
@@ -58,10 +58,19 @@ pub fn tool_is_model_visible(tool: &ToolInfo) -> bool {
|
|||||||
.any(|target| target.as_str() == Some(MCP_UI_MODEL_VISIBILITY))
|
.any(|target| target.as_str() == Some(MCP_UI_MODEL_VISIBILITY))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Catalog identity within one published connection set, including dormant servers.
|
||||||
|
#[derive(PartialEq, Eq)]
|
||||||
|
pub(crate) enum BindingCatalogRevision {
|
||||||
|
Ready(ClientToolCatalogRevision),
|
||||||
|
Dormant(u64),
|
||||||
|
}
|
||||||
|
|
||||||
impl McpConnectionSet {
|
impl McpConnectionSet {
|
||||||
pub(crate) async fn stable_catalog_revisions(
|
pub(crate) async fn stable_catalog_revisions(
|
||||||
&self,
|
&self,
|
||||||
) -> Option<HashMap<String, ClientToolCatalogRevision>> {
|
required_servers: &[String],
|
||||||
|
required_plugins: &HashSet<String>,
|
||||||
|
) -> Option<HashMap<String, BindingCatalogRevision>> {
|
||||||
let mut revisions = HashMap::new();
|
let mut revisions = HashMap::new();
|
||||||
for (server_name, view) in &self.servers {
|
for (server_name, view) in &self.servers {
|
||||||
if !view
|
if !view
|
||||||
@@ -69,9 +78,33 @@ impl McpConnectionSet {
|
|||||||
.client
|
.client
|
||||||
.startup_complete
|
.startup_complete
|
||||||
.load(Ordering::Acquire)
|
.load(Ordering::Acquire)
|
||||||
|
{
|
||||||
|
// A cached catalog can remain stable without starting its server.
|
||||||
|
// Explicit requirements must still start the server during binding capture.
|
||||||
|
if !view.connection.startup_is_dormant()
|
||||||
|
|| view.connection.client.cancel_token.is_cancelled()
|
||||||
|
|| required_servers
|
||||||
|
.iter()
|
||||||
|
.any(|required| required == server_name)
|
||||||
|
|| (self.is_selected_plugin_mcp_server(server_name)
|
||||||
|
&& self
|
||||||
|
.plugin_id_for_mcp_server_name(server_name)
|
||||||
|
.is_some_and(|plugin_id| required_plugins.contains(plugin_id)))
|
||||||
{
|
{
|
||||||
return None;
|
return None;
|
||||||
}
|
}
|
||||||
|
let revision = view
|
||||||
|
.connection
|
||||||
|
.client
|
||||||
|
.tool_catalog_cache_context
|
||||||
|
.as_ref()?
|
||||||
|
.current_revision()?;
|
||||||
|
revisions.insert(
|
||||||
|
server_name.clone(),
|
||||||
|
BindingCatalogRevision::Dormant(revision),
|
||||||
|
);
|
||||||
|
continue;
|
||||||
|
}
|
||||||
let Some(client) = view.connection.client.ready_client() else {
|
let Some(client) = view.connection.client.ready_client() else {
|
||||||
if !view.connection.client.is_codex_apps_mcp_server
|
if !view.connection.client.is_codex_apps_mcp_server
|
||||||
&& self.required_servers.binary_search(server_name).is_err()
|
&& self.required_servers.binary_search(server_name).is_err()
|
||||||
@@ -87,10 +120,10 @@ impl McpConnectionSet {
|
|||||||
let revision = client.tool_catalog.read(|catalog| catalog.revision).await;
|
let revision = client.tool_catalog.read(|catalog| catalog.revision).await;
|
||||||
revisions.insert(
|
revisions.insert(
|
||||||
server_name.clone(),
|
server_name.clone(),
|
||||||
ClientToolCatalogRevision {
|
BindingCatalogRevision::Ready(ClientToolCatalogRevision {
|
||||||
catalog: Arc::clone(&client.tool_catalog),
|
catalog: Arc::clone(&client.tool_catalog),
|
||||||
revision,
|
revision,
|
||||||
},
|
}),
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
Some(revisions)
|
Some(revisions)
|
||||||
@@ -168,6 +201,7 @@ impl McpConnectionSet {
|
|||||||
(tools, errors)
|
(tools, errors)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[instrument(level = "trace", skip_all)]
|
||||||
pub(crate) async fn capture_binding_with_metadata(
|
pub(crate) async fn capture_binding_with_metadata(
|
||||||
self: &Arc<Self>,
|
self: &Arc<Self>,
|
||||||
config: Arc<crate::McpConfig>,
|
config: Arc<crate::McpConfig>,
|
||||||
|
|||||||
@@ -2731,7 +2731,15 @@ async fn capture_binding_skips_pending_optional_servers_after_configured_shared_
|
|||||||
required_manager.required_servers = vec!["pending-selected".to_string()];
|
required_manager.required_servers = vec!["pending-selected".to_string()];
|
||||||
|
|
||||||
let manager = Arc::new(manager);
|
let manager = Arc::new(manager);
|
||||||
assert!(manager.stable_catalog_revisions().await.is_none());
|
assert!(
|
||||||
|
manager
|
||||||
|
.stable_catalog_revisions(
|
||||||
|
/*required_servers*/ &[],
|
||||||
|
/*required_plugins*/ &HashSet::new()
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.is_none()
|
||||||
|
);
|
||||||
let started = tokio::time::Instant::now();
|
let started = tokio::time::Instant::now();
|
||||||
let binding = tokio::time::timeout(
|
let binding = tokio::time::timeout(
|
||||||
Duration::from_millis(500),
|
Duration::from_millis(500),
|
||||||
@@ -2887,9 +2895,25 @@ async fn stable_catalog_revisions_ignore_terminal_optional_server_failures() {
|
|||||||
assert!(failed.client().await.is_err());
|
assert!(failed.client().await.is_err());
|
||||||
manager.insert_test_client("failed", failed);
|
manager.insert_test_client("failed", failed);
|
||||||
|
|
||||||
assert!(manager.stable_catalog_revisions().await.is_some());
|
assert!(
|
||||||
|
manager
|
||||||
|
.stable_catalog_revisions(
|
||||||
|
/*required_servers*/ &[],
|
||||||
|
/*required_plugins*/ &HashSet::new()
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.is_some()
|
||||||
|
);
|
||||||
manager.required_servers.push("failed".to_string());
|
manager.required_servers.push("failed".to_string());
|
||||||
assert!(manager.stable_catalog_revisions().await.is_none());
|
assert!(
|
||||||
|
manager
|
||||||
|
.stable_catalog_revisions(
|
||||||
|
/*required_servers*/ &[],
|
||||||
|
/*required_plugins*/ &HashSet::new()
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.is_none()
|
||||||
|
);
|
||||||
manager.required_servers.clear();
|
manager.required_servers.clear();
|
||||||
|
|
||||||
let binding = capture_binding(&Arc::new(manager)).await;
|
let binding = capture_binding(&Arc::new(manager)).await;
|
||||||
@@ -3840,7 +3864,15 @@ async fn list_all_tools_reconnects_failed_codex_apps_startup_and_reuses_client()
|
|||||||
};
|
};
|
||||||
let manager = Arc::new(manager);
|
let manager = Arc::new(manager);
|
||||||
|
|
||||||
assert!(manager.stable_catalog_revisions().await.is_none());
|
assert!(
|
||||||
|
manager
|
||||||
|
.stable_catalog_revisions(
|
||||||
|
/*required_servers*/ &[],
|
||||||
|
/*required_plugins*/ &HashSet::new()
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.is_none()
|
||||||
|
);
|
||||||
let reconnect_finished_wait = reconnect_finished.notified();
|
let reconnect_finished_wait = reconnect_finished.notified();
|
||||||
let tools = manager.list_all_tools().await;
|
let tools = manager.list_all_tools().await;
|
||||||
assert!(tools.is_empty());
|
assert!(tools.is_empty());
|
||||||
@@ -3855,7 +3887,15 @@ async fn list_all_tools_reconnects_failed_codex_apps_startup_and_reuses_client()
|
|||||||
vec!["drive_search"]
|
vec!["drive_search"]
|
||||||
);
|
);
|
||||||
assert_eq!(attempts.load(std::sync::atomic::Ordering::SeqCst), 1);
|
assert_eq!(attempts.load(std::sync::atomic::Ordering::SeqCst), 1);
|
||||||
assert!(manager.stable_catalog_revisions().await.is_some());
|
assert!(
|
||||||
|
manager
|
||||||
|
.stable_catalog_revisions(
|
||||||
|
/*required_servers*/ &[],
|
||||||
|
/*required_plugins*/ &HashSet::new()
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.is_some()
|
||||||
|
);
|
||||||
|
|
||||||
let step = capture_binding(&manager).await;
|
let step = capture_binding(&manager).await;
|
||||||
let prepared = step
|
let prepared = step
|
||||||
|
|||||||
@@ -46,8 +46,8 @@ use tokio_util::sync::CancellationToken;
|
|||||||
|
|
||||||
use crate::McpConfig;
|
use crate::McpConfig;
|
||||||
use crate::binding::McpBinding;
|
use crate::binding::McpBinding;
|
||||||
use crate::client_tool_catalog::ClientToolCatalogRevision;
|
|
||||||
use crate::client_tool_catalog::CodexAppsToolSnapshot;
|
use crate::client_tool_catalog::CodexAppsToolSnapshot;
|
||||||
|
use crate::connection_manager::BindingCatalogRevision;
|
||||||
use crate::connection_manager::McpConnectionSet;
|
use crate::connection_manager::McpConnectionSet;
|
||||||
use crate::elicitation::ElicitationLifecycle;
|
use crate::elicitation::ElicitationLifecycle;
|
||||||
use crate::elicitation::ElicitationRequestRouter;
|
use crate::elicitation::ElicitationRequestRouter;
|
||||||
@@ -119,7 +119,7 @@ struct PublishedMcpRuntime {
|
|||||||
}
|
}
|
||||||
|
|
||||||
struct CachedMcpBinding {
|
struct CachedMcpBinding {
|
||||||
catalog_revisions: HashMap<String, ClientToolCatalogRevision>,
|
catalog_revisions: HashMap<String, BindingCatalogRevision>,
|
||||||
binding: Arc<McpBinding>,
|
binding: Arc<McpBinding>,
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -363,7 +363,10 @@ impl McpRuntime {
|
|||||||
required_plugins: &HashSet<String>,
|
required_plugins: &HashSet<String>,
|
||||||
) -> Option<Arc<McpBinding>> {
|
) -> Option<Arc<McpBinding>> {
|
||||||
let config = Arc::clone(current.config.as_ref()?);
|
let config = Arc::clone(current.config.as_ref()?);
|
||||||
let stable_catalog_revisions = current.connections.stable_catalog_revisions().await;
|
let stable_catalog_revisions = current
|
||||||
|
.connections
|
||||||
|
.stable_catalog_revisions(required_servers, required_plugins)
|
||||||
|
.await;
|
||||||
if let Some(catalog_revisions) = &stable_catalog_revisions {
|
if let Some(catalog_revisions) = &stable_catalog_revisions {
|
||||||
let cached = current
|
let cached = current
|
||||||
.cached_binding
|
.cached_binding
|
||||||
@@ -390,7 +393,7 @@ impl McpRuntime {
|
|||||||
if let Some(catalog_revisions) = stable_catalog_revisions
|
if let Some(catalog_revisions) = stable_catalog_revisions
|
||||||
&& current
|
&& current
|
||||||
.connections
|
.connections
|
||||||
.stable_catalog_revisions()
|
.stable_catalog_revisions(required_servers, required_plugins)
|
||||||
.await
|
.await
|
||||||
.as_ref()
|
.as_ref()
|
||||||
== Some(&catalog_revisions)
|
== Some(&catalog_revisions)
|
||||||
|
|||||||
@@ -116,7 +116,18 @@ impl Default for ToolCatalogCacheEntry {
|
|||||||
|
|
||||||
impl McpToolCatalogCacheContext {
|
impl McpToolCatalogCacheContext {
|
||||||
pub(crate) fn has_tools(&self) -> bool {
|
pub(crate) fn has_tools(&self) -> bool {
|
||||||
self.current_tools().is_some_and(|tools| !tools.is_empty())
|
self.current_revision().is_some()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Identifies the current usable catalog without cloning its tool definitions.
|
||||||
|
/// The entry is fixed for a connection; accepted publications advance its revision.
|
||||||
|
pub(crate) fn current_revision(&self) -> Option<u64> {
|
||||||
|
let state = lock_unpoisoned(&self.entry.state);
|
||||||
|
let snapshot = state.snapshot.as_ref()?;
|
||||||
|
(!state.disabled_by_server
|
||||||
|
&& !snapshot.tools.is_empty()
|
||||||
|
&& snapshot.published_at.elapsed() <= TOOL_CATALOG_CACHE_TTL)
|
||||||
|
.then_some(state.last_accepted_generation)
|
||||||
}
|
}
|
||||||
|
|
||||||
pub(crate) fn optional_startup_deadline(
|
pub(crate) fn optional_startup_deadline(
|
||||||
|
|||||||
@@ -1,5 +1,6 @@
|
|||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
use std::sync::atomic::AtomicBool;
|
use std::sync::atomic::AtomicBool;
|
||||||
|
use std::sync::atomic::AtomicUsize;
|
||||||
use std::sync::atomic::Ordering;
|
use std::sync::atomic::Ordering;
|
||||||
use std::time::Duration;
|
use std::time::Duration;
|
||||||
|
|
||||||
@@ -49,6 +50,13 @@ use pretty_assertions::assert_eq;
|
|||||||
use serde_json::Value;
|
use serde_json::Value;
|
||||||
use serde_json::json;
|
use serde_json::json;
|
||||||
use test_case::test_case;
|
use test_case::test_case;
|
||||||
|
use tracing::Subscriber;
|
||||||
|
use tracing::span::Attributes;
|
||||||
|
use tracing::span::Id;
|
||||||
|
use tracing_subscriber::Layer;
|
||||||
|
use tracing_subscriber::layer::Context as LayerContext;
|
||||||
|
use tracing_subscriber::layer::SubscriberExt;
|
||||||
|
use tracing_subscriber::util::SubscriberInitExt;
|
||||||
|
|
||||||
use super::rmcp_client::remote_aware_environment_id;
|
use super::rmcp_client::remote_aware_environment_id;
|
||||||
use super::rmcp_client::remote_aware_stdio_server_bin;
|
use super::rmcp_client::remote_aware_stdio_server_bin;
|
||||||
@@ -56,6 +64,18 @@ use super::rmcp_client::remote_aware_stdio_server_bin;
|
|||||||
const SERVER_NAME: &str = "cached_rmcp";
|
const SERVER_NAME: &str = "cached_rmcp";
|
||||||
const NAMESPACE: &str = "mcp__cached_rmcp";
|
const NAMESPACE: &str = "mcp__cached_rmcp";
|
||||||
|
|
||||||
|
struct McpBindingCaptureCounter(Arc<AtomicUsize>);
|
||||||
|
|
||||||
|
impl<S: Subscriber> Layer<S> for McpBindingCaptureCounter {
|
||||||
|
fn on_new_span(&self, attributes: &Attributes<'_>, _id: &Id, _context: LayerContext<'_, S>) {
|
||||||
|
if attributes.metadata().target() == "codex_mcp::connection_manager::tool_catalog"
|
||||||
|
&& attributes.metadata().name() == "capture_binding_with_metadata"
|
||||||
|
{
|
||||||
|
self.0.fetch_add(1, Ordering::SeqCst);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
fn user_turn(prompt: &str) -> TurnInputRequest {
|
fn user_turn(prompt: &str) -> TurnInputRequest {
|
||||||
TurnInputRequest::user_input(vec![UserInput::Text {
|
TurnInputRequest::user_input(vec![UserInput::Text {
|
||||||
text: prompt.to_string(),
|
text: prompt.to_string(),
|
||||||
@@ -558,7 +578,8 @@ async fn cached_http_mcp_starts_lazily_for_subagents(
|
|||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
// Keep spawned tasks on the thread with the scoped tracing subscriber.
|
||||||
|
#[tokio::test(flavor = "current_thread")]
|
||||||
async fn cached_mcp_startup_is_eager_for_root_and_lazy_for_subagents() -> anyhow::Result<()> {
|
async fn cached_mcp_startup_is_eager_for_root_and_lazy_for_subagents() -> anyhow::Result<()> {
|
||||||
skip_if_wine_exec!(
|
skip_if_wine_exec!(
|
||||||
Ok(()),
|
Ok(()),
|
||||||
@@ -566,6 +587,11 @@ async fn cached_mcp_startup_is_eager_for_root_and_lazy_for_subagents() -> anyhow
|
|||||||
);
|
);
|
||||||
skip_if_no_network!(Ok(()));
|
skip_if_no_network!(Ok(()));
|
||||||
|
|
||||||
|
let binding_captures = Arc::new(AtomicUsize::new(0));
|
||||||
|
let _tracing = tracing_subscriber::registry()
|
||||||
|
.with(McpBindingCaptureCounter(Arc::clone(&binding_captures)))
|
||||||
|
.set_default();
|
||||||
|
|
||||||
let responses_server = responses::start_mock_server().await;
|
let responses_server = responses::start_mock_server().await;
|
||||||
let command = remote_aware_stdio_server_bin()?;
|
let command = remote_aware_stdio_server_bin()?;
|
||||||
let environment_id = remote_aware_environment_id();
|
let environment_id = remote_aware_environment_id();
|
||||||
@@ -691,15 +717,29 @@ async fn cached_mcp_startup_is_eager_for_root_and_lazy_for_subagents() -> anyhow
|
|||||||
.await?;
|
.await?;
|
||||||
second_thread.submit(Op::Interrupt).await?;
|
second_thread.submit(Op::Interrupt).await?;
|
||||||
|
|
||||||
|
binding_captures.store(0, Ordering::SeqCst);
|
||||||
let unused_response = mount_sse_once(
|
let unused_response = mount_sse_once(
|
||||||
&responses_server,
|
&responses_server,
|
||||||
responses::sse(vec![
|
responses::sse(vec![
|
||||||
responses::ev_response_created("unused"),
|
responses::ev_response_created("unused"),
|
||||||
responses::ev_assistant_message("unused-message", "done"),
|
responses::ev_function_call(
|
||||||
|
"unused-plan",
|
||||||
|
"update_plan",
|
||||||
|
r#"{"plan":[{"step":"Work without MCP tools","status":"in_progress"}]}"#,
|
||||||
|
),
|
||||||
responses::ev_completed("unused"),
|
responses::ev_completed("unused"),
|
||||||
]),
|
]),
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
|
let unused_done_response = mount_sse_once(
|
||||||
|
&responses_server,
|
||||||
|
responses::sse(vec![
|
||||||
|
responses::ev_response_created("unused-done"),
|
||||||
|
responses::ev_assistant_message("unused-message", "done"),
|
||||||
|
responses::ev_completed("unused-done"),
|
||||||
|
]),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
second_thread
|
second_thread
|
||||||
.start_or_steer_turn(user_turn("Do not call any MCP tools."))
|
.start_or_steer_turn(user_turn("Do not call any MCP tools."))
|
||||||
.await?;
|
.await?;
|
||||||
@@ -718,11 +758,18 @@ async fn cached_mcp_startup_is_eager_for_root_and_lazy_for_subagents() -> anyhow
|
|||||||
!reported_ready_before_startup,
|
!reported_ready_before_startup,
|
||||||
"a dormant MCP server must not be reported as ready"
|
"a dormant MCP server must not be reported as ready"
|
||||||
);
|
);
|
||||||
|
for response in [&unused_response, &unused_done_response] {
|
||||||
assert_definition(
|
assert_definition(
|
||||||
&unused_response,
|
response,
|
||||||
&format!("Use the tools from {cached_process}."),
|
&format!("Use the tools from {cached_process}."),
|
||||||
&format!("Echo from {cached_process}."),
|
&format!("Echo from {cached_process}."),
|
||||||
);
|
);
|
||||||
|
}
|
||||||
|
assert_eq!(
|
||||||
|
binding_captures.load(Ordering::SeqCst),
|
||||||
|
1,
|
||||||
|
"both model steps must share one binding while the cached server stays dormant"
|
||||||
|
);
|
||||||
let (mcp_config, _) = second_thread.current_mcp_config_and_runtime_context().await;
|
let (mcp_config, _) = second_thread.current_mcp_config_and_runtime_context().await;
|
||||||
assert_eq!(
|
assert_eq!(
|
||||||
second_thread
|
second_thread
|
||||||
@@ -847,6 +894,10 @@ async fn cached_mcp_startup_is_eager_for_root_and_lazy_for_subagents() -> anyhow
|
|||||||
.await?;
|
.await?;
|
||||||
let expected_error = format!("MCP tool `{SERVER_NAME}/cwd` is not available to the model");
|
let expected_error = format!("MCP tool `{SERVER_NAME}/cwd` is not available to the model");
|
||||||
assert_eq!(cached_turn.await??, second_process);
|
assert_eq!(cached_turn.await??, second_process);
|
||||||
|
assert!(
|
||||||
|
binding_captures.load(Ordering::SeqCst) > 1,
|
||||||
|
"starting the cached server must invalidate the dormant binding"
|
||||||
|
);
|
||||||
assert_definition(
|
assert_definition(
|
||||||
&cached_done_response,
|
&cached_done_response,
|
||||||
&format!("Use the tools from {second_process}."),
|
&format!("Use the tools from {second_process}."),
|
||||||
|
|||||||
Reference in new issue
Block a user