codex/codex-rs/codex-mcp/src/elicitation.rs

//! MCP elicitation request tracking and policy handling.
//!
//! RMCP clients call into this module when a server asks Codex to elicit data
//! from the user. It decides whether the request can be automatically accepted,
//! must be declined by policy, or should be surfaced as a Codex protocol event
//! and later resolved through the stored responder.
//! Explicit MCP elicitation requests can be surfaced by root and non-root agents.

use std::collections::HashMap;
use std::sync::Arc;
use std::sync::Mutex as StdMutex;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::AtomicU64;
use std::sync::atomic::Ordering;

use crate::McpConfig;
use crate::mcp::McpPermissionPromptAutoApproveContext;
use crate::mcp::mcp_permission_prompt_is_auto_approved;
use anyhow::Context;
use anyhow::Result;
use anyhow::anyhow;
use async_channel::Sender;
use codex_protocol::approvals::ElicitationAbandonedEvent;
use codex_protocol::approvals::ElicitationRequest;
use codex_protocol::approvals::ElicitationRequestEvent;
use codex_protocol::mcp::ClientMcpExtensions;
use codex_protocol::mcp::OPENAI_ELICITATION_EXTENSION_ID;
use codex_protocol::mcp::RequestId as ProtocolRequestId;
use codex_protocol::mcp_approval_meta::APPROVAL_KIND_KEY;
use codex_protocol::mcp_approval_meta::APPROVAL_KIND_TOOL_SUGGESTION;
use codex_protocol::mcp_approval_meta::APPROVALS_REVIEWER_KEY;
use codex_protocol::mcp_approval_meta::STRICT_AUTO_REVIEW_KEY;
use codex_protocol::protocol::AskForApproval;
use codex_protocol::protocol::Event;
use codex_protocol::protocol::EventMsg;
use codex_rmcp_client::Elicitation;
use codex_rmcp_client::ElicitationResponse;
use codex_rmcp_client::SendElicitation;
use futures::future::BoxFuture;
use futures::future::FutureExt;
use rmcp::model::ElicitationAction;
use rmcp::model::RequestId;
use serde_json::Value;
use tokio::sync::oneshot;

static NEXT_ELICITATION_REQUEST_ID: AtomicU64 = AtomicU64::new(0);

const STRICT_AUTO_REVIEW_DECLINE_MESSAGE: &str = "Automated review of this operation failed. Do not proceed without asking the user for explicit approval.";

#[path = "user_verification_elicitation.rs"]
mod user_verification_elicitation;

#[derive(Debug, Clone)]
pub struct ElicitationReviewRequest {
    pub server_name: String,
    pub request_id: RequestId,
    pub elicitation: Elicitation,
}

pub trait ElicitationReviewer: Send + Sync {
    fn review(
        &self,
        request: ElicitationReviewRequest,
    ) -> BoxFuture<'static, Result<Option<ElicitationResponse>>>;
}

pub type ElicitationReviewerHandle = Arc<dyn ElicitationReviewer>;

/// Holds an owner-provided registration while an MCP elicitation is waiting for a response.
#[derive(Clone)]
pub struct ElicitationLifecycle {
    register: Arc<dyn Fn() -> Box<dyn Send + Sync> + Send + Sync>,
}

impl ElicitationLifecycle {
    pub fn new<T>(register: impl Fn() -> T + Send + Sync + 'static) -> Self
    where
        T: Send + Sync + 'static,
    {
        Self {
            register: Arc::new(move || Box::new(register())),
        }
    }

    fn start(&self) -> ActiveElicitation {
        ActiveElicitation {
            _registration: (self.register)(),
        }
    }
}

struct ActiveElicitation {
    _registration: Box<dyn Send + Sync>,
}

/// 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(crate) struct ElicitationRequestRouter {
    requests: Arc<StdMutex<PendingResponders>>,
    terminal_actions_finished: Arc<tokio::sync::Notify>,
    auto_deny: Arc<AtomicBool>,
    full_access_form_input_enabled: Arc<AtomicBool>,
}

