Keep MCP elicitation routable across runtime refreshes (#30127)

## Why

An MCP tool call can still be waiting for an elicitation response when
an environment update replaces the thread's MCP runtime.

Before this change:

```text
runtime A starts a tool call and asks the user
environment becomes ready, so runtime B is published
client answers the prompt through runtime B
runtime B cannot find runtime A's pending responder
```

The response is lost and the original tool call stays blocked.

## What changed

All MCP runtimes for one thread now share a small elicitation router:

```text
runtime A ---\
               shared router: response token -> exact pending responder
runtime B ---/
```

When Codex surfaces an MCP elicitation, it assigns a unique opaque
response token. The router records which pending request owns that
token. A replacement runtime reuses the same router, so the latest
runtime can deliver a response to a request started by the previous
runtime.

The Codex-owned token also prevents two runtime connections that reuse
the same MCP server request ID from receiving each other's responses.

This does not retain or search old MCP managers. Only the pending
responder map is shared.

## Covered scenario

The integration test exercises the complete failure mode:

1. A thread starts while its selected environment is still unavailable.
2. A configured MCP server starts a tool call and asks the client for
input.
3. The environment becomes ready, causing Codex to publish a replacement
MCP runtime.
4. The client answers the original prompt after the replacement.
5. The original tool call receives that answer and completes.

A focused routing test also creates two runtimes with the same server
request ID and verifies that each response reaches the exact request
that emitted its token.

## Scope

This PR changes only elicitation response routing across MCP runtime
replacement. It does not change when runtimes are rebuilt, which
environments contribute MCP configuration, or how environment
availability is detected.
This commit is contained in:
jif authored and GitHub committed 2026-06-26 01:28:14 +00:00
1 parent 723b23efd0
commit fb8598df3f
10 files changed
+398 -26

No files matched your search

+49 -21
View File
@@ -8,6 +8,8 @@
use std::collections::HashMap;
use std::sync::Arc;
use std::sync::Mutex as StdMutex;
use std::sync::atomic::AtomicU64;
use std::sync::atomic::Ordering;
use crate::mcp::McpPermissionPromptAutoApproveContext;
use crate::mcp::mcp_permission_prompt_is_auto_approved;
@@ -32,6 +34,8 @@ use rmcp::model::RequestId;
use tokio::sync::Mutex;
use tokio::sync::oneshot;
static NEXT_ELICITATION_REQUEST_ID: AtomicU64 = AtomicU64::new(0);
#[derive(Debug, Clone)]
pub struct ElicitationReviewRequest {
pub server_name: String,
@@ -48,9 +52,36 @@ pub trait ElicitationReviewer: Send + Sync {
pub type ElicitationReviewerHandle = Arc<dyn ElicitationReviewer>;
/// Routes model-visible elicitation response tokens to their exact pending responders.
///
/// One router is shared by every MCP runtime created for a thread. The public response token is
/// generated by Codex rather than copied from the MCP connection, so separate runtimes may reuse
/// the same server request ID without colliding.
#[derive(Clone, Default)]
pub struct ElicitationRequestRouter {
requests: Arc<Mutex<ResponderMap>>,
}
impl ElicitationRequestRouter {
async fn resolve(
&self,
server_name: String,
id: RequestId,
response: ElicitationResponse,
) -> Result<()> {
self.requests
.lock()
.await
.remove(&(server_name, id))
.ok_or_else(|| anyhow!("elicitation request not found"))?
.send(response)
.map_err(|e| anyhow!("failed to send elicitation response: {e:?}"))
}
}
#[derive(Clone)]
pub(crate) struct ElicitationRequestManager {
requests: Arc<Mutex<ResponderMap>>,
router: ElicitationRequestRouter,
pub(crate) approval_policy: Arc<StdMutex<AskForApproval>>,
pub(crate) permission_profile: Arc<StdMutex<PermissionProfile>>,
auto_deny: Arc<StdMutex<bool>>,
@@ -62,9 +93,10 @@ impl ElicitationRequestManager {
approval_policy: AskForApproval,
permission_profile: PermissionProfile,
reviewer: Option<ElicitationReviewerHandle>,
router: ElicitationRequestRouter,
) -> Self {
Self {
requests: Arc::new(Mutex::new(HashMap::new())),
router,
approval_policy: Arc::new(StdMutex::new(approval_policy)),
permission_profile: Arc::new(StdMutex::new(permission_profile)),
auto_deny: Arc::new(StdMutex::new(false)),
@@ -91,13 +123,11 @@ impl ElicitationRequestManager {
id: RequestId,
response: ElicitationResponse,
) -> Result<()> {
self.requests
.lock()
.await
.remove(&(server_name, id))
.ok_or_else(|| anyhow!("elicitation request not found"))?
.send(response)
.map_err(|e| anyhow!("failed to send elicitation response: {e:?}"))
self.router.resolve(server_name, id, response).await
}
pub(crate) fn router(&self) -> ElicitationRequestRouter {
self.router.clone()
}
pub(crate) fn make_sender(
@@ -105,13 +135,13 @@ impl ElicitationRequestManager {
server_name: String,
tx_event: Sender<Event>,
) -> SendElicitation {
let elicitation_requests = self.requests.clone();
let router = self.router.clone();
let approval_policy = self.approval_policy.clone();
let permission_profile = self.permission_profile.clone();
let auto_deny = self.auto_deny.clone();
let reviewer = self.reviewer.clone();
Box::new(move |id, elicitation| {
let elicitation_requests = elicitation_requests.clone();
let router = router.clone();
let tx_event = tx_event.clone();
let server_name = server_name.clone();
let approval_policy = approval_policy.clone();
@@ -171,6 +201,11 @@ impl ElicitationRequestManager {
}
}
let public_request_id = format!(
"codex-mcp-elicitation-{}",
NEXT_ELICITATION_REQUEST_ID.fetch_add(1, Ordering::Relaxed)
);
let routed_request_id = RequestId::String(public_request_id.clone().into());
let request = match elicitation {
Elicitation::Mcp(
rmcp::model::CreateElicitationRequestParams::FormElicitationParams {
@@ -215,8 +250,8 @@ impl ElicitationRequestManager {
};
let (tx, rx) = oneshot::channel();
{
let mut lock = elicitation_requests.lock().await;
lock.insert((server_name.clone(), id.clone()), tx);
let mut lock = router.requests.lock().await;
lock.insert((server_name.clone(), routed_request_id), tx);
}
let _ = tx_event
.send(Event {
@@ -224,14 +259,7 @@ impl ElicitationRequestManager {
msg: EventMsg::ElicitationRequest(ElicitationRequestEvent {
turn_id: None,
server_name,
id: match id.clone() {
rmcp::model::NumberOrString::String(value) => {
ProtocolRequestId::String(value.to_string())
}
rmcp::model::NumberOrString::Number(value) => {
ProtocolRequestId::Integer(value)
}
},
id: ProtocolRequestId::String(public_request_id),
request,
}),
})