Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions codex-rs/codex-mcp/src/connection_manager.rs
Original file line number Diff line number Diff line change
Expand Up @@ -539,6 +539,13 @@ impl McpConnectionSet {
self.servers.contains_key(server_name)
}

pub(crate) async fn wait_for_server_startup(&self, server_name: &str) -> bool {
let Some(view) = self.servers.get(server_name) else {
return false;
};
view.connection.client().await.is_ok()
}

/// Stop all MCP clients owned by this manager and terminate stdio server processes.
pub async fn shutdown(&self) {
let connections = self
Expand Down
14 changes: 14 additions & 0 deletions codex-rs/codex-mcp/src/connection_manager/tool_catalog.rs
Original file line number Diff line number Diff line change
Expand Up @@ -143,6 +143,19 @@ impl McpConnectionSet {
.startup_complete
.load(Ordering::Acquire)
{
if view.connection.client.has_cached_tools() {
if let Some(server_tools) =
view.listed_tools(&self.tool_plugin_provenance).await
{
listed_tools.extend(server_tools.into_iter().map(|mut tool| {
if let Some(annotations) = tool.tool.annotations.as_mut() {
annotations.read_only_hint = None;
}
Self::with_server_metadata(tool, &view.metadata)
}));
}
continue;
}
let _ = view.connection.client.client().await;
}
view.connection.client.reconnect_failed_startup().await;
Expand Down Expand Up @@ -186,6 +199,7 @@ impl McpConnectionSet {
continue;
}
let Some(client) = clients.client(&tool_info.server_name) else {
tools.push(tool_info);
continue;
};
let Some(call) = self.prepare_call(&tool_info, client, Arc::clone(&config), *revision)
Expand Down
50 changes: 39 additions & 11 deletions codex-rs/codex-mcp/src/connection_manager_tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1513,20 +1513,21 @@ async fn list_all_tools_accepts_canonical_namespaced_tool_names() {
}

#[tokio::test]
async fn capture_binding_waits_for_fresh_startup_even_with_cached_tools() {
async fn capture_binding_exposes_cached_tools_before_startup() {
let codex_home = tempdir().expect("tempdir");
let cache_context = create_codex_apps_tools_cache_context(
codex_home.path().to_path_buf(),
Some("account-one"),
Some("user-one"),
);
store_current_tools(
&cache_context,
vec![create_test_tool(
CODEX_APPS_MCP_SERVER_NAME,
"shared_cached_tool",
)],
let mut cached_tool = create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "shared_cached_tool");
cached_tool.tool.annotations = Some(
rmcp::model::ToolAnnotations::new()
.read_only(true)
.destructive(false)
.open_world(false),
);
store_current_tools(&cache_context, vec![cached_tool]);
let startup_complete = Arc::new(std::sync::atomic::AtomicBool::new(false));
let startup_complete_for_client = Arc::clone(&startup_complete);
let (startup_started, wait_for_startup) = tokio::sync::oneshot::channel();
Expand Down Expand Up @@ -1575,14 +1576,41 @@ async fn capture_binding_waits_for_fresh_startup_even_with_cached_tools() {
},
);
let manager = Arc::new(manager);
let manager_for_capture = Arc::clone(&manager);
let capture = tokio::spawn(async move { capture_binding(&manager_for_capture).await });
let cached_binding = capture_binding(&manager).await;
assert_eq!(
cached_binding
.tools()
.iter()
.map(|tool| tool.callable_name.as_str())
.collect::<Vec<_>>(),
vec!["shared_cached_tool"]
);
assert_eq!(
cached_binding.tools()[0].tool.annotations,
Some(
rmcp::model::ToolAnnotations::new()
.destructive(false)
.open_world(false)
)
);
assert!(
cached_binding
.prepare_call(CODEX_APPS_MCP_SERVER_NAME, "shared_cached_tool")
.is_none()
);

let manager_for_startup = Arc::clone(&manager);
let startup = tokio::spawn(async move {
manager_for_startup
.wait_for_server_startup(CODEX_APPS_MCP_SERVER_NAME)
.await
});

wait_for_startup.await.expect("client startup should begin");
assert!(!capture.is_finished());
release_startup.send(()).expect("release client startup");
assert!(startup.await.expect("startup task"));

let step = capture.await.expect("capture task");
let step = capture_binding(&manager).await;
assert_eq!(
step.tools()
.iter()
Expand Down
15 changes: 15 additions & 0 deletions codex-rs/codex-mcp/src/runtime.rs
Original file line number Diff line number Diff line change
Expand Up @@ -240,6 +240,21 @@ impl McpRuntime {
}
}

/// Captures the current runtime after its selected server has finished startup.
pub async fn current_binding_for_call(&self, server: &str) -> Option<Arc<McpBinding>> {
let current = self.current.load_full();
let config = Arc::clone(current.config.as_ref()?);
if !current.connections.wait_for_server_startup(server).await {
return None;
}
Some(Arc::new(
current
.connections
.capture_binding_with_metadata(config, current.plugins_available)
.await,
))
}

/// Returns the latest published configuration without waiting for clients.
pub fn current_config(&self) -> Option<Arc<McpConfig>> {
self.current.load().config.clone()
Expand Down
6 changes: 5 additions & 1 deletion codex-rs/core/src/mcp_tool_call.rs
Original file line number Diff line number Diff line change
Expand Up @@ -141,7 +141,11 @@ pub(crate) async fn handle_mcp_tool_call(
};

sess.refresh_mcp_if_dirty().await;
let current_binding = sess.services.mcp_runtime.current_binding().await;
let current_binding = sess
.services
.mcp_runtime
.current_binding_for_call(&server)
.await;
let Some(prepared_call) = current_binding
.as_ref()
.and_then(|binding| binding.prepare_call(&server, &tool_name))
Expand Down
17 changes: 9 additions & 8 deletions codex-rs/core/tests/suite/mcp_tool_cache.rs
Original file line number Diff line number Diff line change
Expand Up @@ -248,30 +248,31 @@ async fn regular_mcp_definition_cache_preserves_live_session_state() -> anyhow::
.await;
anyhow::Ok(called_process)
});
fixture.codex.shutdown_and_wait().await?;
fs.write_file(&barrier_file, b"ready".to_vec(), /*sandbox*/ None)
.await?;
tokio::time::timeout(Duration::from_secs(2), async {
while cached_response.requests().is_empty() {
tokio::time::sleep(Duration::from_millis(10)).await;
}
})
.await
.context("live MCP definitions should reach inference after initialization")?;
.context("cached MCP definitions should reach inference before initialization")?;
assert_definition(
&cached_response,
&format!("Use the tools from {second_process}."),
&format!("Echo from {second_process}."),
&format!("Tools in the {NAMESPACE} namespace."),
&format!("Echo from {first_process}."),
);

fixture.codex.shutdown_and_wait().await?;
fs.write_file(&barrier_file, b"ready".to_vec(), /*sandbox*/ None)
.await?;
let expected_error = format!("MCP tool `{SERVER_NAME}/cwd` is not available to the model");
assert_eq!(cached_turn.await??, second_process);
let output = cached_done_response
.single_request()
.function_call_output_text(app_only_call_id)
.expect("app-only tool error should be returned to the model");
assert!(
output.contains("is not available to the model") || output.contains("unsupported call"),
"app-only tools must be rejected before reaching the MCP server: {output}"
output.contains(&expected_error),
"model-visible tool output should contain the live visibility error: {output}"
);
let output = cached_done_response
.single_request()
Expand Down
Loading