// Removes and abandons a published request when its native waiter is dropped.
struct ElicitationRequestGuard {
    router: ElicitationRequestRouter,
    key: (String, RequestId),
}

#[derive(Default, PartialEq)]
enum RouterPhase {
    #[default]
    Open,
    Closing,
    Closed,
}

#[derive(Default)]
struct PendingResponders {
    pending: ResponderMap,
    phase: RouterPhase,
    terminal_actions: usize,
}

// Owns response delivery and the original source for abandonment notifications.
struct ElicitationResponder {
    response: oneshot::Sender<ElicitationResponse>,
    abandonment_events: Option<Sender<Event>>,
}

// Removal elects a winner, but shutdown must also wait for its outside-lock delivery.
struct ElicitationDeliveryGuard(ElicitationRequestRouter);

impl Drop for ElicitationDeliveryGuard {
    fn drop(&mut self) {
        let mut state = self
            .0
            .requests
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        state.terminal_actions -= 1;
        if state.phase == RouterPhase::Closing && state.terminal_actions == 0 {
            state.phase = RouterPhase::Closed;
        }
        drop(state);
        self.0.terminal_actions_finished.notify_waiters();
    }
}

impl Drop for ElicitationRequestGuard {
    fn drop(&mut self) {
        if let Ok(Some((responder, _terminal_action))) = self.router.take_responder(&self.key) {
            emit_abandonment(&self.key, responder.abandonment_events.as_ref());
            drop(responder);
        }
    }
}

fn emit_abandonment(key: &(String, RequestId), events: Option<&Sender<Event>>) {
    if let Some(events) = events {
        let event = Event {
            id: "mcp_elicitation_abandoned".to_string(),
            msg: EventMsg::ElicitationAbandoned(ElicitationAbandonedEvent {
                server_name: key.0.clone(),
                id: match &key.1 {
                    RequestId::String(id) => ProtocolRequestId::String(id.to_string()),
                    RequestId::Number(id) => ProtocolRequestId::Integer(*id),
                },
            }),
        };
        match events.try_send(event) {
            Ok(()) | Err(async_channel::TrySendError::Closed(_)) => {}
            Err(async_channel::TrySendError::Full(_)) => {
                unreachable!("validated unbounded elicitation event channel became full")
            }
        }
    }
}

impl ElicitationRequestRouter {
    fn take_responder(
        &self,
        key: &(String, RequestId),
    ) -> Result<Option<(ElicitationResponder, ElicitationDeliveryGuard)>> {
        let mut state = self
            .requests
            .lock()
            .map_err(|_| anyhow!("elicitation request router unavailable"))?;
        Ok(state.pending.remove(key).map(|responder| {
            state.terminal_actions += 1;
            (responder, ElicitationDeliveryGuard(self.clone()))
        }))
    }

    /// Fence unbounded source publication and terminal deliveries before reporting shutdown.
    /// Bounded sources retain their existing request-only behavior and backpressure.
    pub(crate) async fn close(&self) {
        {
            let (pending, _terminal_action) = {
                let mut state = self
                    .requests
                    .lock()
                    .unwrap_or_else(std::sync::PoisonError::into_inner);
                if state.phase == RouterPhase::Closed {
                    return;
                }
                if state.phase == RouterPhase::Open {
                    state.phase = RouterPhase::Closing;
                }
                let pending = state
                    .pending
                    .extract_if(|_, responder| responder.abandonment_events.is_some())
                    .collect::<ResponderMap>();
                state.terminal_actions += 1;
                (pending, ElicitationDeliveryGuard(self.clone()))
            };
            // This synchronous batch completes before the first cancellation point.
            for (key, responder) in pending {
                emit_abandonment(&key, responder.abandonment_events.as_ref());
                drop(responder);
            }
        }
        loop {
            let notified = self.terminal_actions_finished.notified();
            tokio::pin!(notified);
            notified.as_mut().enable();
            if self
                .requests
                .lock()
                .unwrap_or_else(std::sync::PoisonError::into_inner)
                .phase
                == RouterPhase::Closed
            {
                return;
            }
            notified.await;
        }
    }

