mirror of
https://github.com/pchuan98/codex.git
synced 2026-07-01 00:31:56 +08:00
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:
@@ -18,6 +18,7 @@ use crate::codex_apps_cache::CodexAppsToolsCache;
|
||||
use crate::codex_apps_cache::CodexAppsToolsCacheKey;
|
||||
use crate::codex_apps_cache::CodexAppsToolsFetchSource;
|
||||
use crate::elicitation::ElicitationRequestManager;
|
||||
use crate::elicitation::ElicitationRequestRouter;
|
||||
use crate::elicitation::ElicitationReviewerHandle;
|
||||
use crate::mcp::CODEX_APPS_MCP_SERVER_NAME;
|
||||
use crate::mcp::ToolPluginProvenance;
|
||||
@@ -141,6 +142,7 @@ impl McpConnectionManager {
|
||||
tool_plugin_provenance: ToolPluginProvenance,
|
||||
auth: Option<&CodexAuth>,
|
||||
elicitation_reviewer: Option<ElicitationReviewerHandle>,
|
||||
elicitation_router: ElicitationRequestRouter,
|
||||
) -> Self {
|
||||
let mut required_servers = mcp_servers
|
||||
.iter()
|
||||
@@ -155,6 +157,7 @@ impl McpConnectionManager {
|
||||
approval_policy.value(),
|
||||
initial_permission_profile,
|
||||
elicitation_reviewer,
|
||||
elicitation_router,
|
||||
);
|
||||
let tool_plugin_provenance = Arc::new(tool_plugin_provenance);
|
||||
let startup_submit_id = submit_id.clone();
|
||||
@@ -352,6 +355,7 @@ impl McpConnectionManager {
|
||||
approval_policy.value(),
|
||||
permission_profile.clone(),
|
||||
/*reviewer*/ None,
|
||||
ElicitationRequestRouter::default(),
|
||||
),
|
||||
startup_cancellation_token: CancellationToken::new(),
|
||||
}
|
||||
@@ -444,6 +448,10 @@ impl McpConnectionManager {
|
||||
self.elicitation_requests.set_auto_deny(auto_deny);
|
||||
}
|
||||
|
||||
pub fn elicitation_router(&self) -> ElicitationRequestRouter {
|
||||
self.elicitation_requests.router()
|
||||
}
|
||||
|
||||
pub async fn resolve_elicitation(
|
||||
&self,
|
||||
server_name: String,
|
||||
|
||||
@@ -3,6 +3,7 @@ use crate::codex_apps_cache::CodexAppsToolsCache;
|
||||
use crate::codex_apps_cache::CodexAppsToolsCacheContext;
|
||||
use crate::declared_openai_file_input_param_names;
|
||||
use crate::elicitation::ElicitationRequestManager;
|
||||
use crate::elicitation::ElicitationRequestRouter;
|
||||
use crate::elicitation::elicitation_is_rejected_by_policy;
|
||||
use crate::rmcp_client::AsyncManagedClient;
|
||||
use crate::rmcp_client::ManagedClient;
|
||||
@@ -274,6 +275,7 @@ async fn disabled_permissions_auto_accept_elicitation_with_empty_form_schema() {
|
||||
AskForApproval::Never,
|
||||
PermissionProfile::Disabled,
|
||||
/*reviewer*/ None,
|
||||
ElicitationRequestRouter::default(),
|
||||
);
|
||||
let (tx_event, _rx_event) = async_channel::bounded(1);
|
||||
let sender = manager.make_sender("server".to_string(), tx_event);
|
||||
@@ -309,6 +311,7 @@ async fn disabled_permissions_do_not_auto_accept_elicitation_with_requested_fiel
|
||||
AskForApproval::Never,
|
||||
PermissionProfile::Disabled,
|
||||
/*reviewer*/ None,
|
||||
ElicitationRequestRouter::default(),
|
||||
);
|
||||
let (tx_event, _rx_event) = async_channel::bounded(1);
|
||||
let sender = manager.make_sender("server".to_string(), tx_event);
|
||||
@@ -342,6 +345,100 @@ async fn disabled_permissions_do_not_auto_accept_elicitation_with_requested_fiel
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn shared_elicitation_router_targets_the_exact_pending_request() {
|
||||
let router = ElicitationRequestRouter::default();
|
||||
let manager_a = ElicitationRequestManager::new(
|
||||
AskForApproval::OnRequest,
|
||||
PermissionProfile::default(),
|
||||
/*reviewer*/ None,
|
||||
router.clone(),
|
||||
);
|
||||
let manager_b = ElicitationRequestManager::new(
|
||||
AskForApproval::OnRequest,
|
||||
PermissionProfile::default(),
|
||||
/*reviewer*/ None,
|
||||
router,
|
||||
);
|
||||
let (tx_event, rx_event) = async_channel::bounded(2);
|
||||
let sender_a = manager_a.make_sender("server".to_string(), tx_event.clone());
|
||||
let sender_b = manager_b.make_sender("server".to_string(), tx_event);
|
||||
let elicitation = codex_rmcp_client::Elicitation::Mcp(
|
||||
CreateElicitationRequestParams::FormElicitationParams {
|
||||
meta: None,
|
||||
message: "Which runtime?".to_string(),
|
||||
requested_schema: rmcp::model::ElicitationSchema::builder()
|
||||
.required_property(
|
||||
"runtime",
|
||||
rmcp::model::PrimitiveSchema::String(rmcp::model::StringSchema::new()),
|
||||
)
|
||||
.build()
|
||||
.expect("schema should build"),
|
||||
},
|
||||
);
|
||||
|
||||
let pending_a = tokio::spawn(sender_a(NumberOrString::Number(1), elicitation.clone()));
|
||||
let EventMsg::ElicitationRequest(request_a) = rx_event.recv().await.expect("request A").msg
|
||||
else {
|
||||
panic!("expected elicitation request");
|
||||
};
|
||||
let pending_b = tokio::spawn(sender_b(NumberOrString::Number(1), elicitation));
|
||||
let EventMsg::ElicitationRequest(request_b) = rx_event.recv().await.expect("request B").msg
|
||||
else {
|
||||
panic!("expected elicitation request");
|
||||
};
|
||||
let (
|
||||
codex_protocol::mcp::RequestId::String(request_a_id),
|
||||
codex_protocol::mcp::RequestId::String(request_b_id),
|
||||
) = (request_a.id, request_b.id)
|
||||
else {
|
||||
panic!("expected Codex-owned string request IDs");
|
||||
};
|
||||
assert_ne!(request_a_id, request_b_id);
|
||||
|
||||
let response_a = ElicitationResponse {
|
||||
action: ElicitationAction::Accept,
|
||||
content: Some(serde_json::json!({"runtime": "a"})),
|
||||
meta: None,
|
||||
};
|
||||
manager_b
|
||||
.resolve(
|
||||
"server".to_string(),
|
||||
NumberOrString::String(request_a_id.into()),
|
||||
response_a.clone(),
|
||||
)
|
||||
.await
|
||||
.expect("runtime B should route a response to runtime A");
|
||||
let response_b = ElicitationResponse {
|
||||
action: ElicitationAction::Accept,
|
||||
content: Some(serde_json::json!({"runtime": "b"})),
|
||||
meta: None,
|
||||
};
|
||||
manager_a
|
||||
.resolve(
|
||||
"server".to_string(),
|
||||
NumberOrString::String(request_b_id.into()),
|
||||
response_b.clone(),
|
||||
)
|
||||
.await
|
||||
.expect("runtime A should route a response to runtime B");
|
||||
|
||||
assert_eq!(
|
||||
pending_a
|
||||
.await
|
||||
.expect("request A task")
|
||||
.expect("request A response"),
|
||||
response_a
|
||||
);
|
||||
assert_eq!(
|
||||
pending_b
|
||||
.await
|
||||
.expect("request B task")
|
||||
.expect("request B response"),
|
||||
response_b
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_normalize_tools_short_non_duplicated_names() {
|
||||
let tools = vec![
|
||||
@@ -1158,6 +1255,7 @@ async fn no_local_runtime_fails_local_stdio_but_keeps_local_http_server() {
|
||||
ToolPluginProvenance::default(),
|
||||
/*auth*/ None,
|
||||
/*elicitation_reviewer*/ None,
|
||||
ElicitationRequestRouter::default(),
|
||||
)
|
||||
.await;
|
||||
|
||||
|
||||
@@ -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,
|
||||
}),
|
||||
})
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
pub use connection_manager::McpConnectionManager;
|
||||
pub use connection_manager::tool_is_model_visible;
|
||||
pub use elicitation::ElicitationRequestRouter;
|
||||
pub use elicitation::ElicitationReviewRequest;
|
||||
pub use elicitation::ElicitationReviewer;
|
||||
pub use elicitation::ElicitationReviewerHandle;
|
||||
|
||||
@@ -339,6 +339,7 @@ pub async fn read_mcp_resource(
|
||||
tool_plugin_provenance(config),
|
||||
auth,
|
||||
/*elicitation_reviewer*/ None,
|
||||
crate::elicitation::ElicitationRequestRouter::default(),
|
||||
)
|
||||
.await;
|
||||
|
||||
@@ -415,6 +416,7 @@ pub async fn collect_mcp_server_status_snapshot_with_detail(
|
||||
tool_plugin_provenance,
|
||||
auth,
|
||||
/*elicitation_reviewer*/ None,
|
||||
crate::elicitation::ElicitationRequestRouter::default(),
|
||||
)
|
||||
.await;
|
||||
|
||||
|
||||
Reference in New Issue
Block a user