openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-29 19:08:43 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add per-chat leases for shared MCP connections (#57978)

Browse Source 
## What changed

Negotiate the `codex/connection-lease` extension for pooled stdio connections. For servers that advertise it, open a unique lease per chat and attach its host-owned ID to tool calls and resource reads via `codex/connectionLease`, overriding caller-supplied lease metadata.

When a chat releases its lease, write cancellations for pending requests before sending `codex/connection/release`. Preserve open-before-release ordering across startup timeouts, bound unload waits, and remove connections from the pool when lease cleanup fails or times out. Servers without the extension retain existing behavior.

## Testing

Add coverage for cancellation-before-release ordering, late open acknowledgements, and lease drops outside the runtime. Extend app-server sharing tests to verify distinct chat leases, metadata replacement, release on unload, and continued use by another chat.

GitOrigin-RevId: f4c8b606cb98b5697acab1c515e992ff80ec8425
This commit is contained in:
Charlie Marsh
2026-10-29 03:47:52 +00:00
committed by copyberry
parent 2e764743a6
commit 8fdf647295
10 changed files with 568 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/Cargo.lock
Generated
+1
View File
@@ -4036,6 +4036,7 @@ dependencies = [
"tokio-util",
"tracing",
"url",
"uuid",
]
[[package]]
codex-rs/app-server/tests/suite/v2/mcp_connection_sharing.rs
+31 -1
View File
@@ -31,6 +31,7 @@ async fn stdio_sharing_follows_protocol_and_enablement_for_loaded_chats() -> Res
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 lease_log = codex_home.path().join("leases.jsonl");
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");
@@ -39,8 +40,15 @@ async fn stdio_sharing_follows_protocol_and_enablement_for_loaded_chats() -> Res
"[mcp_servers.{name}]\ncommand = {command}\nenabled_tools = [\"echo\"]\n\
[mcp_servers.{name}.env]\nMCP_TEST_DYNAMIC_SERVER_METADATA = \"1\"\n"
));
// Exercise lease accounting on the shared server in the mixed-protocol case.
if name == "alpha" && !modern {
let log = toml::Value::String(lease_log.to_string_lossy().into_owned());
config.push_str(&format!("MCP_TEST_LEASE_EVENTS_FILE = {log}\n"));
}
if name == "beta" && modern {
// An experimental-only advertisement must leave this server usable without leases.
config.push_str("CODEX_MCP_PROTOCOL_VERSION = \"2026-07-28\"\n");
config.push_str("MCP_TEST_EXPERIMENTAL_LEASE_ONLY = \"1\"\n");
} else if name == "beta" {
config.push_str("CODEX_MCP_PROTOCOL_VERSION = \"2025-06-18\"\n");
}
@@ -78,6 +86,28 @@ async fn stdio_sharing_follows_protocol_and_enablement_for_loaded_chats() -> Res
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?;
if !modern {
let log = std::fs::read_to_string(&lease_log)?;
let events = log
.lines()
.map(serde_json::from_str::<(&str, &str)>)
.collect::<Result<Vec<_>, _>>()?;
assert_eq!(events.len(), 6, "{events:?}");
let (a, b) = (events[0].1, events[2].1);
assert_ne!(a, b);
assert!(events.iter().all(|(_, id)| *id != "forged"));
assert_eq!(
events,
[
("codex/connection/open", a),
("tools/call", a),
("codex/connection/open", b),
("tools/call", b),
("codex/connection/release", a),
("tools/call", b),
]
);
}
#[cfg(unix)]
for process in first_processes.values().chain(second_processes.values()) {
core_test_support::process::wait_for_process_exit(
@@ -129,7 +159,7 @@ async fn call_servers(
server: name.to_string(),
tool: "echo".to_string(),
arguments: Some(json!({ "message": "ping" })),
meta: None,
meta: Some(json!({"codex/connectionLease": "forged"})),
},
})
.await?;
codex-rs/codex-mcp/Cargo.toml
+1
View File
@@ -44,6 +44,7 @@ tokio = { workspace = true, features = ["io-util", "macros", "rt-multi-thread"]
tokio-util = { workspace = true, features = ["rt"] }
tracing = { workspace = true }
url = { workspace = true }
uuid = { workspace = true, features = ["v4"] }
[dev-dependencies]
assert_matches = { workspace = true }
codex-rs/codex-mcp/src/connection_pool.rs
+142 -14
View File
@@ -15,12 +15,17 @@ use codex_connectors::ConnectorRuntimeContextKey;
use codex_login::AuthChangeState;
use codex_login::AuthManager;
use codex_rmcp_client::ElicitationResponse;
use codex_rmcp_client::RequestScope;
use codex_rmcp_client::RmcpClient;
use codex_rmcp_client::SendElicitation;
use futures::FutureExt;
use futures::future::BoxFuture;
use futures::future::Shared;
use rmcp::model::ElicitationAction;
use rmcp::model::ProtocolVersion;
use tokio_util::sync::CancellationToken;
use tokio_util::task::TaskTracker;
use uuid::Uuid;
use crate::client_tool_catalog::ClientToolCatalog;
use crate::client_tool_catalog::ToolCatalogSnapshot;
@@ -31,6 +36,40 @@ use crate::rmcp_client::list_tools_for_client_uncached;
use crate::server::McpServerConnectionIdentity;
use crate::tool_catalog_cache::McpToolCatalogCacheContext;
/// Opt-in resource lifecycle for negotiated stdio connections. Tool calls and resource reads
/// carry a host-owned lease ID. Catalog requests remain independent of a particular lease.
pub(crate) const LEASE_CAPABILITY: &str = "codex/connection-lease";
pub(crate) const LEASE_META_KEY: &str = "codex/connectionLease";
const LEASE_REQUEST_TIMEOUT: Duration = Duration::from_secs(5);
#[derive(Clone)]
pub(crate) struct ServerLease {
pub(crate) id: Uuid,
pub(crate) requests: RequestScope,
opened: Shared<BoxFuture<'static, Result<(), StartupOutcomeError>>>,
}
impl ServerLease {
async fn release(self, client: &RmcpClient, entry: &Arc<Entry>) {
if self.opened.await.is_err() {
return;
}
// MCP cancellation has no acknowledgement. The server owns any handler still
// winding down after it receives these cancellation writes and then release.
self.requests.cancel_and_wait().await;
if let Err(error) = client
.send_custom_request(
"codex/connection/release",
Some(serde_json::json!({"leaseId": self.id})),
)
.await
{
entry.remove_from_pool();
tracing::warn!(%error, "failed to release a server-side MCP lease");
}
}
}
/// Process-owned registry for compatible local stdio servers using MCP 2026-07-28.
#[derive(Clone, Default)]
pub struct McpConnectionPool {
@@ -68,6 +107,7 @@ impl PartialEq for ConnectionKey {
}
struct Entry {
runtime: tokio::runtime::Handle,
key: ConnectionKey,
client: ManagedClientFuture,
capabilities: Arc<Mutex<Option<serde_json::Value>>>,
@@ -154,6 +194,7 @@ pub(crate) struct ConnectionLease {
owner: Mutex<Option<Arc<LeaseOwners>>>,
cancellation: CancellationToken,
slot: Arc<ConnectionSlot>,
server_lease: OnceLock<ServerLease>,
}
/// Tracks an acquisition without retaining its lease, including before startup completes.
@@ -202,6 +243,16 @@ pub(crate) enum ConnectionSharing {
Shared(Arc<ConnectionLease>),
}
impl ConnectionSharing {
/// Identifies this chat to servers that negotiated per-chat leases.
pub(crate) fn server_lease(&self) -> Option<&ServerLease> {
match self {
Self::Shared(lease) => lease.server_lease.get(),
_ => None,
}
}
}
impl McpConnectionPool {
pub(crate) async fn connect(
&self,
@@ -233,6 +284,11 @@ impl McpConnectionPool {
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.clone();
let server_tracks_leases = shared_capabilities
.as_ref()
.and_then(|caps| caps.get("extensions"))
.and_then(|caps| caps.get(LEASE_CAPABILITY))
.is_some();
*capabilities
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner) = shared_capabilities;
@@ -249,6 +305,13 @@ impl McpConnectionPool {
.or_cancel(&cancellation)
.await
.unwrap_or(Err(StartupOutcomeError::Cancelled))?;
if server_tracks_leases {
lease
.open_on_server(Arc::clone(&client.client))
.or_cancel(&cancellation)
.await
.unwrap_or(Err(StartupOutcomeError::Cancelled))?;
}
ConnectionSharing::Shared(lease)
} else {
lease.entry.remove_from_pool();
@@ -316,6 +379,7 @@ impl McpConnectionPool {
tokio::spawn(client.clone());
let owner = Arc::new(LeaseOwners);
let entry = Arc::new(Entry {
runtime: tokio::runtime::Handle::current(),
key: key.clone(),
client,
capabilities: Arc::clone(&capabilities),
@@ -333,6 +397,7 @@ impl McpConnectionPool {
owner: Mutex::new(Some(owner)),
cancellation: lease_cancellation,
slot: Arc::clone(slot),
server_lease: OnceLock::new(),
});
*state = SlotState::Acquired(Arc::downgrade(&lease.entry));
(lease, role)
@@ -354,29 +419,92 @@ impl ConnectionLease {
&self.cancellation
}
/// Returns the slot tracking shutdown when this chat releases the last lease.
/// Register a unique logical owner before publishing the chat's client. Release retains
/// the same future so cancellation during open cannot reorder the two lifecycle requests.
async fn open_on_server(&self, client: Arc<RmcpClient>) -> Result<(), StartupOutcomeError> {
let id = Uuid::new_v4();
let opened = async move {
client
.send_custom_request(
"codex/connection/open",
Some(serde_json::json!({"leaseId": id})),
)
.await
.map(|_| ())
.map_err(StartupOutcomeError::from)
}
.boxed()
.shared();
let opened = {
// Publish while we still own the lease so a racing unload will wait for this open.
let owner = self
.owner
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if owner.is_none() {
return Err(StartupOutcomeError::Cancelled);
}
self.server_lease
.get_or_init(|| ServerLease {
id,
opened,
requests: RequestScope::new(
self.entry.runtime.clone(),
self.cancellation.clone(),
),
})
.opened
.clone()
};
// Bound chat startup while preserving the actual reply for a racing release.
tokio::time::timeout(LEASE_REQUEST_TIMEOUT, opened)
.await
.map_err(|_| {
self.entry.remove_from_pool();
anyhow::anyhow!("timed out waiting for the MCP server to open a connection lease")
})?
}
/// Track this chat's server-side release or the final shutdown of the physical connection.
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()?;
let server_lease = if Arc::into_inner(owner.take()?).is_some() {
self.entry.cancellation.cancel();
self.entry.remove_from_pool();
None
} else {
Some(self.server_lease.get()?.clone())
};
let entry = Arc::clone(&self.entry);
// 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,
let completion = tasks.token();
self.entry.runtime.spawn(async move {
match (entry.client.clone().await, server_lease) {
(Ok(client), None) => client.client.shutdown().await,
(Ok(client), Some(lease)) => {
let release = lease
.release(&client.client, &entry)
.or_cancel(&entry.cancellation);
tokio::pin!(release);
if tokio::time::timeout(LEASE_REQUEST_TIMEOUT, &mut release)
.await
.is_err()
{
// Let the chat unload. A late open reply or cancellation write
// must still precede release while other chats own the connection.
entry.remove_from_pool();
drop(completion);
let _ = release.await;
}
}
},
&runtime,
));
(Err(error), _) => error.wait_for_shutdown().await,
}
});
Some(&self.slot)
}
}
codex-rs/codex-mcp/src/connection_pool_tests.rs
+200 -1
View File
@@ -34,8 +34,12 @@ use std::collections::HashMap;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering;
use tokio::io::AsyncBufReadExt;
use tokio::io::AsyncWriteExt;
use tokio::io::BufReader;
use tokio::sync::Notify;
use tokio::sync::Semaphore;
use tokio::sync::mpsc;
#[derive(Clone)]
pub(crate) struct TestServer {
@@ -170,7 +174,7 @@ pub(crate) fn startup(
starts: Arc<AtomicUsize>,
first_started: Arc<Notify>,
release_first: Arc<Notify>,
server: TestServer,
server: impl InProcessTransportFactory + Clone + 'static,
) -> impl Fn(CancellationToken, Option<CancellationToken>) -> ManagedClientFuture {
move |cancellation, _sharing_declined| {
let sharing = sharing.clone();
@@ -236,6 +240,201 @@ pub(crate) fn startup(
}
}
struct LeaseFrame {
message: serde_json::Value,
replies: mpsc::UnboundedSender<serde_json::Value>,
}
impl LeaseFrame {
fn reply(&self, result: serde_json::Value) {
self.replies
.send(json!({"jsonrpc": "2.0", "id": self.message["id"], "result": result}))
.unwrap();
}
}
#[derive(Clone)]
struct LeaseTransport(mpsc::UnboundedSender<LeaseFrame>);
impl InProcessTransportFactory for LeaseTransport {
fn open(&self) -> BoxFuture<'static, std::io::Result<tokio::io::DuplexStream>> {
let (client, server) = tokio::io::duplex(4096);
let (reader, mut writer) = tokio::io::split(server);
let mut lines = BufReader::new(reader).lines();
let frames = self.0.clone();
let (replies, mut responses) = mpsc::unbounded_channel();
tokio::spawn(async move {
loop {
tokio::select! {
Some(response) = responses.recv() => {
writer.write_all(format!("{response}\n").as_bytes()).await.unwrap();
}
line = lines.next_line() => {
let Some(line) = line.unwrap() else { break };
let frame = LeaseFrame {
message: serde_json::from_str(&line).unwrap(),
replies: replies.clone(),
};
match frame.message["method"].as_str().unwrap() {
"initialize" => frame.reply(json!({
"protocolVersion": "2025-06-18", "capabilities": {"tools": {}},
"serverInfo": {"name": "test", "version": "1"}
})),
"notifications/initialized" => {}
"tools/list" => frame.reply(json!({"tools": []})),
_ => frames.send(frame).unwrap(),
}
}
}
}
});
futures::future::ready(Ok(client)).boxed()
}
}
struct LeasePeer {
pool: McpConnectionPool,
key: ConnectionKey,
transport: LeaseTransport,
frames: mpsc::UnboundedReceiver<LeaseFrame>,
}
impl LeasePeer {
fn new() -> Self {
let (tx, frames) = mpsc::unbounded_channel();
Self {
pool: McpConnectionPool::default(),
key: key(&context(), &config()),
transport: LeaseTransport(tx),
frames,
}
}
fn attach(&self) -> tokio::task::JoinHandle<Result<ManagedClient, StartupOutcomeError>> {
let pool = self.pool.clone();
let key = self.key.clone();
let transport = self.transport.clone();
tokio::spawn(async move {
let ready = Arc::<Notify>::default();
ready.notify_one();
let launch = startup(
ConnectionSharing::Negotiated,
Arc::default(),
Arc::default(),
ready,
transport,
);
pool.connect(
Arc::default(),
key,
Arc::new(Mutex::new(Some(
json!({"extensions": {LEASE_CAPABILITY: {}}}),
))),
/*tool_catalog_cache_context*/ None,
CancellationToken::new(),
launch,
)
.await
})
}
async fn next(&mut self, method: &str) -> LeaseFrame {
let frame = tokio::time::timeout(Duration::from_secs(2), self.frames.recv())
.await
.expect("expected a frame from the client")
.unwrap();
assert_eq!(frame.message["method"], method);
frame
}
async fn chat(&mut self) -> (ManagedClient, serde_json::Value) {
let attach = self.attach();
let open = self.next("codex/connection/open").await;
open.reply(json!({}));
(
attach.await.unwrap().unwrap(),
open.message["params"]["leaseId"].clone(),
)
}
}
#[tokio::test(start_paused = true)]
async fn releasing_lease_cancels_pending_call_before_release() {
let mut peer = LeasePeer::new();
let (first, first_id) = peer.chat().await;
let (other, other_id) = peer.chat().await;
let call = first.call_tool(
"wait".into(),
/*arguments*/ None,
|_| Ok(None),
/*timeout*/ None,
);
tokio::pin!(call);
let request = tokio::select! {
result = &mut call => panic!("call completed without a reply: {result:?}"),
frame = peer.next("tools/call") => frame,
};
let ConnectionSharing::Shared(lease) = &first.connection_sharing else {
panic!("expected shared lease")
};
lease.release();
// Leave the caller unpolled: the release must drive the cancellation write itself.
let cancel = peer.next("notifications/cancelled").await;
assert_eq!(cancel.message["params"]["requestId"], request.message["id"]);
let release = peer.next("codex/connection/release").await;
assert_eq!(release.message["params"]["leaseId"], first_id);
release.reply(json!({}));
assert!(call.await.is_err());
let (result, ()) = tokio::join!(
other.call_tool(
"test".into(),
/*arguments*/ None,
|_| Ok(None),
/*timeout*/ None,
),
async {
let request = peer.next("tools/call").await;
assert_eq!(request.message["params"]["_meta"][LEASE_META_KEY], other_id);
request.reply(json!({"content": []}));
}
);
result.unwrap();
}
#[tokio::test(start_paused = true)]
async fn dropping_lease_outside_runtime_releases_it() {
let mut peer = LeasePeer::new();
let (_other, _) = peer.chat().await;
let (chat, id) = peer.chat().await;
std::thread::spawn(move || drop(chat)).join().unwrap();
let release = peer.next("codex/connection/release").await;
assert_eq!(release.message["params"]["leaseId"], id);
release.reply(json!({}));
}
#[tokio::test(start_paused = true)]
async fn timed_out_open_is_released_after_the_server_acknowledges_it() {
let mut peer = LeasePeer::new();
let (_other, _) = peer.chat().await;
let attach = peer.attach();
let open = peer.next("codex/connection/open").await;
tokio::time::advance(LEASE_REQUEST_TIMEOUT).await;
assert!(attach.await.unwrap().is_err());
assert!(
tokio::time::timeout(Duration::from_secs(1), peer.frames.recv())
.await
.is_err()
);
open.reply(json!({}));
let release = peer.next("codex/connection/release").await;
assert_eq!(
release.message["params"]["leaseId"],
open.message["params"]["leaseId"]
);
release.reply(json!({}));
}
#[tokio::test]
async fn shared_startup_survives_first_chat_cancellation_and_rejects_released_bindings() {
let pool = McpConnectionPool::default();
codex-rs/codex-mcp/src/rmcp_client.rs
+11 -1
View File
@@ -40,6 +40,7 @@ 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::LEASE_CAPABILITY;
use crate::connection_pool::SharingNegotiation;
use crate::elicitation::ElicitationRequestManager;
use crate::executor_environment_http_client::ExecutorEnvironmentHttpClient;
@@ -1140,6 +1141,13 @@ async fn start_server_task(
};
let mut params =
mcp_initialize_request_params(client_elicitation_capability, client_mcp_extensions);
if sharing_negotiation.is_some() {
params
.capabilities
.extensions
.get_or_insert_default()
.insert(LEASE_CAPABILITY.to_owned(), Default::default());
}
if auth_changes.is_some() {
params
.capabilities
@@ -1344,7 +1352,8 @@ pub(crate) fn mcp_initialize_request_params(
}
let extensions = client_mcp_extensions
.iter()
.filter(|(id, _)| *id != MCP_NETWORK_POLICY_META_CAPABILITY)
// Only Core can open and release leases on a pooled connection.
.filter(|(id, _)| !matches!(*id, MCP_NETWORK_POLICY_META_CAPABILITY | LEASE_CAPABILITY))
.filter_map(|(id, settings)| {
settings
.as_object()
@@ -1662,6 +1671,7 @@ mod tests {
),
(OPENAI_FORM_EXTENSION_ID.to_string(), serde_json::json!({})),
(MCP_APP_UI_EXTENSION_ID.to_string(), app_ui.clone()),
(LEASE_CAPABILITY.to_string(), serde_json::json!({})),
]),
);
assert_eq!(
codex-rs/codex-mcp/src/rmcp_client/managed_client.rs
+36 -5
View File
@@ -59,6 +59,18 @@ impl ManagedClient {
}
}
async fn with_lease_request<T>(&self, operation: impl Future<Output = Result<T>>) -> Result<T> {
// Keep one copy of the request future as it passes through both lease scopes.
let operation = std::pin::pin!(operation);
self.with_active_shared_lease(async {
match self.connection_sharing.server_lease() {
Some(lease) => lease.requests.scope(operation).await,
None => operation.await,
}
})
.await?
}
pub(crate) async fn call_tool(
&self,
name: String,
@@ -66,7 +78,20 @@ impl ManagedClient {
meta_for_session: impl Fn(&ServerCapabilities) -> Result<Option<Value>> + Send + Sync,
timeout: Option<Duration>,
) -> Result<rmcp::model::CallToolResult> {
self.with_active_shared_lease(async {
let meta_for_session = |caps: &ServerCapabilities| {
let mut meta = meta_for_session(caps)?;
if let Some(lease) = self.connection_sharing.server_lease() {
meta.get_or_insert_with(|| serde_json::json!({}))
.as_object_mut()
.ok_or_else(|| anyhow!("MCP tool metadata must be an object"))?
.insert(
crate::connection_pool::LEASE_META_KEY.to_owned(),
serde_json::json!(lease.id),
);
}
Ok(meta)
};
self.with_lease_request(async {
if let Some(send_elicitation) = &self.shared_elicitation {
self.client
.call_tool_with_elicitation(
@@ -83,7 +108,7 @@ impl ManagedClient {
.await
}
})
.await?
.await
}
/// Lists resources within this chat's lease lifetime.
@@ -109,10 +134,16 @@ impl ManagedClient {
/// Reads a resource using this chat's input authority and lease lifetime.
pub(crate) async fn read_resource(
&self,
params: rmcp::model::ReadResourceRequestParams,
mut params: rmcp::model::ReadResourceRequestParams,
timeout: Option<Duration>,
) -> Result<rmcp::model::ReadResourceResult> {
self.with_active_shared_lease(async {
if let Some(lease) = self.connection_sharing.server_lease() {
params.meta.get_or_insert_default().insert(
crate::connection_pool::LEASE_META_KEY.to_owned(),
serde_json::json!(lease.id),
);
}
self.with_lease_request(async {
if let Some(send_elicitation) = &self.shared_elicitation {
self.client
.read_resource_with_elicitation(params, timeout, Arc::clone(send_elicitation))
@@ -121,7 +152,7 @@ impl ManagedClient {
self.client.read_resource(params, timeout).await
}
})
.await?
.await
}
pub(crate) async fn listed_catalog(&self) -> (Vec<ToolInfo>, Option<Arc<serde_json::Value>>) {
codex-rs/rmcp-client/src/bin/test_stdio_server.rs
+61
View File
@@ -62,6 +62,23 @@ const DYNAMIC_SERVER_METADATA_ENV: &str = "MCP_TEST_DYNAMIC_SERVER_METADATA";
const INITIALIZE_BARRIER_FILE_ENV: &str = "MCP_TEST_INITIALIZE_BARRIER_FILE";
const SERVER_INSTRUCTIONS_ENV: &str = "MCP_TEST_SERVER_INSTRUCTIONS";
const LEASE_EVENTS_FILE_ENV: &str = "MCP_TEST_LEASE_EVENTS_FILE";
/// Record only the lease protocol for the app-server unload regression.
fn record_lease_event(method: &str, id: &serde_json::Value) -> Result<(), McpError> {
use std::io::Write;
if let Some(path) = std::env::var_os(LEASE_EVENTS_FILE_ENV) {
let event = serde_json::json!([method, id]);
std::fs::OpenOptions::new()
.create(true)
.append(true)
.open(path)
.and_then(|mut file| writeln!(file, "{event}"))
.map_err(|error| McpError::internal_error(error.to_string(), None))?;
}
Ok(())
}
fn dynamic_server_process_label() -> Option<String> {
std::env::var_os(DYNAMIC_SERVER_METADATA_ENV)
.is_some()
@@ -548,6 +565,17 @@ impl ServerHandler for TestToolServer {
JsonObject::new(),
)]));
if std::env::var_os(LEASE_EVENTS_FILE_ENV).is_some() {
capabilities
.extensions
.get_or_insert_default()
.insert("codex/connection-lease".to_owned(), JsonObject::new());
} else if std::env::var_os("MCP_TEST_EXPERIMENTAL_LEASE_ONLY").is_some() {
capabilities
.experimental
.get_or_insert_default()
.insert("codex/connection-lease".to_owned(), JsonObject::new());
}
let server_info = ServerConfig::new(capabilities);
let server_info = match dynamic_server_process_label() {
Some(process_label) => server_info
@@ -689,11 +717,44 @@ impl ServerHandler for TestToolServer {
}
}
async fn on_custom_request(
&self,
request: rmcp::model::CustomRequest,
context: rmcp::service::RequestContext<rmcp::service::RoleServer>,
) -> Result<rmcp::model::CustomResult, McpError> {
if std::env::var_os(LEASE_EVENTS_FILE_ENV).is_none()
|| !matches!(
request.method.as_str(),
"codex/connection/open" | "codex/connection/release"
)
|| !context
.client_capabilities()
.and_then(|caps| caps.extensions)
.is_some_and(|caps| caps.contains_key("codex/connection-lease"))
{
return Err(McpError::new(
rmcp::model::ErrorCode::METHOD_NOT_FOUND,
request.method,
None,
));
}
let id = request
.params
.as_ref()
.and_then(|params| params.get("leaseId"))
.ok_or_else(|| McpError::invalid_params("missing lease ID", None))?;
record_lease_event(&request.method, id)?;
Ok(rmcp::model::CustomResult(serde_json::json!({})))
}
async fn call_tool(
&self,
request: CallToolRequestParams,
context: rmcp::service::RequestContext<rmcp::service::RoleServer>,
) -> Result<rmcp::model::CallToolResponse, McpError> {
if let Some(lease) = context.meta.get("codex/connectionLease") {
record_lease_event("tools/call", lease)?;
}
match request.name.as_ref() {
"search_tools" if std::env::var_os("MCP_TEST_DISCOVER_TOOLS").is_some() => {
self.tools_discovered.store(true, Ordering::Release);
codex-rs/rmcp-client/src/lib.rs
+1
View File
@@ -96,6 +96,7 @@ pub use perform_oauth_login::perform_oauth_login_return_url;
pub use perform_oauth_login::perform_oauth_login_silent;
pub use perform_oauth_login::perform_oauth_login_with_callback_input;
pub use protocol_mode::McpProtocolMode;
pub use request_cancellation::RequestScope;
pub use rmcp::model::ElicitationAction;
pub use rmcp_client::CancellableEventStreamRequest;
pub use rmcp_client::Elicitation;
codex-rs/rmcp-client/src/request_cancellation.rs
+84 -8
View File
@@ -1,5 +1,7 @@
//! Cancel remote MCP work when its local response future is dropped.
use std::future::Future;
use rmcp::Peer;
use rmcp::RoleClient;
use rmcp::model::CancelledNotificationParam;
@@ -8,6 +10,75 @@ use rmcp::model::RequestId;
use rmcp::model::ServerResult;
use rmcp::service::PeerRequestOptions;
use rmcp::service::ServiceError;
use tokio::runtime::Handle;
use tokio_util::sync::CancellationToken;
use tokio_util::task::TaskTracker;
tokio::task_local! {
static REQUEST_SCOPE: RequestScope;
}
/// Tracks cancellation writes for one chat's requests on a shared stdio transport.
#[derive(Clone)]
pub struct RequestScope {
runtime: Handle,
cancellation: CancellationToken,
pending: TaskTracker,
}
impl RequestScope {
pub fn new(runtime: Handle, cancellation: CancellationToken) -> Self {
Self {
runtime,
cancellation,
pending: TaskTracker::new(),
}
}
async fn send(
&self,
peer: Peer<RoleClient>,
request: ClientRequest,
options: PeerRequestOptions,
) -> Result<ServerResult, ServiceError> {
let cancellation = self.cancellation.child_token();
let _cancel_on_drop = cancellation.clone().drop_guard();
// The transport must make progress even if the caller stops polling its future.
// Register before the task checks cancellation so a racing release observes it.
let send = async move {
let handle = tokio::select! {
biased;
() = cancellation.cancelled() => return Err(ServiceError::Cancelled { reason: None }),
result = peer.send_request_with_option(request, options) => result?,
};
let id = handle.id.clone();
tokio::select! {
biased;
result = handle.await_response() => result,
() = cancellation.cancelled() => {
notify_cancelled(peer, id).await;
Err(ServiceError::Cancelled { reason: None })
}
}
};
self.pending
.spawn_on(send, &self.runtime)
.await
.map_err(|_| ServiceError::TransportClosed)?
}
/// Include every wire request made by this operation, including MCP continuations.
pub async fn scope<F: Future>(&self, operation: F) -> F::Output {
REQUEST_SCOPE.scope(self.clone(), operation).await
}
/// Stop admitting requests and wait for the transport to finish writing their cancellations.
pub async fn cancel_and_wait(&self) {
self.cancellation.cancel();
self.pending.close();
self.pending.wait().await;
}
}
/// Whether dropping a caller's future should cancel its request on the peer.
#[derive(Clone, Copy)]
@@ -24,6 +95,9 @@ impl RequestCancellation {
request: ClientRequest,
options: PeerRequestOptions,
) -> Result<ServerResult, ServiceError> {
if let Ok(scope) = REQUEST_SCOPE.try_with(Clone::clone) {
return scope.send(peer.clone(), request, options).await;
}
let handle = peer.send_request_with_option(request, options).await?;
let mut cancellation = CancelOnDrop(match self {
Self::Disabled => None,
@@ -42,14 +116,16 @@ impl Drop for CancelOnDrop {
if let Some((peer, id)) = self.0.take()
&& let Ok(runtime) = tokio::runtime::Handle::try_current()
{
runtime.spawn(async move {
let _ = peer
.notify_cancelled(CancelledNotificationParam::new(
Some(id),
Some("MCP response future dropped".into()),
))
.await;
});
runtime.spawn(notify_cancelled(peer, id));
}
}
}
async fn notify_cancelled(peer: Peer<RoleClient>, id: RequestId) {
let _ = peer
.notify_cancelled(CancelledNotificationParam::new(
Some(id),
Some("MCP request cancelled by caller".into()),
))
.await;
}