    pub(crate) fn auto_deny(&self) -> bool {
        self.auto_deny.load(Ordering::Relaxed)
    }

    pub(crate) fn set_auto_deny(&self, auto_deny: bool) {
        self.auto_deny.store(auto_deny, Ordering::Relaxed);
    }

    pub(crate) fn full_access_form_input_enabled(&self) -> bool {
        self.full_access_form_input_enabled.load(Ordering::Acquire)
    }

    pub(crate) fn enable_full_access_form_input(&self) {
        self.full_access_form_input_enabled
            .store(true, Ordering::Release);
    }

    pub(crate) async fn resolve(
        &self,
        server_name: String,
        id: RequestId,
        response: ElicitationResponse,
    ) -> Result<()> {
        let key = (server_name, id);
        let (responder, _terminal_action) = self
            .take_responder(&key)?
            .ok_or_else(|| anyhow!("elicitation request not found"))?;
        if responder.response.send(response).is_err() {
            emit_abandonment(&key, responder.abandonment_events.as_ref());
            return Err(anyhow!("elicitation response receiver closed"));
        }
        Ok(())
    }

    async fn request_user_interaction(
        &self,
        events: Option<Sender<Event>>,
        authority: &ElicitationAuthority,
        server_name: String,
        request: ElicitationRequest,
    ) -> Result<ElicitationResponse> {
        let Some(events) = events else {
            return Ok(ElicitationResponse {
                action: ElicitationAction::Decline,
                content: None,
                meta: None,
            });
        };
        // Core uses an unbounded source channel, which can atomically publish a request
        // and synchronously enqueue abandonment from Drop. Preserve bounded callers'
        // existing request-only behavior because they require asynchronous backpressure.
        let abandonment_events = events.capacity().is_none().then(|| events.clone());
        let (delivery_context, response_context) =
            if matches!(&request, ElicitationRequest::UserVerification { .. }) {
                (
                    "failed to deliver user-verification request",
                    "user-verification response channel closed",
                )
            } else {
                (
                    "failed to deliver MCP elicitation request",
                    "elicitation request channel closed unexpectedly",
                )
            };
        let public_request_id = format!(
            "codex-mcp-elicitation-{}",
            NEXT_ELICITATION_REQUEST_ID.fetch_add(1, Ordering::Relaxed)
        );
        let request_key = (
            server_name.clone(),
            RequestId::String(public_request_id.clone().into()),
        );
        let (tx, rx) = oneshot::channel();
        let _active_elicitation = authority
            .lifecycle
            .as_ref()
            .map(ElicitationLifecycle::start);
        let event = Event {
            id: "mcp_elicitation_request".to_string(),
            msg: EventMsg::ElicitationRequest(ElicitationRequestEvent {
                turn_id: None,
                server_name,
                id: ProtocolRequestId::String(public_request_id),
                request,
            }),
        };
        let legacy_event = if abandonment_events.is_some() {
            let mut state = self
                .requests
                .lock()
                .map_err(|_| anyhow!("elicitation request router unavailable"))?;
            anyhow::ensure!(
                state.phase == RouterPhase::Open,
                "elicitation request router closed"
            );
            let mut responder =
                state
                    .pending
                    .entry(request_key.clone())
                    .insert_entry(ElicitationResponder {
                        response: tx,
                        abandonment_events: None,
                    });
            // Resolve and close cannot observe a published request without its responder.
            // The validated unbounded send cannot wait or yield while admission is locked.
            if let Err(error) = events.try_send(event) {
                responder.remove();
                return Err(error).context(delivery_context);
            }
            responder.get_mut().abandonment_events = abandonment_events;
            None
        } else {
            self.requests
                .lock()
                .map_err(|_| anyhow!("elicitation request router unavailable"))?
                .pending
                .insert(
                    request_key.clone(),
                    ElicitationResponder {
                        response: tx,
                        abandonment_events: None,
                    },
                );
            Some(event)
        };
        let _pending_request = ElicitationRequestGuard {
            router: self.clone(),
            key: request_key,
        };
        if let Some(event) = legacy_event {
            events.send(event).await.context(delivery_context)?;
        }
        rx.await.context(response_context)
    }
}

