openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-11-14 21:40:39 +00:00
Code Issues Packages Projects Releases Wiki Activity

Deduplicate overlapping Apps startup tool listings (#61220)

Browse Source 
## Why

Concurrent Apps startups with equivalent discovery inputs each issue a `tools/list` request even though they share a live tool catalog.

## What changed

Share successful tool listings across overlapping startups in the same discovery scope while keeping each client's MCP connection independent. Later startups and explicit refreshes still fetch from the server.

Include time spent waiting for a shared listing in the configured startup timeout. If a listing fails or is cancelled, let a waiting startup retry using its own connection. Preserve separate catalogs for different scopes and honor servers that opt out of catalog sharing.

## Testing

Add tests for shared startup listings with independent connections, waiting timeouts, cancellation recovery, and catalog isolation across discovery scopes.

GitOrigin-RevId: cb626f6cb9c38e27b3d75a8da55bda584aaef6f4
This commit is contained in:
Charlie Marsh
2026-11-13 01:08:55 +00:00
committed by copyberry
parent fdad70ba5d
commit e1340f2ea8
5 changed files with 306 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/codex-mcp/src/connection_manager_tests.rs
+24 -10
View File
@@ -433,6 +433,23 @@ impl InProcessTransportFactory for MutableToolsTransportFactory {
}
}
pub(crate) async fn create_test_apps_client(
tool: Tool,
list_started: Option<Arc<Notify>>,
release_list: Option<Arc<Notify>>,
list_requests: Arc<AtomicUsize>,
) -> anyhow::Result<Arc<RmcpClient>> {
let client = RmcpClient::new_in_process_client(Arc::new(RefreshTestTransportFactory {
tool,
list_started,
release_list,
next_cursor: None,
list_requests,
}))
.await?;
Ok(Arc::new(client))
}
struct DisconnectingToolsTransportFactory {
server: MutableToolsServer,
disconnect: CancellationToken,
@@ -856,16 +873,13 @@ pub(crate) async fn create_test_manager_with_ready_apps_client(
release_list: Option<Arc<Notify>>,
) -> anyhow::Result<Arc<McpConnectionSet>> {
let tool = create_test_tool(CODEX_APPS_MCP_SERVER_NAME, tool_name);
let client = Arc::new(
RmcpClient::new_in_process_client(Arc::new(RefreshTestTransportFactory {
tool: tool.tool.clone(),
list_started,
release_list,
next_cursor: None,
list_requests: Arc::new(AtomicUsize::new(0)),
}))
.await?,
);
let client = create_test_apps_client(
tool.tool.clone(),
list_started,
release_list,
Arc::default(),
)
.await?;
client
.initialize(
InitializeRequestParams::new(
codex-rs/codex-mcp/src/rmcp_client.rs
+22 -28
View File
@@ -14,6 +14,9 @@ mod managed_client;
#[path = "rmcp_client/status.rs"]
mod status;
#[cfg(test)]
#[path = "rmcp_client/apps_startup_tests.rs"]
mod apps_startup_tests;
#[cfg(test)]
#[path = "rmcp_client/provenance_tests.rs"]
mod provenance_tests;
@@ -69,7 +72,6 @@ use codex_config::McpServerTransportConfig;
use codex_config::types::AuthKeyringBackendKind;
use codex_config::types::OAuthCredentialsStoreMode;
use codex_connectors::ConnectorRuntimeContext;
use codex_connectors::ConnectorRuntimeFetchSource;
use codex_exec_server::Environment;
use codex_login::AuthChangeState;
use codex_otel::auth_storage::AuthStorageOriginator;
@@ -386,7 +388,7 @@ impl ManagedClientStartup {
sharing_declined,
publish_apps,
is_codex_apps_mcp_server,
startup_timeout: Some(startup_timeout),
startup_timeout,
tx_event,
elicitation_requests,
codex_apps_tools_cache_context,
@@ -1182,7 +1184,7 @@ async fn start_server_task(
let requested_capabilities = params.capabilities.clone();
let started_at = Instant::now();
let initialize_result = client
.initialize(params, startup_timeout, send_elicitation)
.initialize(params, Some(startup_timeout), send_elicitation)
.await;
record_protocol_discovery_metrics(
client.protocol_mode(),
@@ -1255,40 +1257,32 @@ async fn start_server_task(
let list_start = Instant::now();
let server_info =
mcp_server_info_from_implementation(&server_name, initialize_result.server_info);
let fetch_ticket = codex_apps_tools_cache_context
.as_ref()
.map(|context| context.begin_fetch(ConnectorRuntimeFetchSource::Startup));
if let Some(tx) = &publish_apps {
tx.send_replace(Some(Arc::downgrade(&client)));
}
let (tools, meta) = list_tools_for_client_uncached(
let list = list_tools_for_client_uncached(
&server_name,
is_codex_apps_mcp_server,
/*codex_apps_refresh_trigger*/ "initial",
&client,
startup_timeout,
Some(startup_timeout),
catalog_item_limit,
initialize_result.instructions.as_deref(),
)
.await
.map_err(StartupOutcomeError::from)?;
let client_tools: Arc<[ToolInfo]> =
match (codex_apps_tools_cache_context.as_ref(), fetch_ticket) {
(Some(context), Some(ticket)) if server_disables_tool_catalog_cache => {
context.publish_runtime_if_newest_accepted(
ticket,
&server_info,
tools.clone(),
meta.clone(),
);
tools.into()
);
let (client_tools, meta) = match &codex_apps_tools_cache_context {
Some(context) => tokio::select! {
// If the lock and deadline become ready together, skip the expired request.
biased;
() = tokio::time::sleep(startup_timeout) => {
return Err(anyhow!("tools/list startup timed out after {startup_timeout:?}").into());
}
(Some(context), Some(ticket)) => context
.publish_runtime_if_newest_accepted(ticket, &server_info, tools, meta.clone())
.shared_tools(),
(None, None) => tools.into(),
_ => unreachable!("Codex Apps fetch ticket requires cache context"),
};
result = context.fetch_startup_tools(&server_info, list) => result?,
},
None => {
let (tools, meta) = list.await?;
(tools.into(), meta)
}
};
if let (Some(cache_context), Some(mut fetch_ticket)) = (
tool_catalog_cache_context.as_ref(),
tool_catalog_fetch_ticket,
@@ -1420,7 +1414,7 @@ struct StartServerTaskParams {
sharing_declined: Option<CancellationToken>,
server_capabilities: Arc<StdMutex<Option<serde_json::Value>>>,
is_codex_apps_mcp_server: bool,
startup_timeout: Option<Duration>, // TODO: cancel_token should handle this.
startup_timeout: Duration,
tx_event: Option<Sender<Event>>,
elicitation_requests: ElicitationRequestManager,
codex_apps_tools_cache_context: Option<ConnectorRuntimeContext<ToolInfo>>,
codex-rs/codex-mcp/src/rmcp_client/apps_startup_tests.rs
+171
View File
@@ -0,0 +1,171 @@
//! Apps tool discovery across concurrent startups on independent MCP connections.
#![cfg(test)]
use super::*;
use crate::client_tool_catalog::ToolCatalogSnapshot;
use crate::connection_manager::tests::create_test_apps_client;
use crate::connection_manager::tests::create_test_tool;
use crate::mcp::tests::test_mcp_config;
use codex_connectors::ConnectorRuntimeContextKey;
use codex_connectors::ConnectorRuntimeManager;
use pretty_assertions::assert_eq;
use std::path::Path;
use std::sync::atomic::AtomicUsize;
use tokio::sync::Notify;
fn tool_names(catalog: ToolCatalogSnapshot) -> Vec<String> {
catalog
.tools
.iter()
.map(|tool| tool.tool.name.to_string())
.collect()
}
struct AppsStartup {
task: tokio::task::JoinHandle<Result<ManagedClient, StartupOutcomeError>>,
list_started: Arc<Notify>,
release_list: Arc<Notify>,
list_requests: Arc<AtomicUsize>,
}
async fn start_overlapping_apps_clients(
home: &Path,
startup_timeout: Duration,
) -> anyhow::Result<(AppsStartup, AppsStartup)> {
let cache = ConnectorRuntimeManager::new_without_cache().context(
home.to_path_buf(),
ConnectorRuntimeContextKey::personal(Some("account".into()), Some("user".into())),
);
let elicitation_requests = ElicitationRequestManager::new(
Arc::new(test_mcp_config(home.to_path_buf())),
/*reviewer*/ None,
/*lifecycle*/ None,
Default::default(),
);
let start = |tool_name: &str| {
let (publish, mut initialized) = watch::channel(None);
let params = StartServerTaskParams {
server_key: McpServerKey::unscoped(CODEX_APPS_MCP_SERVER_NAME),
publish_apps: Some(publish),
sharing_declined: None,
server_capabilities: Arc::default(),
is_codex_apps_mcp_server: true,
startup_timeout,
tx_event: None,
elicitation_requests: elicitation_requests.clone(),
codex_apps_tools_cache_context: Some(cache.clone()),
tool_catalog_cache_context: None,
tool_catalog_fetch_ticket: None,
client_elicitation_capability: ElicitationCapability::default(),
client_mcp_extensions: ClientMcpExtensions::default(),
auth_changes: None,
catalog_item_limit: crate::pagination::MAX_CODEX_APPS_TOOL_CATALOG_ITEMS,
};
let tool = create_test_tool(CODEX_APPS_MCP_SERVER_NAME, tool_name).tool;
async move {
let list_started = Arc::new(Notify::new());
let release_list = Arc::new(Notify::new());
let list_requests = Arc::new(AtomicUsize::new(0));
let client = create_test_apps_client(
tool,
Some(Arc::clone(&list_started)),
Some(Arc::clone(&release_list)),
Arc::clone(&list_requests),
)
.await?;
let mut startup = Box::pin(start_server_task(
CODEX_APPS_MCP_SERVER_NAME.into(),
client,
params,
));
tokio::select! {
biased;
_ = &mut startup => anyhow::bail!("startup finished before the listing was released"),
ready = initialized.changed() => ready?,
}
Ok::<_, anyhow::Error>(AppsStartup {
task: tokio::spawn(startup),
list_started,
release_list,
list_requests,
})
}
};
let first = start("lookup").await?;
tokio::time::timeout(startup_timeout, first.list_started.notified()).await?;
let second = start("peer_lookup").await?;
Ok((first, second))
}
#[tokio::test(start_paused = true)]
async fn overlapping_apps_startups_share_tools_and_keep_their_own_connections() -> anyhow::Result<()>
{
let home = tempfile::tempdir()?;
let (first, second) =
start_overlapping_apps_clients(home.path(), DEFAULT_STARTUP_TIMEOUT).await?;
first.release_list.notify_one();
second.release_list.notify_one();
let first_client = first.task.await??;
let second_client = second.task.await??;
assert_eq!(first.list_requests.load(Ordering::SeqCst), 1);
assert_eq!(second.list_requests.load(Ordering::SeqCst), 0);
for client in [&first_client, &second_client] {
assert_eq!(client.tool_catalog.read(tool_names).await, ["lookup"]);
}
first_client.client.shutdown().await;
let reply = second_client
.call_tool(
"lookup".into(),
/*arguments*/ None,
|_| Ok(None),
Some(DEFAULT_TOOL_TIMEOUT),
)
.await?;
assert!(!reply.is_error.unwrap_or(false));
second_client.client.shutdown().await;
Ok(())
}
#[tokio::test(start_paused = true)]
async fn waiting_for_apps_tools_uses_the_configured_startup_timeout() -> anyhow::Result<()> {
let home = tempfile::tempdir()?;
let (first, second) =
start_overlapping_apps_clients(home.path(), Duration::from_secs(5)).await?;
// The second server could reply immediately if its client issued another listing.
second.release_list.notify_one();
let (first_result, second_result) = tokio::try_join!(first.task, second.task)?;
for result in [first_result, second_result] {
assert!(
matches!(result, Err(StartupOutcomeError::Failed { error, .. }) if error.contains("timed out"))
);
}
assert_eq!(second.list_requests.load(Ordering::SeqCst), 0);
first.release_list.notify_one();
Ok(())
}
#[tokio::test(start_paused = true)]
async fn cancelling_an_apps_startup_releases_the_waiting_client() -> anyhow::Result<()> {
let home = tempfile::tempdir()?;
let (first, second) =
start_overlapping_apps_clients(home.path(), DEFAULT_STARTUP_TIMEOUT).await?;
let interval = DEFAULT_STARTUP_TIMEOUT / 3;
tokio::select! {
biased;
() = second.list_started.notified() => anyhow::bail!("peer listed before cancellation"),
() = tokio::time::sleep(interval) => {}
}
first.task.abort();
// The peer must start its own request before the abandoned listing would have timed out.
tokio::time::timeout(interval, second.list_started.notified()).await?;
first.release_list.notify_one();
second.release_list.notify_one();
let recovered = second.task.await??;
assert_eq!(
recovered.tool_catalog.read(tool_names).await,
["peer_lookup"]
);
recovered.client.shutdown().await;
Ok(())
}