Share eligible local stdio MCP connections across chats (#57211)

## What changed

Extend `mcp_connection_sharing` to pool compatible local stdio connections that negotiate MCP `2026-07-28`. Keep legacy servers and servers that request input during startup private.

Preserve each chat's tool catalog and route input requests to the calling chat. Release connections independently for each chat, shutting down the shared server after the final release, and retire pooled connections on forced refresh without interrupting other chats.

## Testing

Add pool, runtime, app-server, and core scenario tests covering protocol gating, startup cancellation, refresh, final-release shutdown, and separate tool visibility and simultaneous input responses across chats sharing one process.

GitOrigin-RevId: 896b107267a20ec1bffccdefca6c40fe2945a022
This commit is contained in:
Charlie Marsh
2026-10-26 19:37:32 +00:00
committed by copyberry
parent 0bc73bbeb1
commit 5e3496e0ad
24 changed files with 2197 additions and 142 deletions
@@ -0,0 +1,172 @@
//! Exercises protocol-gated stdio sharing through chat startup, tool dispatch, and shutdown.
use std::collections::BTreeMap;
use std::collections::HashMap;
use std::time::Duration;
use anyhow::Result;
use app_test_support::MockResponsesConfig;
use app_test_support::TestAppServer;
use app_test_support::create_mock_responses_server_sequence_unchecked;
use codex_app_server_protocol::ClientRequest;
use codex_app_server_protocol::ExperimentalFeatureEnablementSetParams;
use codex_app_server_protocol::ExperimentalFeatureEnablementSetResponse;
use codex_app_server_protocol::McpServerToolCallParams;
use codex_app_server_protocol::McpServerToolCallResponse;
use codex_app_server_protocol::ThreadClosedNotification;
use codex_app_server_protocol::ThreadStartParams;
use codex_app_server_protocol::ThreadUnsubscribeParams;
use codex_app_server_protocol::ThreadUnsubscribeResponse;
use codex_app_server_protocol::ThreadUnsubscribeStatus;
use core_test_support::skip_if_remote;
use core_test_support::stdio_server_bin;
use pretty_assertions::assert_eq;
use serde_json::json;
use tempfile::TempDir;
use tokio::time::timeout;
#[tokio::test]
async fn stdio_sharing_follows_protocol_and_enablement_for_loaded_chats() -> Result<()> {
skip_if_remote!(Ok(()), "the local stdio executable is a host-local fixture");
for (modern, modern_without_sharing) in [(false, false), (true, false), (true, true)] {
let responses = create_mock_responses_server_sequence_unchecked(Vec::new()).await;
let codex_home = TempDir::new()?;
let command = toml::Value::String(stdio_server_bin()?);
// Exercise both pool-driven protocol selection and modern connections without pooling.
let mut config = format!("[features]\nmcp_2026_07_28 = {modern_without_sharing}\n");
for name in ["alpha", "beta"] {
config.push_str(&format!(
"[mcp_servers.{name}]\ncommand = {command}\nenabled_tools = [\"echo\"]\n\
[mcp_servers.{name}.env]\nMCP_TEST_DYNAMIC_SERVER_METADATA = \"1\"\n"
));
if name == "beta" && modern {
config.push_str("CODEX_MCP_PROTOCOL_VERSION = \"2026-07-28\"\n");
} else if name == "beta" {
config.push_str("CODEX_MCP_PROTOCOL_VERSION = \"2025-06-18\"\n");
}
}
MockResponsesConfig::new(&responses.uri())
.with_root_config("thread_unload_delay_secs = 0")
.with_extra_config(&config)
.write(codex_home.path())?;
let mut app = TestAppServer::builder()
.with_codex_home(codex_home.path())
.build_initialized()
.await?;
set_sharing(&mut app, /*enabled*/ true).await?;
let first = app.start_thread(ThreadStartParams::default()).await?.thread;
let mut first_processes = call_servers(&mut app, &first.id).await?;
assert_ne!(first_processes["alpha"], first_processes["beta"]);
let second = app.start_thread(ThreadStartParams::default()).await?.thread;
let mut second_processes = call_servers(&mut app, &second.id).await?;
for name in ["alpha", "beta"] {
assert_eq!(
first_processes[name] == second_processes[name],
name == "alpha" || modern
);
}
if modern_without_sharing {
for sharing in [false, true] {
set_sharing(&mut app, sharing).await?;
first_processes = call_servers(&mut app, &first.id).await?;
second_processes = call_servers(&mut app, &second.id).await?;
for name in ["alpha", "beta"] {
assert_eq!(first_processes[name] == second_processes[name], sharing);
}
}
}
unload_thread(&mut app, &first.id).await?;
assert_eq!(call_servers(&mut app, &second.id).await?, second_processes);
unload_thread(&mut app, &second.id).await?;
#[cfg(unix)]
for process in first_processes.values().chain(second_processes.values()) {
core_test_support::process::wait_for_process_exit(
process
.strip_prefix("rmcp-test-process-")
.expect("the fixture identifies its PID"),
)
.await?;
}
#[cfg(not(unix))]
{
let third = app.start_thread(ThreadStartParams::default()).await?.thread;
let third_processes = call_servers(&mut app, &third.id).await?;
for name in ["alpha", "beta"] {
assert_ne!(third_processes[name], first_processes[name]);
}
}
}
Ok(())
}
async fn set_sharing(app: &mut TestAppServer, enabled: bool) -> Result<()> {
let enablement = BTreeMap::from([("mcp_connection_sharing".to_string(), enabled)]);
let response: ExperimentalFeatureEnablementSetResponse = app
.request(
|request_id| ClientRequest::ExperimentalFeatureEnablementSet {
request_id,
params: ExperimentalFeatureEnablementSetParams {
enablement: enablement.clone(),
},
},
)
.await?;
assert_eq!(response.enablement, enablement);
Ok(())
}
async fn call_servers(
app: &mut TestAppServer,
thread_id: &str,
) -> Result<HashMap<&'static str, String>> {
let mut processes = HashMap::new();
for name in ["alpha", "beta"] {
let response: McpServerToolCallResponse = app
.request(|request_id| ClientRequest::McpServerToolCall {
request_id,
params: McpServerToolCallParams {
thread_id: thread_id.to_string(),
server: name.to_string(),
tool: "echo".to_string(),
arguments: Some(json!({ "message": "ping" })),
meta: None,
},
})
.await?;
let result = response.structured_content.expect("echo result");
let process = result["echo"]
.as_str()
.expect("the fixture identifies its process");
processes.insert(name, process.to_owned());
}
Ok(processes)
}
async fn unload_thread(app: &mut TestAppServer, thread_id: &str) -> Result<()> {
let response: ThreadUnsubscribeResponse = app
.request(|request_id| ClientRequest::ThreadUnsubscribe {
request_id,
params: ThreadUnsubscribeParams {
thread_id: thread_id.to_string(),
},
})
.await?;
assert_eq!(
response,
ThreadUnsubscribeResponse {
status: ThreadUnsubscribeStatus::Unsubscribed,
}
);
let closed: ThreadClosedNotification = timeout(
Duration::from_secs(/*secs*/ 10),
app.read_notification("thread/closed"),
)
.await??;
assert_eq!(
closed,
ThreadClosedNotification {
thread_id: thread_id.to_string(),
}
);
Ok(())
}
@@ -87,6 +87,7 @@ mod luna_reserve;
mod marketplace_add;
mod marketplace_remove;
mod marketplace_upgrade;
mod mcp_connection_sharing;
mod mcp_event_stream;
mod mcp_protocol_default;
mod mcp_resource;
+56 -48
View File
@@ -398,55 +398,63 @@ impl PreparedMcpCall {
(server_timeout, requested_timeout) => server_timeout.or(requested_timeout),
};
let tool_name = self.tool_info().tool.name.to_string();
self.client
.tool_catalog
.run_with_snapshot(&self.catalog_snapshot, || async {
let (arguments, meta) = prepare().await?;
let timeout_deadline =
effective_timeout.map(|timeout| tokio::time::Instant::now() + timeout);
let add_trusted_access_context = self.connections.add_trusted_access_context(
self.tool_info(),
&self.server_metadata,
arguments.as_ref(),
meta,
);
let meta = match effective_timeout.zip(timeout_deadline) {
Some((timeout, deadline)) => {
tokio::time::timeout_at(deadline, add_trusted_access_context)
.await
.map_err(|_| {
codex_rmcp_client::request_timeout_error("tools/call", timeout)
})?
}
None => add_trusted_access_context.await,
};
let remaining_timeout = match effective_timeout.zip(timeout_deadline) {
Some((timeout, deadline)) => {
let remaining = deadline.saturating_duration_since(tokio::time::Instant::now());
if remaining.is_zero() {
return Err(codex_rmcp_client::request_timeout_error("tools/call", timeout));
let operation =
self.client
.tool_catalog
.run_with_snapshot(&self.catalog_snapshot, || async {
let (arguments, meta) = prepare().await?;
let timeout_deadline =
effective_timeout.map(|timeout| tokio::time::Instant::now() + timeout);
let add_trusted_access_context = self.connections.add_trusted_access_context(
self.tool_info(),
&self.server_metadata,
arguments.as_ref(),
meta,
);
let meta = match effective_timeout.zip(timeout_deadline) {
Some((timeout, deadline)) => {
tokio::time::timeout_at(deadline, add_trusted_access_context)
.await
.map_err(|_| {
codex_rmcp_client::request_timeout_error("tools/call", timeout)
})?
}
Some(remaining)
}
None => None,
};
self.client
.client
.call_tool(
tool_name.clone(),
arguments,
self.connections.network_policy_meta_for_call(
&self.client,
Some(self.permission_profile()),
self.server_environment_id(),
meta,
),
remaining_timeout,
)
.await
.with_context(|| format!("tool call failed for `{}/{tool_name}`", self.server_name))
})
.await
None => add_trusted_access_context.await,
};
let remaining_timeout = match effective_timeout.zip(timeout_deadline) {
Some((timeout, deadline)) => {
let remaining =
deadline.saturating_duration_since(tokio::time::Instant::now());
if remaining.is_zero() {
return Err(codex_rmcp_client::request_timeout_error(
"tools/call",
timeout,
));
}
Some(remaining)
}
None => None,
};
self.client
.call_tool(
tool_name.clone(),
arguments,
self.connections.network_policy_meta_for_call(
&self.client,
Some(self.permission_profile()),
self.server_environment_id(),
meta,
),
remaining_timeout,
)
.await
.with_context(|| {
format!("tool call failed for `{}/{tool_name}`", self.server_name)
})
});
self.client
.with_active_shared_lease(operation)
.await?
.ok_or_else(|| anyhow::anyhow!(
"tool call rejected because the catalog changed after `{}/{tool_name}` was prepared",
self.server_name
+2 -5
View File
@@ -40,7 +40,6 @@ impl McpBindingClients {
.client(server)
.ok_or_else(|| anyhow!("MCP server '{server}' was not ready for this step"))?;
managed
.client
.list_resources(params, managed.tool_timeout)
.await
.with_context(|| format!("resources/list failed for `{server}`"))
@@ -55,7 +54,6 @@ impl McpBindingClients {
.client(server)
.ok_or_else(|| anyhow!("MCP server '{server}' was not ready for this step"))?;
managed
.client
.list_resource_templates(params, managed.tool_timeout)
.await
.with_context(|| format!("resources/templates/list failed for `{server}`"))
@@ -71,7 +69,6 @@ impl McpBindingClients {
.ok_or_else(|| anyhow!("MCP server '{server}' was not ready for this step"))?;
let uri = params.uri.clone();
managed
.client
.read_resource(params, managed.tool_timeout)
.await
.with_context(|| format!("resources/read failed for `{server}` ({uri})"))
@@ -88,7 +85,7 @@ impl McpBindingClients {
.filter(|(server_name, _)| include_server(server_name))
{
let server_name = server_name.clone();
let client = Arc::clone(&managed.client);
let client = Arc::clone(managed);
let timeout = managed.tool_timeout;
join_set.spawn(async move {
let resources = collect_paginated("resources/list", timeout, |params| {
@@ -116,7 +113,7 @@ impl McpBindingClients {
.filter(|(server_name, _)| include_server(server_name))
{
let server_name = server_name.clone();
let client = Arc::clone(&managed.client);
let client = Arc::clone(managed);
let timeout = managed.tool_timeout;
join_set.spawn(async move {
let templates = collect_paginated("resources/templates/list", timeout, |params| {
+2
View File
@@ -84,6 +84,8 @@ async fn test_step(
/*updates*/ None,
));
let managed_client = Arc::new(ManagedClient {
connection_sharing: Default::default(),
shared_elicitation: None,
_auth_change_notifications: None,
client: Arc::clone(&client),
server_info: McpServerInfo {
@@ -69,6 +69,7 @@ enum CatalogSource {
},
}
#[derive(Clone)]
pub(crate) struct ToolCatalogSnapshot {
pub(crate) revision: u64,
pub(crate) tools: Arc<[ToolInfo]>,
+70 -15
View File
@@ -41,6 +41,8 @@ use std::time::Duration;
use crate::McpProtocolMode;
use crate::binding::call_tool_result_from_rmcp;
use crate::catalog::McpServerSource;
use crate::connection_pool::ConnectionKey;
use crate::connection_pool::ConnectionSharing;
use crate::elicitation::ElicitationRequestManager;
use crate::elicitation::ElicitationRequestRouter;
use crate::event_stream::EventStreamConnectionSettings;
@@ -279,8 +281,8 @@ impl McpServerView {
let timeout = self.refresh_timeout;
let result = client
.tool_catalog
.refresh_after_notifications(&client.client.tool_list_revision(), timeout, || {
crate::rmcp_client::list_tools_for_client_uncached(
.refresh_after_notifications(&client.client.tool_list_revision(), timeout, || async {
let relist = crate::rmcp_client::list_tools_for_client_uncached(
server_name,
/*is_codex_apps_mcp_server*/ false,
"list_changed",
@@ -288,7 +290,8 @@ impl McpServerView {
Some(timeout),
self.catalog_item_limit,
client.server_instructions.as_deref(),
)
);
client.with_active_shared_lease(relist.boxed()).await?
})
.await;
if result.is_err() {
@@ -675,6 +678,12 @@ impl McpConnectionSet {
expected_protocol_mode
};
let protocol_mode = expected_protocol_mode.unwrap_or(protocol_mode);
let connection_pool = runtime_context.connection_pool.as_ref().filter(|_| {
// Sharing a signed-in helper requires tracking account changes.
discover_local_stdio
&& protocol_mode != McpProtocolMode::Legacy
&& (auth.is_none() || auth_manager.is_some())
});
if matches!(configured_config.auth, McpServerAuth::EmaAuth) && existing_ema.is_some() {
if let Some(previous_view) = existing_ema
.filter(|_| elicitation_previous.is_some())
@@ -743,6 +752,10 @@ impl McpConnectionSet {
&& (same_owner || !previous_view.connection.may_use_legacy_fallback())
{
let connection = Arc::clone(&previous_view.connection);
// A failed pool acquisition can leave this chat on its private helper.
let retained_legacy = connection.retained_legacy().is_some();
let can_reuse = retained_legacy
|| connection.client.connection_slot.is_some() == connection_pool.is_some();
let reusable_pending_startup = connection.identity.as_ref()
== Some(&connection_identity)
&& !connection.client.startup_complete.load(Ordering::Acquire)
@@ -770,15 +783,15 @@ impl McpConnectionSet {
.reusable_client(&connection_identity)
.await
.filter(|_| previous_view.catalog_item_limit == catalog_item_limit);
if reusable_pending_startup
|| unchanged_auth_failure.is_some()
|| ready.as_ref().is_some_and(|client| {
expected_protocol_mode.is_some_and(|expected| {
client.client.protocol_mode() == expected
|| (previous_view.protocol_mode == expected
&& connection.retained_legacy().is_some())
})
})
if can_reuse
&& (reusable_pending_startup
|| unchanged_auth_failure.is_some()
|| ready.as_ref().is_some_and(|client| {
expected_protocol_mode.is_some_and(|expected| {
client.client.protocol_mode() == expected
|| (previous_view.protocol_mode == expected && retained_legacy)
})
}))
{
let pending_client =
reusable_pending_startup.then(|| connection.client.clone());
@@ -853,6 +866,7 @@ impl McpConnectionSet {
&& !connection.client.cancel_token.is_cancelled()
&& ready.is_some_and(|client| {
client.client.protocol_mode() == McpProtocolMode::Legacy
&& matches!(client.connection_sharing, ConnectionSharing::Private)
})
{
legacy_to_upgrade = Some(connection);
@@ -861,6 +875,24 @@ impl McpConnectionSet {
let cancel_token = startup_cancellation_token.child_token();
let tool_catalog_cache_context = tool_catalog_cache_context.take();
let has_runtime_auth = runtime_auth_provider.is_some();
let connection_pool = connection_pool.map(|pool| {
(
pool.clone(),
ConnectionKey {
server_name: server_name.clone(),
identity: connection_identity.clone(),
account: codex_apps_tools_cache_key.clone(),
auth_state: auth_manager.as_ref().map(|manager| {
(
Arc::clone(manager),
*manager.auth_change_state_receiver().borrow(),
)
}),
startup_timeout,
catalog_item_limit,
},
)
});
let mut async_managed_client = AsyncManagedClient::new(
ema_auth.clone(),
server_name.clone(),
@@ -890,6 +922,7 @@ impl McpConnectionSet {
.map(|manager| manager.auth_change_state_receiver()),
protocol_mode,
catalog_item_limit,
connection_pool,
);
// A forced reconnect must also wait for the previously published helper to exit.
if let Some(previous) = existing_ema.and_then(|set| set.servers.get(&server_name)) {
@@ -1192,14 +1225,29 @@ impl McpConnectionSet {
|| view.connection.client_for_call().await.is_ok()
}
/// Stop all MCP clients owned by this manager and terminate stdio server processes.
/// Release this manager's clients, terminating shared stdio servers after their final release.
pub async fn shutdown(&self) {
let connections = self
.servers
.values()
.map(|view| Arc::clone(&view.connection))
.collect::<Vec<_>>();
// Keep cleanup alive if an interrupt cancels the refresh that requested it.
Self::shutdown_connections(connections).await;
}
/// Releases pool leases retained by older bindings.
pub(crate) async fn shutdown_shared_connections(&self) {
let connections = self
.servers
.values()
.filter(|view| view.connection.client.connection_slot.is_some())
.map(|view| Arc::clone(&view.connection))
.collect();
Self::shutdown_connections(connections).await;
}
/// Keeps cleanup alive if an interrupt cancels the refresh that requested it.
async fn shutdown_connections(connections: Vec<Arc<McpServerConnection>>) {
let shutdown_task = tokio::spawn(async move {
join_all(connections.iter().map(|connection| connection.shutdown())).await;
});
@@ -1214,6 +1262,14 @@ impl McpConnectionSet {
}
}
pub(crate) fn retire_shared_connections(&self) {
for view in self.servers.values() {
if let Some(slot) = &view.connection.client.connection_slot {
slot.retire();
}
}
}
pub fn plugin_id_for_mcp_server_name(&self, server_name: &str) -> Option<&str> {
self.tool_plugin_context
.plugin_id_for_mcp_server_name(server_name)
@@ -1304,7 +1360,6 @@ impl McpConnectionSet {
meta.remove(ENTITLEMENT_CONTEXT_KEY);
}
let result: rmcp::model::CallToolResult = client
.client
.call_tool(
tool.to_string(),
arguments,
@@ -46,14 +46,13 @@ impl McpConnectionSet {
.filter(|(server_name, _)| include_server(server_name))
{
let server_name = server_name.clone();
let Ok(managed_client) = view.connection.client().await else {
let Ok(client) = view.connection.client().await else {
continue;
};
let timeout = view.tool_timeout;
let client = managed_client.client;
join_set.spawn(async move {
let resources = collect_paginated("resources/list", timeout, |params| {
let client = Arc::clone(&client);
let client = client.clone();
async move {
let response = client.list_resources(params, timeout).await?;
Ok((response.resources, response.next_cursor))
@@ -93,14 +92,13 @@ impl McpConnectionSet {
.filter(|(server_name, _)| include_server(server_name))
{
let server_name = server_name.clone();
let Ok(managed_client) = view.connection.client().await else {
let Ok(client) = view.connection.client().await else {
continue;
};
let timeout = view.tool_timeout;
let client = managed_client.client;
join_set.spawn(async move {
let templates = collect_paginated("resources/templates/list", timeout, |params| {
let client = Arc::clone(&client);
let client = client.clone();
async move {
let response = client.list_resource_templates(params, timeout).await?;
Ok((response.resource_templates, response.next_cursor))
@@ -137,7 +135,6 @@ impl McpConnectionSet {
) -> Result<ListResourcesResult> {
let (managed, timeout) = self.client_by_name(server).await?;
managed
.client
.list_resources(params, timeout)
.await
.with_context(|| format!("resources/list failed for `{server}`"))
@@ -150,7 +147,6 @@ impl McpConnectionSet {
) -> Result<ListResourceTemplatesResult> {
let (managed, timeout) = self.client_by_name(server).await?;
managed
.client
.list_resource_templates(params, timeout)
.await
.with_context(|| format!("resources/templates/list failed for `{server}`"))
@@ -164,7 +160,6 @@ impl McpConnectionSet {
let (managed, timeout) = self.client_by_name(server).await?;
let uri = params.uri.clone();
managed
.client
.read_resource(params, timeout)
.await
.with_context(|| format!("resources/read failed for `{server}` ({uri})"))
@@ -187,7 +187,7 @@ impl McpConnectionSet {
}
}
fn create_test_tool(server_name: &str, tool_name: &str) -> ToolInfo {
pub(crate) fn create_test_tool(server_name: &str, tool_name: &str) -> ToolInfo {
ToolInfo {
server_name: server_name.to_string(),
supports_parallel_tool_calls: false,
@@ -485,6 +485,8 @@ async fn legacy_tool_catalog_rejects_repeated_pagination_cursor() -> anyhow::Res
async fn create_test_managed_client(tools: Vec<ToolInfo>) -> ManagedClient {
ManagedClient {
connection_sharing: Default::default(),
shared_elicitation: None,
_auth_change_notifications: None,
client: Arc::new(
RmcpClient::new_in_process_client(Arc::new(TestInProcessTransportFactory))
@@ -597,6 +599,7 @@ pub(crate) async fn create_ready_async_managed_client(tools: Vec<ToolInfo>) -> A
codex_apps_tools_cache_context: None,
tool_catalog_cache_context: None,
startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(true)),
connection_slot: None,
startup_reconnect: None,
cancel_token: CancellationToken::new(),
}
@@ -770,6 +773,7 @@ fn create_gated_async_managed_client(
codex_apps_tools_cache_context: None,
tool_catalog_cache_context: None,
startup_complete,
connection_slot: None,
startup_reconnect: None,
cancel_token: CancellationToken::new(),
},
@@ -808,6 +812,8 @@ pub(crate) async fn create_test_manager_with_ready_apps_client(
.await?;
let managed_client = ManagedClient {
connection_sharing: Default::default(),
shared_elicitation: None,
_auth_change_notifications: None,
client,
server_info: create_test_server_info("Codex Apps"),
@@ -844,6 +850,7 @@ pub(crate) async fn create_test_manager_with_ready_apps_client(
codex_apps_tools_cache_context: Some(cache_context),
tool_catalog_cache_context: None,
startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(true)),
connection_slot: None,
startup_reconnect: None,
cancel_token: CancellationToken::new(),
},
@@ -899,6 +906,7 @@ fn create_test_manager_with_failed_apps_startup(
codex_apps_tools_cache_context: Some(cache_context),
tool_catalog_cache_context: None,
startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(true)),
connection_slot: None,
startup_reconnect: Some(Arc::new(CodexAppsStartupReconnect::new(reconnect_factory))),
cancel_token: CancellationToken::new(),
},
@@ -2351,6 +2359,7 @@ async fn list_all_tools_uses_shared_codex_apps_cache_while_client_is_pending() {
codex_apps_tools_cache_context: Some(cache_context),
tool_catalog_cache_context: None,
startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)),
connection_slot: None,
startup_reconnect: None,
cancel_token: CancellationToken::new(),
},
@@ -2424,6 +2433,7 @@ async fn capture_binding_preserves_ready_client_catalog_and_shares_prepared_meta
codex_apps_tools_cache_context: Some(cache_context),
tool_catalog_cache_context: None,
startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(true)),
connection_slot: None,
startup_reconnect: None,
cancel_token: CancellationToken::new(),
},
@@ -2821,6 +2831,7 @@ async fn list_available_server_infos_uses_cache_while_client_is_pending() {
codex_apps_tools_cache_context: None,
tool_catalog_cache_context: None,
startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)),
connection_slot: None,
startup_reconnect: None,
cancel_token: CancellationToken::new(),
},
@@ -2919,6 +2930,7 @@ async fn capture_binding_exposes_cached_tools_before_startup() {
codex_apps_tools_cache_context: Some(cache_context),
tool_catalog_cache_context: None,
startup_complete,
connection_slot: None,
startup_reconnect: None,
cancel_token: CancellationToken::new(),
},
@@ -3022,6 +3034,7 @@ async fn capture_binding_skips_pending_optional_servers_after_configured_shared_
codex_apps_tools_cache_context: None,
tool_catalog_cache_context: None,
startup_complete: Arc::new(AtomicBool::new(false)),
connection_slot: None,
startup_reconnect: None,
cancel_token: CancellationToken::new(),
},
@@ -3277,6 +3290,7 @@ async fn capture_binding_shares_optional_startup_grace_across_connection_sets()
codex_apps_tools_cache_context: None,
tool_catalog_cache_context: Some(cache_context.clone()),
startup_complete: Arc::new(AtomicBool::new(false)),
connection_slot: None,
startup_reconnect: None,
cancel_token: CancellationToken::new(),
},
@@ -3967,6 +3981,7 @@ async fn list_all_tools_blocks_while_client_is_pending_without_cached_tools() {
codex_apps_tools_cache_context: None,
tool_catalog_cache_context: None,
startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)),
connection_slot: None,
startup_reconnect: None,
cancel_token: CancellationToken::new(),
},
@@ -4026,6 +4041,7 @@ async fn shutdown_cancels_pending_tool_listing() {
codex_apps_tools_cache_context: None,
tool_catalog_cache_context: None,
startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)),
connection_slot: None,
startup_reconnect: None,
cancel_token,
},
@@ -4074,6 +4090,7 @@ async fn shutdown_continues_after_caller_is_aborted() {
codex_apps_tools_cache_context: None,
tool_catalog_cache_context: None,
startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)),
connection_slot: None,
startup_reconnect: None,
cancel_token: CancellationToken::new(),
},
@@ -4128,6 +4145,7 @@ async fn list_all_tools_does_not_block_when_shared_codex_apps_cache_is_empty() {
codex_apps_tools_cache_context: Some(cache_context),
tool_catalog_cache_context: None,
startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)),
connection_slot: None,
startup_reconnect: None,
cancel_token: CancellationToken::new(),
},
@@ -4184,6 +4202,7 @@ async fn list_all_tools_uses_shared_codex_apps_cache_when_client_startup_fails()
codex_apps_tools_cache_context: Some(cache_context),
tool_catalog_cache_context: None,
startup_complete,
connection_slot: None,
startup_reconnect: None,
cancel_token: CancellationToken::new(),
},
@@ -5357,7 +5376,7 @@ fn reusable_server_runtime_context() -> McpRuntimeContext {
)
}
fn reusable_server_identity(
pub(crate) fn reusable_server_identity(
server_name: &str,
config: &McpServerConfig,
runtime_context: &McpRuntimeContext,
@@ -5815,6 +5834,8 @@ async fn reconciliation_reuses_connection_without_relisting_regular_tools() -> a
)
.await?;
let managed_client = ManagedClient {
connection_sharing: Default::default(),
shared_elicitation: None,
_auth_change_notifications: None,
client: Arc::clone(&client),
server_info: create_test_server_info("Mutable tools"),
@@ -5857,6 +5878,7 @@ async fn reconciliation_reuses_connection_without_relisting_regular_tools() -> a
codex_apps_tools_cache_context: None,
tool_catalog_cache_context: None,
startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(true)),
connection_slot: None,
startup_reconnect: None,
cancel_token: CancellationToken::new(),
},
@@ -6693,6 +6715,7 @@ async fn reconciliation_replaces_closed_connections() -> anyhow::Result<()> {
codex_apps_tools_cache_context: None,
tool_catalog_cache_context: None,
startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(true)),
connection_slot: None,
startup_reconnect: None,
cancel_token: CancellationToken::new(),
},
+445
View File
@@ -0,0 +1,445 @@
//! Shares MCP 2026-07-28 connections while keeping each chat's ownership separate.
//!
//! The protocol requires request-scoped context and explicit handles for application state.
//! Request continuations use the calling chat's elicitation authority. Legacy connections and
//! servers that request input during startup remain private.
use std::sync::Arc;
use std::sync::Mutex;
use std::sync::OnceLock;
use std::sync::Weak;
use std::time::Duration;
use codex_async_utils::OrCancelExt;
use codex_connectors::ConnectorRuntimeContextKey;
use codex_login::AuthChangeState;
use codex_login::AuthManager;
use codex_rmcp_client::ElicitationResponse;
use codex_rmcp_client::SendElicitation;
use futures::FutureExt;
use rmcp::model::ElicitationAction;
use rmcp::model::ProtocolVersion;
use tokio_util::sync::CancellationToken;
use tokio_util::task::TaskTracker;
use crate::client_tool_catalog::ClientToolCatalog;
use crate::client_tool_catalog::ToolCatalogSnapshot;
use crate::rmcp_client::ManagedClient;
use crate::rmcp_client::ManagedClientFuture;
use crate::rmcp_client::StartupOutcomeError;
use crate::rmcp_client::list_tools_for_client_uncached;
use crate::server::McpServerConnectionIdentity;
use crate::tool_catalog_cache::McpToolCatalogCacheContext;
/// Process-owned registry for compatible local stdio servers using MCP 2026-07-28.
#[derive(Clone, Default)]
pub struct McpConnectionPool {
entries: Arc<Mutex<Vec<Weak<Entry>>>>,
}
#[derive(Clone)]
pub(crate) struct ConnectionKey {
pub(crate) server_name: String,
pub(crate) identity: McpServerConnectionIdentity,
pub(crate) account: ConnectorRuntimeContextKey,
pub(crate) auth_state: Option<(Arc<AuthManager>, AuthChangeState)>,
pub(crate) startup_timeout: Duration,
pub(crate) catalog_item_limit: usize,
}
impl PartialEq for ConnectionKey {
fn eq(&self, other: &Self) -> bool {
self.server_name == other.server_name
&& self.identity == other.identity
&& self.account == other.account
&& self.startup_timeout == other.startup_timeout
&& self.catalog_item_limit == other.catalog_item_limit
&& match (&self.auth_state, &other.auth_state) {
(Some((left, left_state)), Some((right, right_state))) => {
Arc::ptr_eq(left, right)
&& left_state.owner_generation == right_state.owner_generation
&& left_state.source_generation == right_state.source_generation
&& left_state.login_generation == right_state.login_generation
}
(None, None) => true,
(Some(_), None) | (None, Some(_)) => false,
}
}
}
struct Entry {
key: ConnectionKey,
client: ManagedClientFuture,
capabilities: Arc<Mutex<Option<serde_json::Value>>>,
// Share immutable definitions; every chat owns its catalog revisions and refresh locks.
catalog: Mutex<Option<ToolCatalogSnapshot>>,
cancellation: CancellationToken,
/// Wakes joiners without cancelling the creator's startup.
sharing_declined: CancellationToken,
owners: Weak<LeaseOwners>,
pool: Weak<Mutex<Vec<Weak<Entry>>>>,
}
/// One strong reference per active chat lease, independent of retained client bindings.
/// Every strong reference must be consumed by `ConnectionLease::release`.
struct LeaseOwners;
#[derive(Clone, Copy, PartialEq, Eq)]
enum AcquisitionRole {
Creator,
Joiner,
}
impl Entry {
fn remove_from_pool(self: &Arc<Self>) {
if let Some(pool) = self.pool.upgrade() {
pool.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.retain(|entry| !entry.ptr_eq(&Arc::downgrade(self)));
}
}
async fn tool_catalog_for_chat(
&self,
client: &ManagedClient,
role: AcquisitionRole,
cache_context: Option<&McpToolCatalogCacheContext>,
) -> Result<Arc<ClientToolCatalog>, StartupOutcomeError> {
// Listing must not update another chat's catalog or wait for its calls.
let listed = if role == AcquisitionRole::Creator {
client.tool_catalog.read(|snapshot| snapshot).await
} else {
let fetch_ticket = cache_context.map(|context| (context, context.begin_fetch()));
let (tools, meta) = list_tools_for_client_uncached(
&self.key.server_name,
self.key.server_name == crate::CODEX_APPS_MCP_SERVER_NAME,
"shared_connection",
&client.client,
Some(self.key.startup_timeout),
self.key.catalog_item_limit,
client.server_instructions.as_deref(),
)
.await?;
if let Some((context, mut ticket)) = fetch_ticket {
if context.track_tool_list_changes {
ticket.tool_list_revision = Some(client.client.tool_list_revision());
}
context.publish_if_newest(ticket, &tools, meta.clone());
}
ToolCatalogSnapshot {
revision: 0,
tools: tools.into(),
meta,
}
};
let mut catalog = self
.catalog
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let catalog = catalog.get_or_insert_with(|| listed.clone());
if catalog.tools != listed.tools || catalog.meta != listed.meta {
*catalog = listed;
}
Ok(Arc::new(ClientToolCatalog::new(
Arc::clone(&catalog.tools),
catalog.meta.clone(),
/*updates*/ None,
)))
}
}
/// A chat releases its lease explicitly on shutdown, even while bindings retain its client.
pub(crate) struct ConnectionLease {
entry: Arc<Entry>,
owner: Mutex<Option<Arc<LeaseOwners>>>,
cancellation: CancellationToken,
slot: Arc<ConnectionSlot>,
}
/// Tracks an acquisition without retaining its lease, including before startup completes.
#[derive(Default)]
pub(crate) struct ConnectionSlot {
state: Mutex<SlotState>,
shutdown_tasks: OnceLock<TaskTracker>,
}
#[derive(Default)]
enum SlotState {
#[default]
Available,
Acquired(Weak<Entry>),
Retired,
}
impl ConnectionSlot {
pub(crate) fn retire(&self) {
let mut state = self
.state
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if let SlotState::Acquired(entry) = std::mem::replace(&mut *state, SlotState::Retired)
&& let Some(entry) = entry.upgrade()
{
entry.remove_from_pool();
}
}
pub(crate) async fn wait_for_shutdown(&self) {
if let Some(tasks) = self.shutdown_tasks.get() {
tasks.close();
tasks.wait().await;
}
}
}
#[derive(Clone, Default)]
pub(crate) enum ConnectionSharing {
#[default]
Private,
Negotiated,
/// The server declined sharing, but its creator still owns the startup entry.
Owned(Arc<ConnectionLease>),
Shared(Arc<ConnectionLease>),
}
impl McpConnectionPool {
pub(crate) async fn connect(
&self,
slot: Arc<ConnectionSlot>,
key: ConnectionKey,
capabilities: Arc<Mutex<Option<serde_json::Value>>>,
tool_catalog_cache_context: Option<McpToolCatalogCacheContext>,
cancellation: CancellationToken,
start: impl Fn(CancellationToken, Option<CancellationToken>) -> ManagedClientFuture + Sync,
) -> Result<ManagedClient, StartupOutcomeError> {
let (lease, role) = self
.acquire(key, &slot, Arc::clone(&capabilities), &cancellation, &start)
.or_cancel(&cancellation)
.await
.unwrap_or(Err(StartupOutcomeError::Cancelled))?;
let outcome = tokio::select! {
outcome = lease.entry.client.clone() => Some(outcome),
() = lease.entry.sharing_declined.cancelled(), if role == AcquisitionRole::Joiner => None,
() = cancellation.cancelled() => return Err(StartupOutcomeError::Cancelled),
};
match outcome {
Some(Ok(mut client))
if role == AcquisitionRole::Creator
|| matches!(client.connection_sharing, ConnectionSharing::Negotiated) =>
{
let shared_capabilities = lease
.entry
.capabilities
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.clone();
*capabilities
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner) = shared_capabilities;
client.connection_sharing =
if matches!(client.connection_sharing, ConnectionSharing::Negotiated) {
client.client.enable_request_cancellation();
client.tool_catalog = lease
.entry
.tool_catalog_for_chat(
&client,
role,
tool_catalog_cache_context.as_ref(),
)
.or_cancel(&cancellation)
.await
.unwrap_or(Err(StartupOutcomeError::Cancelled))?;
ConnectionSharing::Shared(lease)
} else {
lease.entry.remove_from_pool();
ConnectionSharing::Owned(lease)
};
Ok(client)
}
Some(Err(error)) if role == AcquisitionRole::Creator => Err(error),
_ => {
// Older or failed servers leave other chats to initialize privately
// with their own elicitation authority.
drop(lease);
start(cancellation, /*sharing_declined*/ None).await
}
}
}
async fn acquire(
&self,
key: ConnectionKey,
slot: &Arc<ConnectionSlot>,
capabilities: Arc<Mutex<Option<serde_json::Value>>>,
cancellation: &CancellationToken,
start: impl Fn(CancellationToken, Option<CancellationToken>) -> ManagedClientFuture,
) -> Result<(Arc<ConnectionLease>, AcquisitionRole), StartupOutcomeError> {
loop {
let (lease, role) = {
let mut state = slot
.state
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if matches!(&*state, SlotState::Retired) {
return Err(StartupOutcomeError::Cancelled);
}
let mut entries = self
.entries
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
entries.retain(|entry| entry.strong_count() > 0);
let lease_cancellation = cancellation.child_token();
let existing = entries
.iter()
.filter_map(Weak::upgrade)
.filter(|entry| {
entry.key == key
&& !entry.sharing_declined.is_cancelled()
&& entry.client.peek().is_none_or(|result| {
result.as_ref().is_ok_and(|client| {
matches!(
client.connection_sharing,
ConnectionSharing::Negotiated
)
})
})
})
.find_map(|entry| entry.owners.upgrade().map(|owner| (entry, owner)));
let (entry, owner, role) = if let Some((entry, owner)) = existing {
(entry, owner, AcquisitionRole::Joiner)
} else {
let startup_cancellation = CancellationToken::new();
let sharing_declined = CancellationToken::new();
let client =
start(startup_cancellation.clone(), Some(sharing_declined.clone()));
// Startup belongs to the connection, including if its first chat is cancelled.
tokio::spawn(client.clone());
let owner = Arc::new(LeaseOwners);
let entry = Arc::new(Entry {
key: key.clone(),
client,
capabilities: Arc::clone(&capabilities),
catalog: Mutex::new(None),
cancellation: startup_cancellation,
sharing_declined,
owners: Arc::downgrade(&owner),
pool: Arc::downgrade(&self.entries),
});
entries.push(Arc::downgrade(&entry));
(entry, owner, AcquisitionRole::Creator)
};
let lease = Arc::new(ConnectionLease {
entry,
owner: Mutex::new(Some(owner)),
cancellation: lease_cancellation,
slot: Arc::clone(slot),
});
*state = SlotState::Acquired(Arc::downgrade(&lease.entry));
(lease, role)
};
if let Some(Ok(client)) = lease.entry.client.peek()
&& client.client.is_closed().await
{
lease.entry.remove_from_pool();
lease.release();
continue;
}
return Ok((lease, role));
}
}
}
impl ConnectionLease {
pub(crate) fn cancellation(&self) -> &CancellationToken {
&self.cancellation
}
/// Returns the slot tracking shutdown when this chat releases the last lease.
pub(crate) fn release(&self) -> Option<&ConnectionSlot> {
self.cancellation.cancel();
let mut owner = self
.owner
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
Arc::into_inner(owner.take()?)?;
self.entry.cancellation.cancel();
self.entry.remove_from_pool();
let client = self.entry.client.clone();
let runtime = tokio::runtime::Handle::try_current().ok()?;
// The chat must be able to await its final release even when cancellation drops the lease.
let tasks = self.slot.shutdown_tasks.get_or_init(TaskTracker::new);
drop(tasks.spawn_on(
async move {
match client.await {
Ok(client) => client.client.shutdown().await,
Err(error) => error.wait_for_shutdown().await,
}
},
&runtime,
));
Some(&self.slot)
}
}
impl Drop for ConnectionLease {
fn drop(&mut self) {
let _ = self.release();
}
}
/// Releases waiting chats when startup becomes private and detaches callbacks after sharing.
pub(crate) struct SharingNegotiation {
callback: Arc<Mutex<Option<SendElicitation>>>,
sharing_declined: CancellationToken,
}
impl SharingNegotiation {
pub(crate) fn new(
sender: SendElicitation,
sharing_declined: CancellationToken,
) -> (Self, SendElicitation) {
let callback = Arc::new(Mutex::new(Some(sender)));
let negotiation = Self {
callback: Arc::clone(&callback),
sharing_declined: sharing_declined.clone(),
};
let sender: SendElicitation = Box::new(move |id, request| {
let callback = callback
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if let Some(callback) = callback.as_ref() {
// Client input binds this startup to its original chat.
sharing_declined.cancel();
callback(id, request)
} else {
futures::future::ready(Ok(ElicitationResponse {
action: ElicitationAction::Cancel,
content: None,
meta: None,
}))
.boxed()
}
});
(negotiation, sender)
}
pub(crate) fn finish(self, protocol_version: &ProtocolVersion) -> ConnectionSharing {
let mut callback = self
.callback
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if *protocol_version == ProtocolVersion::V_2026_07_28
&& !self.sharing_declined.is_cancelled()
{
callback.take();
ConnectionSharing::Negotiated
} else {
self.sharing_declined.cancel();
ConnectionSharing::Private
}
}
}
#[cfg(test)]
#[path = "connection_pool_tests.rs"]
pub(crate) mod tests;
@@ -0,0 +1,657 @@
//! Tests for shared MCP connection ownership and protocol negotiation.
use super::*;
use crate::binding_clients::McpBindingClients;
use crate::connection_manager::tests::create_ready_async_managed_client;
use crate::connection_manager::tests::create_test_tool;
use crate::connection_manager::tests::reusable_server_identity;
use crate::runtime::McpRuntimeContext;
use codex_config::McpServerConfig;
use codex_exec_server::EnvironmentManager;
use codex_login::CodexAuth;
use codex_rmcp_client::Elicitation;
use codex_rmcp_client::InProcessTransportFactory;
use codex_rmcp_client::RmcpClient;
use futures::future::BoxFuture;
use pretty_assertions::assert_eq;
use rmcp::ServerHandler;
use rmcp::ServiceExt;
use rmcp::model::ClientCapabilities;
use rmcp::model::Implementation;
use rmcp::model::InitializeRequestParams;
use rmcp::model::ListResourceTemplatesResult;
use rmcp::model::ListResourcesResult;
use rmcp::model::ListToolsResult;
use rmcp::model::PaginatedRequestParams;
use rmcp::model::ProtocolVersion;
use rmcp::model::ServerCapabilities;
use rmcp::model::ServerConfig;
use rmcp::model::Tool;
use rmcp::service::RequestContext;
use rmcp::service::RoleServer;
use serde_json::json;
use std::collections::HashMap;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering;
use tokio::sync::Notify;
use tokio::sync::Semaphore;
#[derive(Clone)]
pub(crate) struct TestServer {
tools: Arc<Mutex<Vec<Tool>>>,
block_listing: Arc<AtomicBool>,
started_listings: Arc<Semaphore>,
cancelled_listings: Arc<Semaphore>,
}
impl Default for TestServer {
fn default() -> Self {
Self {
tools: Arc::new(Mutex::new(vec![
create_test_tool("test_server", "test").tool,
])),
block_listing: Arc::default(),
started_listings: Arc::new(Semaphore::new(/*permits*/ 0)),
cancelled_listings: Arc::new(Semaphore::new(/*permits*/ 0)),
}
}
}
impl ServerHandler for TestServer {
fn get_info(&self) -> ServerConfig {
ServerConfig::new(
ServerCapabilities::builder()
.enable_tools()
.enable_resources()
.build(),
)
}
async fn call_tool(
&self,
_request: rmcp::model::CallToolRequestParams,
_context: RequestContext<RoleServer>,
) -> Result<rmcp::model::CallToolResponse, rmcp::ErrorData> {
Ok(rmcp::model::CallToolResult::success(Vec::new()).into())
}
async fn list_tools(
&self,
_request: Option<PaginatedRequestParams>,
context: RequestContext<RoleServer>,
) -> Result<ListToolsResult, rmcp::ErrorData> {
if self.block_listing.load(Ordering::Acquire) {
context.ct.cancelled().await;
self.cancelled_listings.add_permits(1);
return Err(rmcp::ErrorData::internal_error(
"listing cancelled",
/*data*/ None,
));
}
Ok(ListToolsResult::with_all_items(
self.tools.lock().unwrap().clone(),
))
}
async fn list_resources(
&self,
_request: Option<PaginatedRequestParams>,
context: RequestContext<RoleServer>,
) -> Result<ListResourcesResult, rmcp::ErrorData> {
if self.block_listing.load(Ordering::Acquire) {
self.started_listings.add_permits(1);
context.ct.cancelled().await;
self.cancelled_listings.add_permits(1);
}
Ok(ListResourcesResult::with_all_items(Vec::new()))
}
async fn list_resource_templates(
&self,
_request: Option<PaginatedRequestParams>,
context: RequestContext<RoleServer>,
) -> Result<ListResourceTemplatesResult, rmcp::ErrorData> {
if self.block_listing.load(Ordering::Acquire) {
self.started_listings.add_permits(1);
context.ct.cancelled().await;
self.cancelled_listings.add_permits(1);
}
Ok(ListResourceTemplatesResult::with_all_items(Vec::new()))
}
}
impl InProcessTransportFactory for TestServer {
fn open(&self) -> BoxFuture<'static, std::io::Result<tokio::io::DuplexStream>> {
let server = self.clone();
async move {
let (client, transport) = tokio::io::duplex(4096);
tokio::spawn(async move {
server
.serve(transport)
.await
.unwrap()
.waiting()
.await
.unwrap();
});
Ok(client)
}
.boxed()
}
}
pub(crate) fn key(context: &McpRuntimeContext, config: &McpServerConfig) -> ConnectionKey {
ConnectionKey {
server_name: "test_server".to_string(),
identity: reusable_server_identity("test_server", config, context),
account: ConnectorRuntimeContextKey::personal(
/*account_id*/ None, /*chatgpt_user_id*/ None,
),
auth_state: None,
startup_timeout: Duration::from_secs(30),
catalog_item_limit: 100,
}
}
pub(crate) fn context() -> McpRuntimeContext {
McpRuntimeContext::new(
Arc::new(EnvironmentManager::default_for_tests()),
std::env::temp_dir(),
)
}
pub(crate) fn config() -> McpServerConfig {
serde_json::from_value(json!({"command": "test-server", "args": ["server.mjs"]})).unwrap()
}
pub(crate) fn startup(
sharing: ConnectionSharing,
starts: Arc<AtomicUsize>,
first_started: Arc<Notify>,
release_first: Arc<Notify>,
server: TestServer,
) -> impl Fn(CancellationToken, Option<CancellationToken>) -> ManagedClientFuture {
move |cancellation, _sharing_declined| {
let sharing = sharing.clone();
let starts = Arc::clone(&starts);
let first_started = Arc::clone(&first_started);
let release_first = Arc::clone(&release_first);
let server = server.clone();
async move {
async {
if starts.fetch_add(1, Ordering::SeqCst) == 0 {
first_started.notify_one();
release_first.notified().await;
}
let mut client = create_ready_async_managed_client(Vec::new())
.await
.client()
.await?;
client.client = Arc::new(
RmcpClient::new_in_process_client(Arc::new(server))
.await
.expect("create test MCP client"),
);
client
.client
.initialize(
InitializeRequestParams::new(
ClientCapabilities::default(),
Implementation::new("test", "1"),
)
.with_protocol_version(ProtocolVersion::V_2025_06_18),
Some(Duration::from_secs(30)),
Box::new(|_, _| {
futures::future::ready(Ok(ElicitationResponse {
action: ElicitationAction::Decline,
content: None,
meta: None,
}))
.boxed()
}),
)
.await?;
let (tools, meta) = list_tools_for_client_uncached(
"test_server",
/*is_codex_apps_mcp_server*/ false,
"test",
&client.client,
/*timeout*/ None,
/*catalog_item_limit*/ 100,
/*server_instructions*/ None,
)
.await?;
client.tool_catalog =
Arc::new(ClientToolCatalog::new(tools, meta, /*updates*/ None));
client.connection_sharing = sharing;
Ok(client)
}
.or_cancel(&cancellation)
.await
.unwrap_or(Err(StartupOutcomeError::Cancelled))
}
.boxed()
.shared()
}
}
#[tokio::test]
async fn shared_startup_survives_first_chat_cancellation_and_rejects_released_bindings() {
let pool = McpConnectionPool::default();
let key = key(&context(), &config());
let server = TestServer::default();
let starts = Arc::new(AtomicUsize::new(0));
let started = Arc::new(Notify::new());
let release = Arc::new(Notify::new());
let launch = startup(
ConnectionSharing::Negotiated,
Arc::clone(&starts),
Arc::clone(&started),
Arc::clone(&release),
server.clone(),
);
let first_cancel = CancellationToken::new();
let first = pool.connect(
Arc::default(),
key.clone(),
Arc::default(),
/*tool_catalog_cache_context*/ None,
first_cancel.clone(),
&launch,
);
let second = pool.connect(
Arc::default(),
key.clone(),
Arc::default(),
/*tool_catalog_cache_context*/ None,
CancellationToken::new(),
&launch,
);
tokio::pin!(first, second);
assert!(futures::poll!(&mut first).is_pending());
assert!(futures::poll!(&mut second).is_pending());
started.notified().await;
first_cancel.cancel();
assert!(matches!(first.await, Err(StartupOutcomeError::Cancelled)));
release.notify_one();
let client = second.await.unwrap();
assert_eq!(starts.load(Ordering::SeqCst), 1);
assert!(!client.client.is_closed().await);
let other = pool
.connect(
Arc::default(),
key,
Arc::default(),
/*tool_catalog_cache_context*/ None,
CancellationToken::new(),
&launch,
)
.await
.unwrap();
let mut config = crate::mcp::tests::test_mcp_config(std::env::temp_dir());
config
.server_permission_profiles
.insert("test_server".to_string(), config.permission_profile.clone());
let prepared = crate::binding::PreparedMcpCall::new(
Arc::new(crate::connection_manager::McpConnectionSet::empty(
/*prefix_mcp_tool_names*/ true,
)),
Arc::new(client.clone()),
Arc::new(config),
client.tool_catalog.read(Arc::new).await,
create_test_tool("test_server", "test"),
crate::server::McpServerMetadata {
environment_id: "local".to_string(),
pollutes_memory: false,
origin: None,
supports_parallel_tool_calls: true,
default_tools_approval_mode: None,
tool_approval_modes: Default::default(),
},
/*plugin_id*/ None,
/*selected_plugin_server*/ false,
)
.unwrap();
let bindings = McpBindingClients::new(HashMap::from([(
"test_server".to_string(),
Arc::new(client.clone()),
)]));
server.block_listing.store(true, Ordering::Release);
let waiting = async {
tokio::join!(
bindings.list_resources("test_server", /*params*/ None),
bindings.list_resource_templates("test_server", /*params*/ None),
)
};
tokio::pin!(waiting);
tokio::select! {
_ = &mut waiting => panic!("resource listings must wait for cancellation"),
permits = server.started_listings.acquire_many(2) => permits.unwrap().forget(),
}
let mut owned = create_ready_async_managed_client(Vec::new()).await;
owned.client = futures::future::ready(Ok(client.clone())).boxed().shared();
owned.shutdown().await;
let (resources, templates) = waiting.await;
assert!(resources.is_err());
assert!(templates.is_err());
tokio::time::timeout(
Duration::from_secs(1),
server.cancelled_listings.acquire_many(2),
)
.await
.expect("released resource listings must be cancelled")
.unwrap()
.forget();
let prepared_error = prepared
.call_with_preparation(/*requested_timeout*/ None, || async {
panic!("released bindings must not perform irreversible preparation")
})
.await
.unwrap_err();
assert!(
prepared_error
.to_string()
.contains("lease has been released")
);
other
.call_tool(
"test".to_string(),
/*arguments*/ None,
|_| Ok(None),
/*timeout*/ None,
)
.await
.unwrap();
owned.client = futures::future::ready(Ok(other)).boxed().shared();
owned.shutdown().await;
}
#[tokio::test]
async fn callback_during_startup_releases_waiting_chats_without_sharing_authority() {
let pool = McpConnectionPool::default();
let key = key(&context(), &config());
let starts = Arc::new(AtomicUsize::new(0));
let started = Arc::new(Notify::new());
let release = Arc::new(Notify::new());
let second_callback = Arc::new(Notify::new());
let release_second = Arc::new(Notify::new());
let callbacks = Arc::new(Mutex::new(Vec::new()));
let launch = |name| {
let launch = startup(
ConnectionSharing::Private,
Arc::clone(&starts),
Arc::clone(&started),
Arc::clone(&release),
TestServer::default(),
);
let callbacks = Arc::clone(&callbacks);
let second_callback = Arc::clone(&second_callback);
let release_second = Arc::clone(&release_second);
move |cancellation, sharing_declined: Option<CancellationToken>| {
let startup = launch(cancellation, sharing_declined.clone());
let callbacks = Arc::clone(&callbacks);
let second_callback = Arc::clone(&second_callback);
let release_second = Arc::clone(&release_second);
async move {
let mut client = startup.await?;
let (negotiation, callback) = SharingNegotiation::new(
Box::new(move |_, _| {
callbacks.lock().unwrap().push(name);
futures::future::ready(Ok(ElicitationResponse {
action: ElicitationAction::Decline,
content: None,
meta: None,
}))
.boxed()
}),
sharing_declined.unwrap_or_default(),
);
callback(
rmcp::model::RequestId::Number(0),
Elicitation::OpenAiForm {
meta: None,
message: "startup prompt".to_string(),
requested_schema: json!({}),
},
)
.await
.unwrap();
if name == "first" {
second_callback.notified().await;
} else {
second_callback.notify_one();
release_second.notified().await;
}
client.connection_sharing = negotiation.finish(&ProtocolVersion::V_2026_07_28);
Ok(client)
}
.boxed()
.shared()
}
};
let first = pool.connect(
Arc::default(),
key.clone(),
Arc::default(),
/*tool_catalog_cache_context*/ None,
CancellationToken::new(),
launch("first"),
);
let second = pool.connect(
Arc::default(),
key,
Arc::default(),
/*tool_catalog_cache_context*/ None,
CancellationToken::new(),
launch("second"),
);
tokio::pin!(first, second);
assert!(futures::poll!(&mut first).is_pending());
assert!(futures::poll!(&mut second).is_pending());
started.notified().await;
release.notify_one();
let first = tokio::time::timeout(Duration::from_secs(10), async {
tokio::select! {
first = &mut first => first,
_ = &mut second => panic!("the joining chat should still be starting"),
}
})
.await
.expect("the joining chat must start before the first finishes")
.unwrap();
first
.call_tool(
"test".to_string(),
/*arguments*/ None,
|_| Ok(None),
/*timeout*/ None,
)
.await
.unwrap();
let mut owner = create_ready_async_managed_client(Vec::new()).await;
owner.client = futures::future::ready(Ok(first.clone())).boxed().shared();
tokio::time::timeout(Duration::from_secs(10), owner.shutdown())
.await
.expect("private owner shutdown must not wait for the joining chat");
assert!(first.client.is_closed().await);
release_second.notify_one();
let second = second.await.unwrap();
assert_eq!(starts.load(Ordering::SeqCst), 2);
assert_eq!(*callbacks.lock().unwrap(), ["first", "second"]);
second
.call_tool(
"test".to_string(),
/*arguments*/ None,
|_| Ok(None),
/*timeout*/ None,
)
.await
.unwrap();
}
#[tokio::test]
async fn connection_key_changes_with_the_owner_but_survives_credential_refresh() {
let manager =
AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing());
let mut original = key(&context(), &config());
original.auth_state = Some((manager, AuthChangeState::default()));
let mut refreshed = original.clone();
refreshed.auth_state.as_mut().unwrap().1.generation += 1;
assert!(original == refreshed);
refreshed.auth_state.as_mut().unwrap().1.owner_generation += 1;
assert!(original != refreshed);
}
#[tokio::test(start_paused = true)]
async fn new_chat_relists_tools_without_invalidating_or_waiting_for_existing_calls() {
let pool = McpConnectionPool::default();
let key = key(&context(), &config());
let server = TestServer::default();
let starts = Arc::new(AtomicUsize::new(0));
let release = Arc::new(Notify::new());
release.notify_one();
let launch = startup(
ConnectionSharing::Negotiated,
Arc::clone(&starts),
Arc::default(),
release,
server.clone(),
);
let first = pool
.connect(
Arc::default(),
key.clone(),
Arc::default(),
/*tool_catalog_cache_context*/ None,
CancellationToken::new(),
&launch,
)
.await
.unwrap();
let before = first.tool_catalog.read(|snapshot| snapshot).await;
let release_call = Notify::new();
let call = first
.tool_catalog
.run_with_snapshot(&before, || release_call.notified());
tokio::pin!(call);
assert!(futures::poll!(&mut call).is_pending());
for tool_name in ["test", "new"] {
*server.tools.lock().unwrap() = vec![create_test_tool("test_server", tool_name).tool];
let second = tokio::time::timeout(
Duration::from_secs(1),
pool.connect(
Arc::default(),
key.clone(),
Arc::default(),
/*tool_catalog_cache_context*/ None,
CancellationToken::new(),
&launch,
),
)
.await
.expect("new chats must not wait for existing calls")
.unwrap();
assert_eq!(starts.load(Ordering::SeqCst), 1);
assert_eq!(
second
.tool_catalog
.read(|snapshot| snapshot.tools[0].tool.name.to_string())
.await,
tool_name
);
tokio::time::timeout(
Duration::from_secs(1),
second.tool_catalog.refresh(
|| async { Ok((vec![create_test_tool("test_server", "refreshed")], None, ())) },
|_, _, ()| {},
),
)
.await
.expect("one chat's refresh must not wait for another chat's calls")
.unwrap();
assert_eq!(
first
.tool_catalog
.read(|snapshot| snapshot.tools[0].tool.name.to_string())
.await,
"test"
);
}
release_call.notify_one();
assert_eq!(call.await, Some(()));
server.block_listing.store(true, Ordering::Release);
let failed = pool
.connect(
Arc::default(),
key,
Arc::default(),
/*tool_catalog_cache_context*/ None,
CancellationToken::new(),
&launch,
)
.await;
assert!(failed.is_err());
let _cancelled =
tokio::time::timeout(Duration::from_secs(1), server.cancelled_listings.acquire())
.await
.expect("timed-out listings must be cancelled")
.unwrap();
first
.call_tool(
"test".to_string(),
/*arguments*/ None,
|_| Ok(None),
/*timeout*/ None,
)
.await
.unwrap();
}
#[tokio::test]
async fn negotiation_requires_the_stateless_protocol_and_detaches_startup_callbacks() {
for (protocol_version, expected_sharing, expected_action) in [
(
ProtocolVersion::V_2025_06_18,
false,
ElicitationAction::Decline,
),
(
ProtocolVersion::V_2026_07_28,
true,
ElicitationAction::Cancel,
),
] {
let (negotiation, callback) = SharingNegotiation::new(
Box::new(|_, _| {
futures::future::ready(Ok(ElicitationResponse {
action: ElicitationAction::Decline,
content: None,
meta: None,
}))
.boxed()
}),
CancellationToken::new(),
);
let sharing = negotiation.finish(&protocol_version);
assert_eq!(
matches!(sharing, ConnectionSharing::Negotiated),
expected_sharing
);
let response = callback(
rmcp::model::RequestId::Number(0),
Elicitation::OpenAiForm {
meta: None,
message: "unexpected".to_string(),
requested_schema: json!({}),
},
)
.await
.unwrap();
assert_eq!(response.action, expected_action);
}
}
+2
View File
@@ -6,6 +6,7 @@ pub use codex_rmcp_client::McpProtocolMode;
pub use connection_manager::NetworkPolicyMetadata;
pub use connection_manager::NetworkPolicyMetadataProvider;
pub use connection_manager::tool_is_model_visible;
pub use connection_pool::McpConnectionPool;
pub use elicitation::ElicitationLifecycle;
pub use elicitation::ElicitationReviewRequest;
pub use elicitation::ElicitationReviewer;
@@ -118,6 +119,7 @@ mod client_capabilities;
mod client_tool_catalog;
pub(crate) mod codex_apps;
pub(crate) mod connection_manager;
mod connection_pool;
pub(crate) mod elicitation;
mod ema_account;
mod ema_auth;
+96 -50
View File
@@ -9,6 +9,8 @@
//! Higher-level aggregation and resource/tool APIs live in
//! [`crate::connection_manager`].
#[path = "rmcp_client/managed_client.rs"]
mod managed_client;
#[path = "rmcp_client/status.rs"]
mod status;
@@ -16,6 +18,8 @@ mod status;
#[path = "rmcp_client/provenance_tests.rs"]
mod provenance_tests;
pub(crate) use managed_client::ManagedClient;
use std::borrow::Cow;
use std::collections::BTreeMap;
use std::collections::HashMap;
@@ -33,6 +37,10 @@ use crate::codex_apps::normalize_codex_apps_callable_name;
use crate::codex_apps::normalize_codex_apps_callable_namespace;
use crate::codex_apps::normalize_codex_apps_tool_title;
use crate::codex_apps::prepare_openai_file_params_for_model;
use crate::connection_pool::ConnectionKey;
use crate::connection_pool::ConnectionSharing;
use crate::connection_pool::ConnectionSlot;
use crate::connection_pool::SharingNegotiation;
use crate::elicitation::ElicitationRequestManager;
use crate::executor_environment_http_client::ExecutorEnvironmentHttpClient;
use crate::mcp::CODEX_APPS_MCP_SERVER_NAME;
@@ -94,7 +102,6 @@ use rmcp::model::Tool as RmcpTool;
use tokio::sync::watch;
use tokio::time::Instant as TokioInstant;
use tokio_util::sync::CancellationToken;
use tokio_util::task::AbortOnDropHandle;
use tracing::Instrument;
use tracing::instrument;
use tracing::warn;
@@ -125,45 +132,6 @@ const UNTRUSTED_CONNECTOR_META_KEYS: &[&str] = &[
"connectorDescription",
];
#[derive(Clone)]
pub(crate) struct ManagedClient {
pub(crate) _auth_change_notifications: Option<Arc<AbortOnDropHandle<()>>>,
pub(crate) client: Arc<RmcpClient>,
pub(crate) server_info: McpServerInfo,
pub(crate) tool_catalog: Arc<ClientToolCatalog>,
pub(crate) tool_timeout: Option<Duration>,
pub(crate) server_instructions: Option<String>,
pub(crate) server_supports_sandbox_state_meta_capability: bool,
pub(crate) client_supports_network_policy_meta_capability: bool,
pub(crate) codex_apps_tools_cache_context: Option<ConnectorRuntimeContext<ToolInfo>>,
}
impl ManagedClient {
pub(crate) async fn listed_catalog(&self) -> (Vec<ToolInfo>, Option<Arc<serde_json::Value>>) {
let total_start = Instant::now();
self.tool_catalog
.read(|catalog| {
// Discovery may use the shared cache until this client is refreshed.
// Executable bindings always capture this client's own catalog.
if catalog.revision == 0
&& let Some(cache_context) = &self.codex_apps_tools_cache_context
{
let tools = cache_context.current_catalog_snapshot();
emit_duration(
MCP_TOOLS_LIST_DURATION_METRIC,
total_start.elapsed(),
&[("cache", if tools.is_some() { "hit" } else { "miss" })],
);
if let Some(tools) = tools {
return (tools.tools().to_vec(), tools.shared_tool_list_meta());
}
}
(catalog.tools.to_vec(), catalog.meta)
})
.await
}
}
pub(crate) type ManagedClientFuture =
Shared<BoxFuture<'static, Result<ManagedClient, StartupOutcomeError>>>;
type StartupShutdown = Shared<BoxFuture<'static, Result<(), String>>>;
@@ -328,7 +296,7 @@ struct ManagedClientStartup {
}
impl ManagedClientStartup {
fn start(&self) -> ManagedClientFuture {
fn start(&self, sharing_declined: Option<CancellationToken>) -> ManagedClientFuture {
// A new attempt must not expose capabilities from an earlier connection.
*self
.server_capabilities
@@ -408,6 +376,7 @@ impl ManagedClientStartup {
server_name.clone(),
client,
StartServerTaskParams {
sharing_declined,
is_codex_apps_mcp_server,
startup_timeout: Some(startup_timeout),
tx_event,
@@ -491,6 +460,7 @@ pub(crate) struct AsyncManagedClient {
pub(crate) tool_catalog_cache_context: Option<McpToolCatalogCacheContext>,
pub(crate) startup_complete: Arc<AtomicBool>,
pub(crate) startup_reconnect: Option<Arc<CodexAppsStartupReconnect>>,
pub(crate) connection_slot: Option<Arc<ConnectionSlot>>,
pub(crate) cancel_token: CancellationToken,
}
@@ -520,6 +490,7 @@ impl AsyncManagedClient {
auth_changes: Option<watch::Receiver<AuthChangeState>>,
protocol_mode: McpProtocolMode,
catalog_item_limit: usize,
connection_pool: Option<(crate::McpConnectionPool, ConnectionKey)>,
) -> Self {
let is_codex_apps_mcp_server = server_name == CODEX_APPS_MCP_SERVER_NAME;
let reconnect_server_name = server_name.clone();
@@ -569,19 +540,64 @@ impl AsyncManagedClient {
startup_complete: Arc::clone(&startup_complete),
server_capabilities: Arc::clone(&server_capabilities),
});
// Startup can publish a legacy fallback. Readers wait for the negotiated catalog.
// A discovery attempt may publish a legacy catalog. It cannot read that catalog
// before negotiation or publish a modern joiner's tool list into it.
let tool_catalog_cache_context = tool_catalog_cache_context
.filter(|_| protocol_mode != McpProtocolMode::DiscoverLocalStdio);
let client = startup.start();
let (client, connection_slot) = if let Some((pool, key)) = connection_pool {
let slot = Arc::new(ConnectionSlot::default());
let connection_slot = Arc::clone(&slot);
let startup = Arc::clone(&startup);
let cancel_token = cancel_token.clone();
let startup_complete = Arc::clone(&startup_complete);
let server_capabilities = Arc::clone(&server_capabilities);
let tool_catalog_cache_context = tool_catalog_cache_context.clone();
let client = async move {
let outcome = pool
.connect(
slot,
key,
server_capabilities,
tool_catalog_cache_context,
cancel_token,
|cancellation, sharing_declined| {
let mut shared_startup = (*startup).clone();
shared_startup.cancel_token = cancellation;
shared_startup.start(sharing_declined)
},
)
.await
.map(|mut client| {
if matches!(client.connection_sharing, ConnectionSharing::Shared(_)) {
client.shared_elicitation =
Some(Arc::new(startup.elicitation_requests.make_sender(
startup.server_name.clone(),
startup.tx_event.clone(),
&startup.client_mcp_extensions,
)));
}
client
});
startup_complete.store(true, Ordering::Release);
outcome
}
.boxed()
.shared();
(client, Some(connection_slot))
} else {
(startup.start(/*sharing_declined*/ None), None)
};
let startup_reconnect = is_codex_apps_mcp_server.then(|| {
let startup = Arc::clone(&startup);
Arc::new(
CodexAppsStartupReconnect::new(Arc::new(move || startup.start()))
.with_startup_status_context(
startup_submit_id,
reconnect_server_name,
reconnect_tx_event,
),
CodexAppsStartupReconnect::new(Arc::new(move || {
startup.start(/*sharing_declined*/ None)
}))
.with_startup_status_context(
startup_submit_id,
reconnect_server_name,
reconnect_tx_event,
),
)
});
Self {
@@ -594,6 +610,7 @@ impl AsyncManagedClient {
tool_catalog_cache_context,
startup_complete,
startup_reconnect,
connection_slot,
cancel_token,
}
}
@@ -669,13 +686,29 @@ impl AsyncManagedClient {
pub(crate) async fn shutdown(&self) {
self.cancel_token.cancel();
match self.client().await {
Ok(client) => client.client.shutdown().await,
Ok(client) => match &client.connection_sharing {
ConnectionSharing::Owned(lease) | ConnectionSharing::Shared(lease) => {
if let Some(slot) = lease.release() {
return slot.wait_for_shutdown().await;
}
if matches!(&client.connection_sharing, ConnectionSharing::Owned(_)) {
// Waiters from startup may still hold leases even though sharing was declined.
client.client.shutdown().await;
}
}
ConnectionSharing::Private | ConnectionSharing::Negotiated => {
client.client.shutdown().await;
}
},
Err(StartupOutcomeError::Cancelled) => {}
Err(error) => {
warn!("failed to initialize MCP client during shutdown: {error:#}");
error.wait_for_shutdown().await;
}
}
if let Some(slot) = &self.connection_slot {
slot.wait_for_shutdown().await;
}
}
pub(crate) fn has_cached_tools(&self) -> bool {
@@ -1081,6 +1114,7 @@ async fn start_server_task(
params: StartServerTaskParams,
) -> Result<ManagedClient, StartupOutcomeError> {
let StartServerTaskParams {
sharing_declined,
is_codex_apps_mcp_server,
startup_timeout,
tx_event,
@@ -1098,6 +1132,12 @@ async fn start_server_task(
elicitation_requests.make_sender(server_name.clone(), tx_event, &client_mcp_extensions);
let client_supports_network_policy_meta =
client_mcp_extensions.contains(MCP_NETWORK_POLICY_META_CAPABILITY);
let (sharing_negotiation, send_elicitation) = if let Some(sharing_declined) = sharing_declined {
let (negotiation, sender) = SharingNegotiation::new(send_elicitation, sharing_declined);
(Some(negotiation), sender)
} else {
(None, send_elicitation)
};
let mut params =
mcp_initialize_request_params(client_elicitation_capability, client_mcp_extensions);
if auth_changes.is_some() {
@@ -1130,6 +1170,9 @@ async fn start_server_task(
{
context.invalidate_if_newest(ticket);
}
let connection_sharing = sharing_negotiation
.map(|negotiation| negotiation.finish(&initialize_result.protocol_version))
.unwrap_or_default();
*server_capabilities
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner) =
@@ -1234,6 +1277,8 @@ async fn start_server_task(
);
}
let managed = ManagedClient {
connection_sharing,
shared_elicitation: None,
_auth_change_notifications: auth_change_notifications,
client: Arc::clone(&client),
server_info,
@@ -1338,6 +1383,7 @@ fn mcp_server_info_from_implementation(
}
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.
@@ -0,0 +1,150 @@
//! A chat's access to an initialized MCP client and its tool catalog.
//! A pooled connection honors this chat's lease and routes input requests back to it.
use std::future::Future;
use std::sync::Arc;
use std::time::Duration;
use std::time::Instant;
use super::MCP_TOOLS_LIST_DURATION_METRIC;
use crate::client_tool_catalog::ClientToolCatalog;
use crate::connection_pool::ConnectionSharing;
use crate::runtime::emit_duration;
use crate::tools::ToolInfo;
use anyhow::Result;
use anyhow::anyhow;
use codex_connectors::ConnectorRuntimeContext;
use codex_protocol::mcp::McpServerInfo;
use codex_rmcp_client::RmcpClient;
use codex_rmcp_client::SendElicitation;
use rmcp::model::ListResourceTemplatesResult;
use rmcp::model::ListResourcesResult;
use rmcp::model::PaginatedRequestParams;
use rmcp::model::ServerCapabilities;
use serde_json::Value;
use tokio_util::task::AbortOnDropHandle;
#[derive(Clone)]
pub(crate) struct ManagedClient {
pub(crate) connection_sharing: ConnectionSharing,
/// The calling chat's input handler, separate from the shared transport's startup handler.
pub(crate) shared_elicitation: Option<Arc<SendElicitation>>,
pub(crate) _auth_change_notifications: Option<Arc<AbortOnDropHandle<()>>>,
pub(crate) client: Arc<RmcpClient>,
pub(crate) server_info: McpServerInfo,
pub(crate) tool_catalog: Arc<ClientToolCatalog>,
pub(crate) tool_timeout: Option<Duration>,
pub(crate) server_instructions: Option<String>,
pub(crate) server_supports_sandbox_state_meta_capability: bool,
pub(crate) client_supports_network_policy_meta_capability: bool,
pub(crate) codex_apps_tools_cache_context: Option<ConnectorRuntimeContext<ToolInfo>>,
}
impl ManagedClient {
pub(crate) async fn with_active_shared_lease<T>(
&self,
operation: impl Future<Output = T>,
) -> Result<T> {
match &self.connection_sharing {
ConnectionSharing::Shared(lease) => tokio::select! {
biased;
() = lease.cancellation().cancelled() => {
Err(anyhow!("MCP connection lease has been released"))
}
result = operation => Ok(result),
},
ConnectionSharing::Private
| ConnectionSharing::Negotiated
| ConnectionSharing::Owned(_) => Ok(operation.await),
}
}
pub(crate) async fn call_tool(
&self,
name: String,
arguments: Option<Value>,
meta_for_session: impl Fn(&ServerCapabilities) -> Result<Option<Value>> + Send + Sync,
timeout: Option<Duration>,
) -> Result<rmcp::model::CallToolResult> {
self.with_active_shared_lease(async {
if let Some(send_elicitation) = &self.shared_elicitation {
self.client
.call_tool_with_elicitation(
name,
arguments,
meta_for_session,
timeout,
Arc::clone(send_elicitation),
)
.await
} else {
self.client
.call_tool(name, arguments, meta_for_session, timeout)
.await
}
})
.await?
}
/// Lists resources within this chat's lease lifetime.
pub(crate) async fn list_resources(
&self,
params: Option<PaginatedRequestParams>,
timeout: Option<Duration>,
) -> Result<ListResourcesResult> {
self.with_active_shared_lease(self.client.list_resources(params, timeout))
.await?
}
/// Lists resource templates within this chat's lease lifetime.
pub(crate) async fn list_resource_templates(
&self,
params: Option<PaginatedRequestParams>,
timeout: Option<Duration>,
) -> Result<ListResourceTemplatesResult> {
self.with_active_shared_lease(self.client.list_resource_templates(params, timeout))
.await?
}
/// Reads a resource using this chat's input authority and lease lifetime.
pub(crate) async fn read_resource(
&self,
params: rmcp::model::ReadResourceRequestParams,
timeout: Option<Duration>,
) -> Result<rmcp::model::ReadResourceResult> {
self.with_active_shared_lease(async {
if let Some(send_elicitation) = &self.shared_elicitation {
self.client
.read_resource_with_elicitation(params, timeout, Arc::clone(send_elicitation))
.await
} else {
self.client.read_resource(params, timeout).await
}
})
.await?
}
pub(crate) async fn listed_catalog(&self) -> (Vec<ToolInfo>, Option<Arc<serde_json::Value>>) {
let total_start = Instant::now();
self.tool_catalog
.read(|catalog| {
// Discovery may use the shared cache until this client is refreshed.
// Executable bindings always capture this client's own catalog.
if catalog.revision == 0
&& let Some(cache_context) = &self.codex_apps_tools_cache_context
{
let tools = cache_context.current_catalog_snapshot();
emit_duration(
MCP_TOOLS_LIST_DURATION_METRIC,
total_start.elapsed(),
&[("cache", if tools.is_some() { "hit" } else { "miss" })],
);
if let Some(tools) = tools {
return (tools.tools().to_vec(), tools.shared_tool_list_meta());
}
}
(catalog.tools.to_vec(), catalog.meta)
})
.await
}
}
+51 -3
View File
@@ -8,6 +8,10 @@
mod status;
#[cfg(test)]
#[path = "runtime_shared_connection_tests.rs"]
mod shared_connection_tests;
use std::collections::HashMap;
use std::collections::HashSet;
use std::path::PathBuf;
@@ -46,6 +50,7 @@ use rmcp::model::ReadResourceResult;
use rmcp::model::RequestId;
use tokio::sync::watch;
use tokio_util::sync::CancellationToken;
use tracing::warn;
use crate::McpConfig;
use crate::binding::McpBinding;
@@ -54,6 +59,7 @@ use crate::client_tool_catalog::CodexAppsToolSnapshot;
use crate::connection_manager::BindingCatalogRevision;
use crate::connection_manager::McpConnectionSet;
use crate::connection_manager::PreviousMcpConnections;
use crate::connection_pool::McpConnectionPool;
use crate::elicitation::ElicitationLifecycle;
use crate::elicitation::ElicitationRequestRouter;
use crate::elicitation::ElicitationReviewerHandle;
@@ -112,6 +118,8 @@ pub struct McpRuntimeInput {
/// their exact connections and configuration for as long as they are needed.
pub struct McpRuntime {
current: ArcSwap<PublishedMcpRuntime>,
// Old bindings keep leases alive after publication; cancellation alone does not release them.
superseded_connections: Mutex<Vec<Weak<McpConnectionSet>>>,
event_stream_cancellation: Mutex<EventStreamCancellation>,
reconnect_pending: AtomicBool,
resource_cache_generation: AtomicU64,
@@ -269,6 +277,7 @@ impl McpRuntime {
ready_environments: HashMap::new(),
cached_binding: Mutex::new(None),
}),
superseded_connections: Mutex::default(),
event_stream_cancellation: Mutex::new(EventStreamCancellation {
event_server_available: false,
cancel_event_streams_on_server_removal: watch::channel(()).0,
@@ -398,6 +407,9 @@ impl McpRuntime {
let auth = input.auth.clone();
let auth_manager = input.auth_manager.clone();
let auth_token = auth.as_ref().and_then(|auth| auth.get_token().ok());
if !reuse_previous {
current.connections.retire_shared_connections();
}
let auth_matches = current.auth_matches(auth.as_ref());
let auth_owner = AuthOwner::capture(auth.as_ref(), auth_manager.as_deref());
let same_owner = match (current.auth_owner, auth_owner) {
@@ -421,8 +433,9 @@ impl McpRuntime {
let ready_selected_capability_roots = input.ready_selected_capability_roots.clone();
let environment_selections = Arc::clone(&input.runtime_context.environment_selections);
let ready_environments = input.runtime_context.ready_environments.clone();
// Keep MCP startup state out of the caller's nested session-initialization future.
let connections = Arc::new(
McpConnectionSet::new(
Box::pin(McpConnectionSet::new(
ema_auth,
retain_ema.then_some(current.connections.as_ref()),
reuse_previous.then_some(PreviousMcpConnections {
@@ -432,7 +445,7 @@ impl McpRuntime {
publication_gate,
input,
self.elicitation_router.clone(),
)
))
.await,
);
let hosted_event_server_retained = connections.contains_server(CODEX_APPS_MCP_SERVER_NAME)
@@ -448,6 +461,14 @@ impl McpRuntime {
.event_stream_cancellation
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
{
let mut superseded = self
.superseded_connections
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
superseded.retain(|connections| connections.strong_count() > 0);
superseded.push(Arc::downgrade(&current.connections));
}
self.current.store(Arc::new(PublishedMcpRuntime {
connections,
config: Some(config),
@@ -923,7 +944,26 @@ impl McpRuntime {
pub async fn shutdown(&self) {
self.elicitation_router.close().await;
self.latest_connections().shutdown().await;
let superseded = std::mem::take(
&mut *self
.superseded_connections
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner),
);
let current = self.latest_connections();
// Keep every publication's cleanup alive if the caller is interrupted.
let shutdown = tokio::spawn(async move {
for connections in superseded
.into_iter()
.filter_map(|connections| connections.upgrade())
{
connections.shutdown_shared_connections().await;
}
current.shutdown().await;
});
if let Err(error) = shutdown.await {
warn!("MCP runtime shutdown task failed: {error}");
}
}
}
@@ -933,6 +973,7 @@ impl McpRuntime {
/// environment registry, captured thread selections, and the host-local cwd used by MCP processes.
#[derive(Clone)]
pub struct McpRuntimeContext {
pub(crate) connection_pool: Option<McpConnectionPool>,
environment_manager: Arc<EnvironmentManager>,
environment_selections: Arc<[TurnEnvironmentSelection]>,
ready_environments: HashMap<String, Arc<Environment>>,
@@ -981,6 +1022,7 @@ impl McpRuntimeContext {
.with_tls_backend_fallback(),
);
Self {
connection_pool: None,
environment_manager,
environment_selections: Arc::default(),
ready_environments: HashMap::new(),
@@ -989,6 +1031,12 @@ impl McpRuntimeContext {
}
}
/// Uses this process's pool for compatible local stdio servers using MCP 2026-07-28.
pub fn with_connection_pool(mut self, pool: McpConnectionPool) -> Self {
self.connection_pool = Some(pool);
self
}
/// Pins all captured selections and the concrete handles for environments already ready.
pub fn with_selected_environments(
mut self,
@@ -0,0 +1,172 @@
use super::*;
use crate::connection_manager::tests::create_ready_async_managed_client;
use crate::connection_pool::ConnectionSharing;
use crate::connection_pool::ConnectionSlot;
use crate::connection_pool::McpConnectionPool;
use crate::connection_pool::tests as fixture;
use futures::FutureExt;
use pretty_assertions::assert_eq;
use std::sync::atomic::AtomicUsize;
use tokio::sync::Notify;
#[tokio::test]
async fn forced_refresh_replaces_pending_connection_without_interrupting_other_chats() {
let pool = McpConnectionPool::default();
let key = fixture::key(&fixture::context(), &fixture::config());
let starts = Arc::new(AtomicUsize::new(0));
let started = Arc::new(Notify::new());
let release = Arc::new(Notify::new());
let launch = Arc::new(fixture::startup(
ConnectionSharing::Negotiated,
Arc::clone(&starts),
Arc::clone(&started),
Arc::clone(&release),
fixture::TestServer::default(),
));
let slot = Arc::new(ConnectionSlot::default());
let first = {
let pool = pool.clone();
let key = key.clone();
let slot = Arc::clone(&slot);
let launch = Arc::clone(&launch);
async move {
pool.connect(
slot,
key,
Arc::default(),
/*tool_catalog_cache_context*/ None,
CancellationToken::new(),
&*launch,
)
.await
}
.boxed()
.shared()
};
let mut poll_first = first.clone();
assert!(futures::poll!(&mut poll_first).is_pending());
let other = pool.connect(
Arc::default(),
key.clone(),
Arc::default(),
/*tool_catalog_cache_context*/ None,
CancellationToken::new(),
&*launch,
);
tokio::pin!(other);
assert!(futures::poll!(&mut other).is_pending());
started.notified().await;
let mut owner = create_ready_async_managed_client(Vec::new()).await;
owner.client = first;
owner.connection_slot = Some(slot);
let mut connections = McpConnectionSet::empty(/*prefix_mcp_tool_names*/ true);
connections.insert_test_client("test_server", owner);
let connections = Arc::new(connections);
let config = Arc::new(crate::mcp::tests::test_mcp_config(std::env::temp_dir()));
let runtime = McpRuntime::empty(/*prefix_mcp_tool_names*/ true);
runtime.current.store(Arc::new(PublishedMcpRuntime {
connections: Arc::clone(&connections),
config: Some(Arc::clone(&config)),
auth: None,
auth_manager: None,
auth_token: None,
auth_owner: None,
auth_generation: Arc::new(()),
plugins_available: false,
ready_selected_capability_roots: Vec::new(),
environment_selections: Arc::default(),
ready_environments: HashMap::new(),
cached_binding: Mutex::new(None),
}));
let input = McpRuntimeInput {
network_policy_metadata: None,
startup_policy: McpStartupPolicy::Eager,
config,
plugins_available: false,
ready_selected_capability_roots: Vec::new(),
mcp_servers: HashMap::new(),
submit_id: "refresh".to_string(),
tx_event: None,
startup_cancellation_token: CancellationToken::new(),
runtime_context: fixture::context(),
codex_apps_tools_cache: ConnectorRuntimeManager::default(),
tool_catalog_cache: McpToolCatalogCache::default(),
codex_apps_tools_cache_key: ConnectorRuntimeContextKey::personal(
/*account_id*/ None, /*chatgpt_user_id*/ None,
),
client_mcp_extensions: ClientMcpExtensions::default(),
auth: None,
auth_manager: None,
elicitation_reviewer: None,
elicitation_lifecycle: None,
};
let binding = runtime.current_binding_for_call("test_server");
tokio::pin!(binding);
assert!(futures::poll!(&mut binding).is_pending());
runtime.reconnect_on_next_refresh();
runtime.replace(input).await;
let _replacement = tokio::time::timeout(
Duration::from_secs(1),
pool.connect(
Arc::default(),
key.clone(),
Arc::default(),
/*tool_catalog_cache_context*/ None,
CancellationToken::new(),
&*launch,
),
)
.await
.expect("forced refresh must not wait for the old startup")
.unwrap();
release.notify_one();
let other = other.await.unwrap();
other
.call_tool(
"test".to_string(),
/*arguments*/ None,
|_| Ok(None),
/*timeout*/ None,
)
.await
.unwrap();
// A second refresh of the old owner must not evict the newer replacement.
connections.retire_shared_connections();
let _joined = pool
.connect(
Arc::default(),
key,
Arc::default(),
/*tool_catalog_cache_context*/ None,
CancellationToken::new(),
&*launch,
)
.await
.unwrap();
assert_eq!(starts.load(Ordering::SeqCst), 2);
let binding = binding.await.expect("binding from the old publication");
runtime.shutdown().await;
assert!(
binding
.list_resources("test_server", /*params*/ None)
.await
.is_err()
);
other
.call_tool(
"test".to_string(),
/*arguments*/ None,
|_| Ok(None),
/*timeout*/ None,
)
.await
.unwrap();
let transport = Arc::clone(&other.client);
let mut other_owner = create_ready_async_managed_client(Vec::new()).await;
other_owner.client = futures::future::ready(Ok(other)).boxed().shared();
other_owner.shutdown().await;
assert!(transport.is_closed().await);
drop(binding);
}
+6
View File
@@ -111,6 +111,7 @@ enum OrderedMcpOverlay {
#[derive(Clone)]
pub struct McpManager {
connection_pool: codex_mcp::McpConnectionPool,
plugins_manager: Arc<PluginsManager>,
extensions: Arc<ExtensionRegistry<Config>>,
codex_apps_tools_cache: ConnectorRuntimeManager<ToolInfo>,
@@ -135,6 +136,7 @@ impl McpManager {
codex_apps_tools_cache: ConnectorRuntimeManager<ToolInfo>,
) -> Self {
Self {
connection_pool: codex_mcp::McpConnectionPool::default(),
plugins_manager,
extensions,
codex_apps_tools_cache,
@@ -152,6 +154,10 @@ impl McpManager {
self.tool_catalog_cache.clone()
}
pub(crate) fn connection_pool(&self) -> codex_mcp::McpConnectionPool {
self.connection_pool.clone()
}
pub(crate) fn shared_definitions(&self) -> &SharedMcpDefinitions {
&self.shared_definitions
}
+2 -1
View File
@@ -175,7 +175,7 @@ impl Session {
self.mcp_runtime_context(&environments, &host_fallback_cwd)
}
fn mcp_runtime_context(
pub(super) fn mcp_runtime_context(
&self,
environments: &TurnEnvironmentSnapshot,
host_fallback_cwd: &std::path::Path,
@@ -188,6 +188,7 @@ impl Session {
self.services.turn_environments.environment_manager(),
local_process_cwd,
)
.with_connection_pool(self.services.mcp_manager.connection_pool())
.with_selected_environments(
environments.all_selections().into(),
environments.ready_environment_handles(),
+2 -8
View File
@@ -358,14 +358,8 @@ impl Session {
}),
);
let mcp_config = Arc::new(config);
let runtime_context = McpRuntimeContext::new(
self.services.turn_environments.environment_manager(),
desired.local_process_cwd.clone(),
)
.with_selected_environments(
desired.environments.all_selections().into(),
desired.environments.ready_environment_handles(),
);
let runtime_context =
self.mcp_runtime_context(&desired.environments, &desired.local_process_cwd);
McpRuntimeInput {
network_policy_metadata: None,
startup_policy: if matches!(desired.session_source, SessionSource::SubAgent(_)) {
+3
View File
@@ -152,6 +152,9 @@ mod mcp_resource_messages;
#[path = "scenarios_mcp_ema.rs"]
mod mcp_ema;
#[path = "scenarios_mcp_sharing.rs"]
mod mcp_sharing;
#[path = "scenarios_guardian_agent_messages_tests.rs"]
mod guardian_agent_messages;
@@ -0,0 +1,179 @@
//! Shared MCP transports preserve each chat's model-visible tools and input responses.
use super::*;
use codex_protocol::approvals::ElicitationAction;
use codex_protocol::approvals::ElicitationRequest;
use codex_protocol::config_types::ApprovalsReviewer;
use codex_protocol::models::PermissionProfile;
use codex_protocol::protocol::AskForApproval;
use codex_protocol::protocol::ThreadSettingsOverrides;
use core_test_support::responses;
use pretty_assertions::assert_eq;
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn shared_mcp_keeps_tools_and_model_requested_input_per_chat() -> Result<()> {
skip_if_no_network!(Ok(()));
core_test_support::skip_if_remote!(Ok(()), "the shared stdio fixture runs on the host");
skip_if_wine_exec!(Ok(()), "requires a Windows test_stdio_server binary");
let first_server = start_mock_server().await;
let second_server = start_mock_server().await;
let command = stdio_server_bin()?;
let test = test_codex()
.with_config(move |config| {
super::configure_scenario_catalog(config);
config.approvals_reviewer = ApprovalsReviewer::User;
config
.features
.enable(Feature::McpConnectionSharing)
.expect("enable sharing");
config
.features
.disable(Feature::Mcp20260728)
.expect("let sharing select the protocol");
config
.mcp_servers
.set(
serde_json::from_value(json!({
"shared": {
"command": command,
"enabled_tools": ["echo"],
"default_tools_approval_mode": "approve",
"env": {
"MCP_TEST_ECHO_INPUT": "1",
},
},
}))
.expect("shared MCP config"),
)
.expect("set MCP config");
})
.with_model_info_override("gpt-5.5", |model| {
model.tool_mode = Some(ToolMode::Direct);
model.supports_search_tool = false;
model.use_responses_lite = false;
})
.build_with_auto_env(&first_server)
.await?;
wait_for_mcp_server(&test.codex, "shared").await?;
let mut options = test.start_thread_options().await;
options.config.model_provider.base_url = Some(format!("{}/v1", second_server.uri()));
let mut servers = options.config.mcp_servers.get().clone();
servers
.get_mut("shared")
.expect("shared server")
.enabled_tools = Some(vec!["echo-tool".to_string()]);
options.config.mcp_servers.set(servers)?;
let second = test.thread_manager.start_thread(options).await?.thread;
wait_for_mcp_server(&second, "shared").await?;
let chats = [
(
&test.codex,
&first_server,
"echo",
"First task owner?",
"Mira",
),
(
&second,
&second_server,
"echo_tool",
"Second task owner?",
"Sora",
),
];
let mut mocks = Vec::new();
for (thread, server, tool, question, _) in chats {
mocks.push(
mount_sse_sequence(
server,
vec![
sse(vec![
ev_response_created("ask-owner"),
ev_function_call_with_namespace(
"ask-owner",
"mcp__shared",
tool,
&json!({ "message": question }).to_string(),
),
ev_completed("ask-owner"),
]),
responses::sse_completed("done"),
],
)
.await,
);
thread
.start_or_steer_turn(
TurnInputRequest::user_input(vec![UserInput::Text {
text: format!("Ask the user: {question}"),
text_elements: Vec::new(),
}])
.with_thread_settings(ThreadSettingsOverrides {
approval_policy: Some(AskForApproval::OnRequest),
permission_profile: Some(PermissionProfile::read_only()),
..Default::default()
}),
)
.await?;
}
let mut inputs = Vec::new();
let mut processes = Vec::new();
for (thread, _, _, question, _) in chats {
let input = wait_for_event_match(thread, |event| match event {
EventMsg::ElicitationRequest(input) => Some(input.clone()),
_ => None,
})
.await;
let ElicitationRequest::Form { meta, message, .. } = &input.request else {
panic!("expected a form input");
};
assert_eq!(message, question);
processes.push(meta.as_ref().expect("fixture metadata")["fixture/processId"].clone());
inputs.push(input);
}
assert!(processes[0].is_u64());
assert_eq!(
processes[0], processes[1],
"the chats must share one process"
);
// Both calls use the same server input key; resolve the joiner's input first.
for ((thread, _, _, _, answer), input) in chats.into_iter().zip(inputs).rev() {
thread
.submit(Op::ResolveElicitation {
server_name: input.server_name,
request_id: input.id,
decision: ElicitationAction::Accept,
content: Some(json!({ "answer": answer })),
meta: None,
})
.await?;
}
let mut histories = Vec::new();
for ((thread, ..), mock) in chats.into_iter().zip(mocks) {
wait_for_event(thread, |event| matches!(event, EventMsg::TurnComplete(_))).await;
histories.push(mock.requests());
}
let entries = histories
.iter()
.zip(["creator", "joiner"])
.flat_map(|(requests, label)| {
requests
.iter()
.map(move |request| SnapshotEntry::captured(request).labeled(label))
})
.collect::<Vec<_>>();
insta::assert_snapshot!(context_snapshot::format_context_snapshot(
"Two chats share one MCP process, expose different tools, and answer simultaneous model-requested questions in reverse order.",
&entries,
&ContextSnapshotOptions::default()
.rewrite_known_segments()
.include_request_settings(),
));
Ok(())
}
@@ -0,0 +1,68 @@
---
source: core/tests/suite/scenarios_mcp_sharing.rs
expression: "context_snapshot::format_context_snapshot(\"Two chats share one MCP process, expose different tools, and answer simultaneous model-requested questions in reverse order.\",\n&entries,\n&ContextSnapshotOptions::default().rewrite_known_segments().include_request_settings(),)"
---
Scenario: Two chats share one MCP process, expose different tools, and answer simultaneous model-requested questions in reverse order.
## Window 1
Settings:
include: ["reasoning.encrypted_content"]
model: "gpt-5.5"
parallel_tool_calls: true
prompt_cache_key: "<PROMPT_CACHE_KEY 1>"
reasoning: {"effort":"medium"}
store: false
stream: true
text: {"verbosity":"low"}
tool_choice: "auto"
tools (11; hash=e4ad6947d6c88cac):
- function/exec_command: Runs a command in a PTY, returning output or a session ID for ongoing interaction.; args=[cmd, justification, login, max_output_tokens, prefix_rule, sandbox_permissions, shell, tty, workdir, yield_time_ms]; hash=c04108e51a8548c5
- function/write_stdin: Writes characters to an existing unified exec session and returns recent output.; args=[chars, max_output_tokens, session_id, yield_time_ms]; hash=7fb62d8a3e2df4e2
- function/list_mcp_resources: Lists resources provided by MCP servers. Resources allow servers to share data that provides cont...; args=[cursor, server]; hash=6792edaff13561bd
- function/list_mcp_resource_templates: Lists resource templates provided by MCP servers. Parameterized resource templates allow servers ...; args=[cursor, server]; hash=60f17c90ef7a4ebb
- function/read_mcp_resource: Read a specific resource from an MCP server given the server name and resource URI.; args=[server, uri]; hash=24246896b39008ae
- function/request_user_input: Request user input for one to three short questions and wait for the response. This tool is only ...; args=[questions]; hash=64a729b25c424fc6
- custom/apply_patch: The `apply_patch` tool can be used to edit files. This is a FREEFORM tool, so do not wrap the pat...; hash=dea9f12b0111a4ba
- function/view_image: View a local image file from the filesystem when visual inspection is needed. Use this for images...; args=[detail, path]; hash=3b766d5bdec5d50e
- namespace/multi_agent_v1: Tools for spawning and managing sub-agents.; hash=2327142bda177c50
- function/close_agent
- function/resume_agent
- function/send_input
- function/spawn_agent
- function/wait_agent
- namespace/mcp__shared: Use these tools to exercise the rmcp test server.; hash=c2583dfa504bbe3f
- function/echo
- web_search; hash=f103e327c6870a46
-- request 1 (turn; creator) --
00:message/developer:
<MODEL_INSTRUCTIONS>
01:message/developer:
<PERMISSIONS_INSTRUCTIONS>
02:message/user:
<ENVIRONMENT_CONTEXT>
03:message/user:
Ask the user: First task owner?
-- request 2 (turn; creator) --
04:function_call/mcp__shared.echo:{"message":"First task owner?"}
05:function_call_output:Wall time: <DURATION> seconds
Output: | Mira
## Window 2 (after request 2: input diverged at item 03)
Settings: relative to window 1
prompt_cache_key: "<PROMPT_CACHE_KEY 2>"
tools (11; hash=b0f271ad4e160a41):
~ namespace/mcp__shared: Use these tools to exercise the rmcp test server.; hash=a9e155b340098fac
- function/echo_tool
-- request 3 (turn; joiner) --
00:message/developer:
<MODEL_INSTRUCTIONS>
01:message/developer:
<PERMISSIONS_INSTRUCTIONS>
02:message/user:
<ENVIRONMENT_CONTEXT>
03:message/user:
Ask the user: Second task owner?
-- request 4 (turn; joiner) --
04:function_call/mcp__shared.echo_tool:{"message":"Second task owner?"}
05:function_call_output:Wall time: <DURATION> seconds
Output: | Sora
+1 -1
View File
@@ -269,7 +269,7 @@ pub enum Feature {
Mcp20260728,
/// Honor ordinary MCP server tool-list notifications within a turn.
McpToolListChanged,
/// Discover the MCP protocol for eligible local stdio servers.
/// Discover and share eligible local stdio connections across chats.
McpConnectionSharing,
/// Enable MCP protocol version 2026-07-28 for the host-owned Codex Apps server.
CodexAppsMcp20260728,
@@ -13,9 +13,15 @@ use rmcp::ServiceExt;
use rmcp::handler::server::ServerHandler;
use rmcp::model::CallToolRequestParams;
use rmcp::model::CallToolResult;
use rmcp::model::ContentBlock;
use rmcp::model::ElicitRequest;
use rmcp::model::ElicitRequestParams;
use rmcp::model::ElicitationSchema;
use rmcp::model::Implementation;
use rmcp::model::InitializeRequestParams;
use rmcp::model::InitializeResult;
use rmcp::model::InputRequest;
use rmcp::model::InputRequiredResult;
use rmcp::model::JsonObject;
use rmcp::model::ListResourceTemplatesResult;
use rmcp::model::ListResourcesResult;
@@ -781,6 +787,30 @@ impl ServerHandler for TestToolServer {
}
};
if std::env::var_os("MCP_TEST_ECHO_INPUT").is_some() {
if let Some(responses) = request.input_responses {
let answer = responses["confirmation"]["content"]["answer"]
.as_str()
.ok_or_else(|| McpError::invalid_params("missing answer text", None))?;
return Ok(CallToolResult::success(vec![ContentBlock::text(answer)]).into());
}
let mut meta = MetaObject::new();
meta.insert("fixture/processId".into(), json!(std::process::id()));
let form = ElicitRequestParams::FormElicitationParams {
meta: Some(meta.into()),
message: args.message,
requested_schema: ElicitationSchema::builder()
.required_string_property("answer", |s| s)
.build()
.map_err(|error| McpError::internal_error(error.to_string(), None))?,
};
let input = InputRequest::Elicitation(ElicitRequest::new(form));
return Ok(InputRequiredResult::from_input_requests(BTreeMap::from([
("confirmation".into(), input),
]))
.into());
}
let env_snapshot: HashMap<String, String> = std::env::vars().collect();
let env_name = args.env_var.as_deref().unwrap_or("MCP_TEST_VALUE");
let echo = dynamic_server_process_label()