#[derive(Clone)]
pub(crate) struct ElicitationAuthority {
    pub(crate) config: Arc<McpConfig>,
    reviewer: Option<ElicitationReviewerHandle>,
    lifecycle: Option<ElicitationLifecycle>,
}

#[derive(Clone, Default)]
pub(crate) struct ElicitationRequestManager {
    router: ElicitationRequestRouter,
    pub(crate) authority: Arc<StdMutex<Option<ElicitationAuthority>>>,
}

impl ElicitationRequestManager {
    pub(crate) fn new(
        config: Arc<McpConfig>,
        reviewer: Option<ElicitationReviewerHandle>,
        lifecycle: Option<ElicitationLifecycle>,
        router: ElicitationRequestRouter,
    ) -> Self {
        Self {
            router,
            authority: Arc::new(StdMutex::new(Some(ElicitationAuthority {
                config,
                reviewer,
                lifecycle,
            }))),
        }
    }

    pub(crate) fn update(
        &self,
        config: Arc<McpConfig>,
        reviewer: Option<ElicitationReviewerHandle>,
        lifecycle: Option<ElicitationLifecycle>,
    ) -> bool {
        let Ok(mut authority) = self.authority.lock() else {
            return false;
        };
        *authority = Some(ElicitationAuthority {
            config,
            reviewer,
            lifecycle,
        });
        true
    }

