openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-08-23 13:09:46 +00:00
Code Issues Packages Projects Releases Wiki Activity

Make MCP resource clients follow the latest runtime (#34733)

Browse Source 
## What changed

- Make `McpResourceClient` resolve resource operations and cache identity from
  the latest `McpRuntime` connection snapshot.
- Remove step-bound resource clients from `McpBinding` and the associated
  per-binding client identity tracking.

GitOrigin-RevId: ee59f5867308c5a63e6e232384a50cf1e0c2a011
This commit is contained in:
jif
2026-07-22 11:04:29 +00:00
committed by copyberry
parent 6278742c41
commit 84d2b203ed
4 changed files with 18 additions and 134 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
codex-rs/codex-mcp/src/binding.rs
View File

@@ -22,7 +22,6 @@ use tokio::sync::RwLock;
use crate::McpConfig;
use crate::binding_clients::McpBindingClients;
use crate::connection_manager::McpConnectionSet;
use crate::resource_client::McpResourceClient;
use crate::rmcp_client::ManagedClient;
use crate::server::McpServerMetadata;
use crate::tools::ToolInfo;
@@ -92,11 +91,6 @@ impl McpBinding {
self.connections.has_servers()
}
/// Returns resource access bound to this binding's exact connection set.
pub fn resource_client(&self) -> McpResourceClient {
McpResourceClient::for_binding(Arc::clone(&self.clients))
}
pub async fn list_resources(
&self,
server: &str,

37
codex-rs/codex-mcp/src/binding_clients.rs
View File

@@ -1,11 +1,9 @@
use std::collections::HashMap;
use std::sync::Arc;
use std::sync::Weak;
use anyhow::Context;
use anyhow::Result;
use anyhow::anyhow;
use codex_rmcp_client::RmcpClient;
use rmcp::model::ListResourceTemplatesResult;
use rmcp::model::ListResourcesResult;
use rmcp::model::PaginatedRequestParams;
@@ -21,50 +19,17 @@ use crate::rmcp_client::ManagedClient;
/// The ready clients captured for one model step.
pub(crate) struct McpBindingClients {
clients: HashMap<String, Arc<ManagedClient>>,
identity: McpBindingClientIdentity,
}
#[derive(Clone)]
pub(crate) struct McpBindingClientIdentity(Vec<(String, Weak<RmcpClient>)>);
impl PartialEq for McpBindingClientIdentity {
fn eq(&self, other: &Self) -> bool {
self.0.len() == other.0.len()
&& self.0.iter().zip(&other.0).all(
|((server, client), (other_server, other_client))| {
server == other_server && client.ptr_eq(other_client)
},
)
}
}
impl Eq for McpBindingClientIdentity {}
impl McpBindingClients {
pub(crate) fn new(clients: HashMap<String, Arc<ManagedClient>>) -> Self {
let mut identity = clients
.iter()
.map(|(server, client)| (server.clone(), Arc::downgrade(&client.client)))
.collect::<Vec<_>>();
identity.sort_by(|left, right| left.0.cmp(&right.0));
Self {
clients,
identity: McpBindingClientIdentity(identity),
}
Self { clients }
}
pub(crate) fn client(&self, server: &str) -> Option<Arc<ManagedClient>> {
self.clients.get(server).cloned()
}
pub(crate) fn contains_server(&self, server: &str) -> bool {
self.clients.contains_key(server)
}
pub(crate) fn identity(&self) -> McpBindingClientIdentity {
self.identity.clone()
}
pub(crate) async fn list_resources(
&self,
server: &str,

23
codex-rs/codex-mcp/src/connection_manager_tests.rs
View File

@@ -1323,11 +1323,6 @@ async fn capture_binding_waits_for_fresh_startup_even_with_cached_tools() {
.collect::<Vec<_>>(),
vec!["client_local_tool"]
);
assert!(
step.resource_client()
.has_server(CODEX_APPS_MCP_SERVER_NAME)
.await
);
}
#[tokio::test]
@@ -1877,12 +1872,6 @@ async fn tool_lists_do_not_block_and_share_codex_apps_startup_reconnect() {
pending_step.tools().is_empty(),
"a model step must not advertise cached tools without an exact ready client"
);
let pending_resources = pending_step.resource_client();
assert!(
!pending_resources
.has_server(CODEX_APPS_MCP_SERVER_NAME)
.await
);
release_reconnect.notify_one();
tokio::task::yield_now().await;
@@ -1908,18 +1897,6 @@ async fn tool_lists_do_not_block_and_share_codex_apps_startup_reconnect() {
.prepare_call(CODEX_APPS_MCP_SERVER_NAME, "drive_search")
.is_some()
);
let recovered_resources = recovered_step.resource_client();
assert!(
recovered_resources
.has_server(CODEX_APPS_MCP_SERVER_NAME)
.await
);
assert!(
!pending_resources
.has_server(CODEX_APPS_MCP_SERVER_NAME)
.await
);
assert!(pending_resources.cache_key() != recovered_resources.cache_key());
assert_eq!(attempts.load(std::sync::atomic::Ordering::SeqCst), 1);
}

86
codex-rs/codex-mcp/src/resource_client.rs
View File

@@ -9,8 +9,6 @@ use rmcp::model::PaginatedRequestParams;
use rmcp::model::ReadResourceRequestParams;
use crate::McpRuntime;
use crate::binding_clients::McpBindingClientIdentity;
use crate::binding_clients::McpBindingClients;
use crate::connection_manager::McpConnectionSet;
/// One page of resources returned by an MCP server.
@@ -29,48 +27,19 @@ pub struct McpResourceReadResult {
pub contents: Vec<ResourceContent>,
}
/// Access to MCP resources through either the latest runtime or one exact step.
/// Access to MCP resources through the latest runtime.
#[derive(Clone)]
pub struct McpResourceClient {
source: McpResourceSource,
runtime: Arc<McpRuntime>,
}
/// Opaque identity for the connection set currently used by an MCP resource client.
#[derive(Clone)]
pub struct McpResourceClientCacheKey(McpResourceClientCacheKeyInner);
#[derive(Clone)]
enum McpResourceClientCacheKeyInner {
Latest(Weak<McpConnectionSet>),
Exact(McpBindingClientIdentity),
}
#[derive(Clone)]
enum McpResourceSource {
Latest(Arc<McpRuntime>),
Exact(Arc<McpBindingClients>),
}
pub struct McpResourceClientCacheKey(Weak<McpConnectionSet>);
impl PartialEq for McpResourceClientCacheKey {
fn eq(&self, other: &Self) -> bool {
match (&self.0, &other.0) {
(
McpResourceClientCacheKeyInner::Latest(left),
McpResourceClientCacheKeyInner::Latest(right),
) => left.ptr_eq(right),
(
McpResourceClientCacheKeyInner::Exact(left),
McpResourceClientCacheKeyInner::Exact(right),
) => left == right,
(
McpResourceClientCacheKeyInner::Latest(_),
McpResourceClientCacheKeyInner::Exact(_),
)
| (
McpResourceClientCacheKeyInner::Exact(_),
McpResourceClientCacheKeyInner::Latest(_),
) => false,
}
self.0.ptr_eq(&other.0)
}
}
@@ -87,38 +56,19 @@ impl std::fmt::Debug for McpResourceClient {
impl McpResourceClient {
/// Creates a resource client that follows the thread's latest published runtime.
pub fn new(runtime: Arc<McpRuntime>) -> Self {
Self {
source: McpResourceSource::Latest(runtime),
}
}
pub(crate) fn for_binding(clients: Arc<McpBindingClients>) -> Self {
Self {
source: McpResourceSource::Exact(clients),
}
Self { runtime }
}
/// Returns the identity of the connection set used by this client.
pub fn cache_key(&self) -> McpResourceClientCacheKey {
let key = match &self.source {
McpResourceSource::Latest(runtime) => {
McpResourceClientCacheKeyInner::Latest(Arc::downgrade(&runtime.snapshot()))
}
McpResourceSource::Exact(clients) => {
McpResourceClientCacheKeyInner::Exact(clients.identity())
}
};
McpResourceClientCacheKey(key)
McpResourceClientCacheKey(Arc::downgrade(&self.runtime.snapshot()))
}
/// Returns whether this client can address the named server.
///
/// This does not wait for server startup.
pub async fn has_server(&self, server: &str) -> bool {
match &self.source {
McpResourceSource::Latest(runtime) => runtime.snapshot().contains_server(server),
McpResourceSource::Exact(clients) => clients.contains_server(server),
}
self.runtime.snapshot().contains_server(server)
}
/// Lists one resource page from the named server.
@@ -129,12 +79,11 @@ impl McpResourceClient {
) -> Result<McpResourcePage> {
let params =
cursor.map(|cursor| PaginatedRequestParams::default().with_cursor(Some(cursor)));
let result = match &self.source {
McpResourceSource::Latest(runtime) => {
runtime.snapshot().list_resources(server, params).await
}
McpResourceSource::Exact(clients) => clients.list_resources(server, params).await,
}?;
let result = self
.runtime
.snapshot()
.list_resources(server, params)
.await?;
let resources = result
.resources
.into_iter()
@@ -149,12 +98,11 @@ impl McpResourceClient {
/// Reads one resource from the named server.
pub async fn read_resource(&self, server: &str, uri: &str) -> Result<McpResourceReadResult> {
let params = ReadResourceRequestParams::new(uri.to_string());
let result = match &self.source {
McpResourceSource::Latest(runtime) => {
runtime.snapshot().read_resource(server, params).await
}
McpResourceSource::Exact(clients) => clients.read_resource(server, params).await,
}?;
let result = self
.runtime
.snapshot()
.read_resource(server, params)
.await?;
let contents = result
.contents
.into_iter()