diff --git a/codex-rs/codex-mcp/src/connection_manager_tests.rs b/codex-rs/codex-mcp/src/connection_manager_tests.rs index 8230ff0334..63a9cc3675 100644 --- a/codex-rs/codex-mcp/src/connection_manager_tests.rs +++ b/codex-rs/codex-mcp/src/connection_manager_tests.rs @@ -883,6 +883,94 @@ async fn shared_elicitation_router_targets_the_exact_pending_request() { assert_eq!(outstanding.load(std::sync::atomic::Ordering::SeqCst), 0); } +#[tokio::test] +async fn cancelled_elicitation_is_removed_without_affecting_other_pending_requests() { + let router = ElicitationRequestRouter::default(); + let manager = ElicitationRequestManager::new( + AskForApproval::OnRequest, + PermissionProfile::default(), + /*reviewer*/ None, + /*lifecycle*/ None, + router.clone(), + ); + let (tx_event, rx_event) = async_channel::bounded(2); + let sender = manager.make_sender("server".to_string(), Some(tx_event)); + let elicitation = + codex_rmcp_client::Elicitation::Mcp(ElicitRequestParams::FormElicitationParams { + meta: None, + message: "Confirm?".to_string(), + requested_schema: + rmcp::model::ElicitationSchema::builder() + .required_property( + "answer", + rmcp::model::PrimitiveSchemaDefinition::String( + rmcp::model::StringSchema::new(), + ), + ) + .build() + .expect("schema should build"), + }); + + let cancelled = tokio::spawn(sender(NumberOrString::Number(1), elicitation.clone())); + let EventMsg::ElicitationRequest(cancelled_request) = + rx_event.recv().await.expect("cancelled request event").msg + else { + panic!("expected elicitation request"); + }; + let pending = tokio::spawn(sender(NumberOrString::Number(2), elicitation)); + let EventMsg::ElicitationRequest(pending_request) = + rx_event.recv().await.expect("pending request event").msg + else { + panic!("expected elicitation request"); + }; + let ( + codex_protocol::mcp::RequestId::String(cancelled_id), + codex_protocol::mcp::RequestId::String(pending_id), + ) = (cancelled_request.id, pending_request.id) + else { + panic!("expected Codex-owned string request IDs"); + }; + + cancelled.abort(); + assert!( + cancelled + .await + .expect_err("cancelled request should be aborted") + .is_cancelled() + ); + + let response = ElicitationResponse { + action: ElicitationAction::Accept, + content: Some(serde_json::json!({"answer": "yes"})), + meta: None, + }; + let error = router + .resolve( + "server".to_string(), + NumberOrString::String(cancelled_id.into()), + response.clone(), + ) + .await + .expect_err("cancelled request should be removed immediately"); + assert_eq!(error.to_string(), "elicitation request not found"); + + router + .resolve( + "server".to_string(), + NumberOrString::String(pending_id.into()), + response.clone(), + ) + .await + .expect("another pending request should remain routable"); + assert_eq!( + pending + .await + .expect("pending request task") + .expect("pending request response"), + response + ); +} + #[test] fn test_normalize_tools_short_non_duplicated_names() { let tools = vec![ diff --git a/codex-rs/codex-mcp/src/elicitation.rs b/codex-rs/codex-mcp/src/elicitation.rs index c61fd5c0d9..e8b542393e 100644 --- a/codex-rs/codex-mcp/src/elicitation.rs +++ b/codex-rs/codex-mcp/src/elicitation.rs @@ -32,7 +32,6 @@ use futures::future::BoxFuture; use futures::future::FutureExt; use rmcp::model::ElicitationAction; use rmcp::model::RequestId; -use tokio::sync::Mutex; use tokio::sync::oneshot; static NEXT_ELICITATION_REQUEST_ID: AtomicU64 = AtomicU64::new(0); @@ -87,10 +86,27 @@ struct ActiveElicitation { /// the same server request ID without colliding. #[derive(Clone, Default)] pub(crate) struct ElicitationRequestRouter { - requests: Arc>, + requests: Arc>, auto_deny: Arc, } +struct PendingElicitationRequest { + router: ElicitationRequestRouter, + key: (String, RequestId), +} + +impl Drop for PendingElicitationRequest { + fn drop(&mut self) { + let responder = self + .router + .requests + .lock() + .ok() + .and_then(|mut requests| requests.remove(&self.key)); + drop(responder); + } +} + impl ElicitationRequestRouter { pub(crate) fn auto_deny(&self) -> bool { self.auto_deny.load(Ordering::Relaxed) @@ -106,11 +122,13 @@ impl ElicitationRequestRouter { id: RequestId, response: ElicitationResponse, ) -> Result<()> { - self.requests + let responder = self + .requests .lock() - .await + .map_err(|_| anyhow!("elicitation request router unavailable"))? .remove(&(server_name, id)) - .ok_or_else(|| anyhow!("elicitation request not found"))? + .ok_or_else(|| anyhow!("elicitation request not found"))?; + responder .send(response) .map_err(|e| anyhow!("failed to send elicitation response: {e:?}")) } @@ -294,10 +312,16 @@ impl ElicitationRequestManager { }; let (tx, rx) = oneshot::channel(); let _active_elicitation = lifecycle.as_ref().map(ElicitationLifecycle::start); - { - let mut lock = router.requests.lock().await; - lock.insert((server_name.clone(), routed_request_id), tx); - } + let request_key = (server_name.clone(), routed_request_id); + router + .requests + .lock() + .map_err(|_| anyhow!("elicitation request router unavailable"))? + .insert(request_key.clone(), tx); + let _pending_request = PendingElicitationRequest { + router: router.clone(), + key: request_key, + }; let _ = tx_event .send(Event { id: "mcp_elicitation_request".to_string(),