    pub(crate) fn make_sender(
        &self,
        server_name: String,
        tx_event: Option<Sender<Event>>,
        client_mcp_extensions: &ClientMcpExtensions,
    ) -> SendElicitation {
        // Event receivers such as codex mcp-server do not necessarily handle verification.
        // Only wait for a response when trusted host activation enabled this exact route.
        let user_verification_enabled = client_mcp_extensions
            .get(OPENAI_ELICITATION_EXTENSION_ID)
            .and_then(|settings| settings.get("userVerification"))
            .is_some_and(Value::is_object);
        let router = self.router.clone();
        let authority = self.authority.clone();
        Box::new(move |id, elicitation| {
            let router = router.clone();
            let tx_event = tx_event.clone();
            let server_name = server_name.clone();
            let authority = authority.clone();
            async move {
                let authority = authority
                    .lock()
                    .ok()
                    .and_then(|authority| authority.clone());
                if let Elicitation::UserVerification {
                    meta,
                    title,
                    description,
                    challenge,
                } = elicitation
                {
                    if !user_verification_enabled {
                        return Ok(ElicitationResponse {
                            action: ElicitationAction::Cancel,
                            content: None,
                            meta: Some(
                                codex_rmcp_client::UserVerificationReason::ApprovalUnavailable
                                    .into_meta(),
                            ),
                        });
                    }
                    return user_verification_elicitation::route(
                        router,
                        tx_event,
                        authority,
                        server_name,
                        ElicitationRequest::UserVerification {
                            meta,
                            title,
                            description,
                            challenge,
                        },
                    )
                    .await;
                }
                if router.auto_deny() {
                    return Ok(ElicitationResponse {
                        action: ElicitationAction::Decline,
                        content: None,
                        meta: None,
                    });
                }

                if elicitation
                    .meta()
                    .and_then(|meta| meta.get(APPROVAL_KIND_KEY))
                    .and_then(serde_json::Value::as_str)
                    == Some(APPROVAL_KIND_TOOL_SUGGESTION)
                {
                    return Ok(ElicitationResponse {
                        action: ElicitationAction::Decline,
                        content: None,
                        meta: None,
                    });
                }

                let Some(authority) = authority else {
                    return Ok(ElicitationResponse {
                        action: ElicitationAction::Decline,
                        content: None,
                        meta: None,
                    });
                };
                let ElicitationAuthority {
                    config,
                    reviewer,
                    lifecycle,
                    ..
                } = &authority;
                let approval_policy = config.approval_policy.value();
                let Some(permission_profile) = config.permission_profile_for_server(&server_name)
                else {
                    return Ok(ElicitationResponse {
                        action: ElicitationAction::Decline,
                        content: None,
                        meta: None,
                    });
                };

                match elicitation
                    .meta()
                    .and_then(|meta| meta.get(STRICT_AUTO_REVIEW_KEY))
                {
                    Some(Value::Bool(true)) => {
                        if matches!(
                            approval_policy,
                            AskForApproval::Granular(config) if !config.allows_mcp_elicitations()
                        ) {
                            return Ok(strict_auto_review_decline());
                        }
                        let Some(reviewer) = reviewer.as_ref() else {
                            return Ok(strict_auto_review_decline());
                        };
                        let _active_elicitation =
                            lifecycle.as_ref().map(ElicitationLifecycle::start);
                        return Ok(
                            match reviewer
                                .review(ElicitationReviewRequest {
                                    server_name,
                                    request_id: id,
                                    elicitation,
                                })
                                .await
                            {
                                Ok(Some(response))
                                    if (response.action == ElicitationAction::Accept
                                        && response.content == Some(serde_json::json!({}))
                                        || matches!(
                                            response.action,
                                            ElicitationAction::Decline | ElicitationAction::Cancel
                                        ) && response.content.is_none())
                                        && response
                                            .meta
                                            .as_ref()
                                            .and_then(Value::as_object)
                                            .and_then(|meta| meta.get(APPROVALS_REVIEWER_KEY))
                                            .and_then(Value::as_str)
                                            == Some("auto_review") =>
                                {
                                    response
                                }
                                Ok(Some(_)) | Ok(None) | Err(_) => strict_auto_review_decline(),
                            },
                        );
                    }
                    None | Some(Value::Bool(false)) => {}
                    Some(_) => return Ok(strict_auto_review_decline()),
                }

                let permission_prompt_is_auto_approved = mcp_permission_prompt_is_auto_approved(
                    approval_policy,
                    permission_profile,
                    McpPermissionPromptAutoApproveContext::default(),
                );
                if permission_prompt_is_auto_approved && can_auto_accept_elicitation(&elicitation) {
                    return Ok(ElicitationResponse {
                        action: ElicitationAction::Accept,
                        content: Some(serde_json::json!({})),
                        meta: None,
                    });
                }

                let should_surface_form_in_full_access = router.full_access_form_input_enabled()
                    && permission_prompt_is_auto_approved
                    && !elicitation
                        .meta()
                        .is_some_and(|meta| meta.contains_key(APPROVAL_KIND_KEY))
                    && match &elicitation {
                        Elicitation::Mcp(
                            rmcp::model::ElicitRequestParams::FormElicitationParams {
                                requested_schema,
                                ..
                            },
                        ) => !requested_schema.properties.is_empty(),
                        Elicitation::OpenAiElicitationForm {
                            requested_schema, ..
                        } => requested_schema
                            .get("properties")
                            .and_then(Value::as_object)
                            .is_some_and(|properties| !properties.is_empty()),
                        Elicitation::Mcp(_)
                        | Elicitation::OpenAiForm { .. }
                        | Elicitation::UserVerification { .. } => false,
                    };

                if !should_surface_form_in_full_access {
                    if elicitation_is_rejected_by_policy(approval_policy) {
                        return Ok(ElicitationResponse {
                            action: ElicitationAction::Decline,
                            content: None,
                            meta: None,
                        });
                    }

                    if let Some(reviewer) = reviewer {
                        let request = ElicitationReviewRequest {
                            server_name: server_name.clone(),
                            request_id: id.clone(),
                            elicitation: elicitation.clone(),
                        };
                        if let Some(response) = reviewer.review(request).await? {
                            return Ok(response);
                        }
                    }
                }

                let request = match elicitation {
                    Elicitation::UserVerification { .. } => {
                        return Ok(ElicitationResponse {
                            action: ElicitationAction::Cancel,
                            content: None,
                            meta: None,
                        });
                    }
                    Elicitation::Mcp(rmcp::model::ElicitRequestParams::FormElicitationParams {
                        meta,
                        message,
                        requested_schema,
                    }) => ElicitationRequest::Form {
                        meta: meta
                            .map(serde_json::to_value)
                            .transpose()
                            .context("failed to serialize MCP elicitation metadata")?,
                        message,
                        requested_schema: serde_json::to_value(requested_schema)
                            .context("failed to serialize MCP elicitation schema")?,
                    },
                    Elicitation::Mcp(rmcp::model::ElicitRequestParams::UrlElicitationParams {
                        meta,
                        message,
                        url,
                        elicitation_id,
                    }) => ElicitationRequest::Url {
                        meta: meta
                            .map(serde_json::to_value)
                            .transpose()
                            .context("failed to serialize MCP elicitation metadata")?,
                        message,
                        url,
                        elicitation_id,
                    },
                    Elicitation::Mcp(_) => {
                        return Ok(ElicitationResponse {
                            action: ElicitationAction::Decline,
                            content: None,
                            meta: None,
                        });
                    }
                    Elicitation::OpenAiForm {
                        meta,
                        message,
                        requested_schema,
                    } => ElicitationRequest::OpenAiForm {
                        meta,
                        message,
                        requested_schema,
                    },
                    Elicitation::OpenAiElicitationForm {
                        meta,
                        message,
                        requested_schema,
                    } => ElicitationRequest::OpenAiElicitationForm {
                        meta,
                        message,
                        requested_schema,
                    },
                };
                router
                    .request_user_interaction(tx_event, &authority, server_name, request)
                    .await
            }
            .boxed()
        })
    }
}

fn strict_auto_review_decline() -> ElicitationResponse {
    ElicitationResponse {
        action: ElicitationAction::Decline,
        content: None,
        meta: Some(serde_json::json!({
            "message": STRICT_AUTO_REVIEW_DECLINE_MESSAGE,
        })),
    }
}

pub(crate) fn elicitation_is_rejected_by_policy(approval_policy: AskForApproval) -> bool {
    match approval_policy {
        AskForApproval::Never => true,
        AskForApproval::OnRequest => false,
        AskForApproval::UnlessTrusted => false,
        AskForApproval::Granular(granular_config) => !granular_config.allows_mcp_elicitations(),
    }
}

type ResponderMap = HashMap<(String, RequestId), ElicitationResponder>;

fn can_auto_accept_elicitation(elicitation: &Elicitation) -> bool {
    match elicitation {
        Elicitation::Mcp(rmcp::model::ElicitRequestParams::FormElicitationParams {
            requested_schema,
            ..
        }) => {
            // Auto-accept confirm/approval elicitations without schema requirements.
            requested_schema.properties.is_empty()
        }
        Elicitation::Mcp(_)
        | Elicitation::OpenAiForm { .. }
        | Elicitation::OpenAiElicitationForm { .. }
        | Elicitation::UserVerification { .. } => false,
    }
}

#[cfg(test)]
#[path = "elicitation_tests.rs"]
mod tests;