codex/codex-rs/rmcp-client/src/elicitation_client_service.rs

use std::collections::HashMap;
use std::collections::HashSet;
use std::sync::Arc;
use std::sync::Mutex;
use std::sync::PoisonError;

use codex_protocol::mcp::OPENAI_ELICITATION_EXTENSION_ID;

use rmcp::RoleClient;
use rmcp::model::ClientInfo;
use rmcp::model::ClientResult;
use rmcp::model::CustomRequest;
use rmcp::model::CustomResult;
use rmcp::model::ElicitResult;
use rmcp::model::ElicitationAction;
use rmcp::model::MetaObject;
use rmcp::model::ProtocolVersion;
use rmcp::model::RequestId;
use rmcp::model::RequestMetaObject;
use rmcp::model::RequestParamsMeta;
use rmcp::model::ServerNotification;
use rmcp::model::ServerRequest;
use rmcp::service::NotificationContext;
use rmcp::service::RequestContext;
use rmcp::service::Service;
use serde::Deserialize;
use serde::Serialize;
use serde_json::Map;
use serde_json::Value;
use tokio::sync::oneshot;

use crate::logging_client_handler::LoggingClientHandler;
use crate::rmcp_client::Elicitation;
use crate::rmcp_client::ElicitationPauseState;
use crate::rmcp_client::ElicitationResponse;
use crate::rmcp_client::SendElicitation;

const MCP_PROGRESS_TOKEN_META_KEY: &str = "progressToken";
const MCP_ELICITATION_CREATE_METHOD: &str = "elicitation/create";
const OPENAI_FORM_METHOD: &str = "openai/form";
const OPENAI_ELICITATION_METHOD: &str = "openai/elicitation/create";

#[derive(Deserialize)]
#[serde(tag = "mode")]
enum OpenAiElicitationRequestParams {
    #[serde(rename = "form")]
    Form(OpenAiFormRequestParams),
}

#[derive(Deserialize)]
#[serde(rename_all = "camelCase")]
struct OpenAiFormRequestParams {
    #[serde(rename = "_meta")]
    meta: Option<Value>,
    message: String,
    requested_schema: Value,
}

#[derive(Clone)]
pub(crate) struct ElicitationClientService {
    handler: LoggingClientHandler,
    supports_openai_form: bool,
    supports_openai_elicitation_form: bool,
    supports_user_verification: bool,
    send_elicitation: Arc<SendElicitation>,
    pause_state: ElicitationPauseState,
    pending_verifications: Arc<Mutex<VerificationCancellations>>,
}

// A notification handler can run before its request handler. Never evict an early
// cancellation: after saturation, cancel new elicitations for this connection.
const MAX_EARLY_CANCELLATIONS: usize = 1024;

#[derive(Default)]
struct VerificationCancellations {
    pending: HashMap<RequestId, oneshot::Sender<()>>,
    early: HashSet<RequestId>,
    saturated: bool,
}

struct PendingVerification {
    request_id: RequestId,
    cancellations: Arc<Mutex<VerificationCancellations>>,
}

impl Drop for PendingVerification {
    fn drop(&mut self) {
        self.cancellations
            .lock()
            .unwrap_or_else(PoisonError::into_inner)
            .pending
            .remove(&self.request_id);
    }
}

impl ElicitationClientService {
    pub(crate) fn new(
        client_info: ClientInfo,
        send_elicitation: SendElicitation,
        pause_state: ElicitationPauseState,
    ) -> Self {
        let supports_openai_form = client_info
            .capabilities
            .extensions
            .as_ref()
            .is_some_and(|extensions| extensions.contains_key(OPENAI_FORM_METHOD));
        let supports_openai_elicitation_form = client_info
            .capabilities
            .extensions
            .as_ref()
            .and_then(|extensions| extensions.get(OPENAI_ELICITATION_EXTENSION_ID))
            .and_then(|settings| settings.get("form"))
            .is_some_and(Value::is_object);
        let send_elicitation = Arc::new(send_elicitation);
        let supports_user_verification = client_info
            .capabilities
            .extensions
            .as_ref()
            .and_then(|extensions| extensions.get(OPENAI_ELICITATION_EXTENSION_ID))
            .and_then(|settings| settings.get("userVerification"))
            .is_some_and(Value::is_object);
        Self {
            handler: LoggingClientHandler::new(
                client_info,
                clone_send_elicitation(Arc::clone(&send_elicitation)),
            ),
            supports_openai_form,
            supports_openai_elicitation_form,
            supports_user_verification,
            send_elicitation,
            pause_state,
            pending_verifications: Arc::default(),
        }
    }

    async fn create_elicitation(
        &self,
        request: Elicitation,
        context: RequestContext<RoleClient>,
    ) -> Result<ElicitationResponse, rmcp::ErrorData> {
        let RequestContext { id, meta, ct, .. } = context;
        let request = restore_context_meta(request, meta);
        let user_verification = matches!(&request, Elicitation::UserVerification { .. });
        let (cancel_tx, cancel_rx) = oneshot::channel();
        let _pending = {
            let mut cancellations = self
                .pending_verifications
                .lock()
                .unwrap_or_else(PoisonError::into_inner);
            if cancellations.saturated || cancellations.early.remove(&id) {
                return Ok(ElicitationResponse {
                    action: ElicitationAction::Cancel,
                    content: None,
                    meta: None,
                });
            }
            cancellations.pending.insert(id.clone(), cancel_tx);
            PendingVerification {
                request_id: id.clone(),
                cancellations: Arc::clone(&self.pending_verifications),
            }
        };
        let _pause = self.pause_state.enter();
        let response = tokio::select! {
            biased;
            _ = ct.cancelled() => {
                return Ok(ElicitationResponse {
                    action: ElicitationAction::Cancel,
                    content: None,
                    meta: None,
                });
            }
            _ = cancel_rx => {
                return Ok(ElicitationResponse {
                    action: ElicitationAction::Cancel,
                    content: None,
                    meta: None,
                });
            }
            response = (self.send_elicitation)(id, request) => response,
        }
        .map_err(|err| rmcp::ErrorData::internal_error(err.to_string(), None))?;
        Ok(if user_verification {
            crate::user_verification::validate_response(response)
        } else {
            response
        })
    }
}

fn clone_send_elicitation(send_elicitation: Arc<SendElicitation>) -> SendElicitation {
    Box::new(move |request_id, request| send_elicitation(request_id, request))
}

impl Service<RoleClient> for ElicitationClientService {
    async fn handle_request(
        &self,
        request: ServerRequest,
        context: RequestContext<RoleClient>,
    ) -> Result<ClientResult, rmcp::ErrorData> {
        match request {
            ServerRequest::ElicitRequest(request) => {
                let modern_session = context
                    .peer
                    .peer_info()
                    .is_some_and(|info| info.protocol_version >= ProtocolVersion::V_2026_07_28);
                let response = self
                    .create_elicitation(Elicitation::Mcp(request.params), context)
                    .await?;
                if modern_session {
                    Ok(ClientResult::ElicitResult(typed_elicitation_result(
                        response,
                    )?))
                } else {
                    Ok(ClientResult::CustomResult(elicitation_response_result(
                        response,
                    )?))
                }
            }
            ServerRequest::CustomRequest(request)
                if request.method == MCP_ELICITATION_CREATE_METHOD =>
            {
                let modern_session = context
                    .peer
                    .peer_info()
                    .is_some_and(|info| info.protocol_version >= ProtocolVersion::V_2026_07_28);
                let response = self
                    .create_elicitation(custom_mcp_elicitation(request)?, context)
                    .await?;
                if modern_session {
                    Ok(ClientResult::ElicitResult(typed_elicitation_result(
                        response,
                    )?))
                } else {
                    Ok(ClientResult::CustomResult(elicitation_response_result(
                        response,
                    )?))
                }
            }
            ServerRequest::CustomRequest(request)
                if request.method == OPENAI_ELICITATION_METHOD
                    && self.supports_user_verification
                    && request
                        .params
                        .as_ref()
                        .and_then(|params| params.get("mode"))
                        .and_then(Value::as_str)
                        == Some(crate::user_verification::MODE) =>
            {
                let response = self
                    .create_elicitation(crate::user_verification::parse_request(request)?, context)
                    .await?;
                Ok(ClientResult::CustomResult(elicitation_response_result(
                    response,
                )?))
            }
            ServerRequest::CustomRequest(request)
                if request.method == OPENAI_FORM_METHOD && self.supports_openai_form =>
            {
                let response = self
                    .create_elicitation(openai_form_elicitation(request)?, context)
                    .await?;
                Ok(ClientResult::CustomResult(elicitation_response_result(
                    response,
                )?))
            }
            ServerRequest::CustomRequest(request)
                if request.method == OPENAI_ELICITATION_METHOD
                    && self.supports_openai_elicitation_form =>
            {
                let params = request
                    .params_as::<OpenAiElicitationRequestParams>()
                    .map_err(|err| {
                        rmcp::ErrorData::invalid_params(err.to_string(), /*data*/ None)
                    })?
                    .ok_or_else(|| {
                        rmcp::ErrorData::invalid_params("missing params", /*data*/ None)
                    })?;
                let OpenAiElicitationRequestParams::Form(params) = params;
                let response = self
                    .create_elicitation(
                        Elicitation::OpenAiElicitationForm {
                            meta: params.meta,
                            message: params.message,
                            requested_schema: params.requested_schema,
                        },
                        context,
                    )
                    .await?;
                Ok(ClientResult::CustomResult(elicitation_response_result(
                    response,
                )?))
            }
            ServerRequest::CustomRequest(request)
                if request.method == OPENAI_ELICITATION_METHOD
                    && self.supports_user_verification =>
            {
                Err(rmcp::ErrorData::invalid_params(
                    "invalid elicitation mode",
                    /*data*/ None,
                ))
            }
            request => {
                <LoggingClientHandler as Service<RoleClient>>::handle_request(
                    &self.handler,
                    request,
                    context,
                )
                .await
            }
        }
    }

    async fn handle_notification(
        &self,
        notification: ServerNotification,
        context: NotificationContext<RoleClient>,
    ) -> Result<(), rmcp::ErrorData> {
        if let ServerNotification::CancelledNotification(cancelled) = &notification
            && let Some(request_id) = cancelled.params.request_id.as_ref()
        {
            let mut cancellations = self
                .pending_verifications
                .lock()
                .unwrap_or_else(PoisonError::into_inner);
            if let Some(cancel) = cancellations.pending.remove(request_id) {
                let _ = cancel.send(());
            } else if !cancellations.saturated && !cancellations.early.contains(request_id) {
                // Bound both the number and size of retained request IDs.
                if cancellations.early.len() == MAX_EARLY_CANCELLATIONS
                    || matches!(request_id, RequestId::String(id) if id.len() > 1024)
                {
                    cancellations.saturated = true;
                    cancellations.early.clear();
                } else {
                    cancellations.early.insert(request_id.clone());
                }
            }
        }
        <LoggingClientHandler as Service<RoleClient>>::handle_notification(
            &self.handler,
            notification,
            context,
        )
        .await
    }

    fn get_info(&self) -> ClientInfo {
        <LoggingClientHandler as Service<RoleClient>>::get_info(&self.handler)
    }
}

fn custom_mcp_elicitation(request: CustomRequest) -> Result<Elicitation, rmcp::ErrorData> {
    let raw_params = request
        .params
        .ok_or_else(|| rmcp::ErrorData::invalid_params("missing params", None))?;
    let params: rmcp::model::ElicitRequestParams = serde_json::from_value(raw_params)
        .map_err(|err| rmcp::ErrorData::invalid_params(err.to_string(), None))?;
    Ok(Elicitation::Mcp(params))
}

fn openai_form_elicitation(request: CustomRequest) -> Result<Elicitation, rmcp::ErrorData> {
    let params = request
        .params_as::<OpenAiFormRequestParams>()
        .map_err(|err| rmcp::ErrorData::invalid_params(err.to_string(), None))?
        .ok_or_else(|| rmcp::ErrorData::invalid_params("missing params", None))?;
    Ok(Elicitation::OpenAiForm {
        meta: params.meta,
        message: params.message,
        requested_schema: params.requested_schema,
    })
}

fn restore_context_meta(
    mut request: Elicitation,
    mut context_meta: RequestMetaObject,
) -> Elicitation {
    // RMCP lifts JSON-RPC `_meta` into RequestContext before invoking services.
    context_meta.remove(MCP_PROGRESS_TOKEN_META_KEY);
    if context_meta.is_empty() {
        return request;
    }

    match &mut request {
        Elicitation::UserVerification { .. } => {}
        Elicitation::Mcp(request) => request
            .meta_mut()
            .get_or_insert_with(RequestMetaObject::new)
            .extend(context_meta),
        Elicitation::OpenAiForm { meta, .. } | Elicitation::OpenAiElicitationForm { meta, .. } => {
            let meta = meta
                .get_or_insert_with(|| Value::Object(Map::new()))
                .as_object_mut();
            if let Some(meta) = meta {
                meta.extend(context_meta.0.0);
            }
        }
    }
    request
}

#[derive(Serialize)]
#[serde(rename_all = "camelCase")]
struct CreateElicitationResultWithMeta {
    action: ElicitationAction,
    #[serde(skip_serializing_if = "Option::is_none")]
    content: Option<Value>,
    #[serde(rename = "_meta", skip_serializing_if = "Option::is_none")]
    meta: Option<Value>,
}

fn elicitation_response_result(
    response: ElicitationResponse,
) -> Result<CustomResult, rmcp::ErrorData> {
    let ElicitationResponse {
        action,
        content,
        meta,
    } = response;
    let result = CreateElicitationResultWithMeta {
        action,
        content,
        meta,
    };

    serde_json::to_value(result)
        .map(CustomResult)
        .map_err(|err| rmcp::ErrorData::internal_error(err.to_string(), None))
}

fn typed_elicitation_result(
    response: ElicitationResponse,
) -> Result<ElicitResult, rmcp::ErrorData> {
    let ElicitationResponse {
        action,
        content,
        meta,
    } = response;
    let mut result = ElicitResult::new(action);
    result.content = content;
    result.meta = match meta {
        None => None,
        Some(Value::Object(meta)) => Some(MetaObject::from(meta)),
        Some(meta) => {
            return Err(rmcp::ErrorData::invalid_params(
                format!("MCP elicitation response _meta must be an object, got {meta}"),
                None,
            ));
        }
    };
    Ok(result)
}

#[cfg(test)]
mod tests {
    use pretty_assertions::assert_eq;
    use rmcp::model::BooleanSchema;
    use rmcp::model::ElicitRequestParams;
    use rmcp::model::ElicitationSchema;
    use rmcp::model::PrimitiveSchemaDefinition;
    use serde_json::Value;
    use serde_json::json;

    use super::*;

    #[test]
    fn restore_context_meta_adds_elicitation_meta_and_removes_progress_token() {
        let request = restore_context_meta(
            Elicitation::Mcp(form_request(/*meta*/ None)),
            meta(json!({
                "progressToken": "progress-token",
                "persist": ["session", "always"],
            })),
        );

        assert_eq!(
            request,
            Elicitation::Mcp(form_request(Some(meta(json!({
                "persist": ["session", "always"],
            })))))
        );
    }

    #[test]
    fn legacy_sep1034_elicitation_without_mode_preserves_schema_defaults() {
        let request = json!({
            "method": "elicitation/create",
            "params": {
                "message": "Confirm the default values",
                "requestedSchema": {
                    "type": "object",
                    "properties": {
                        "name": {"type": "string", "default": "John Doe"},
                        "age": {"type": "integer", "default": 30},
                        "score": {"type": "number", "default": 95.5},
                        "status": {
                            "type": "string",
                            "enum": ["active", "inactive"],
                            "default": "active",
                        },
                        "verified": {"type": "boolean", "default": true},
                    },
                    "required": [],
                },
            },
        });

        let request = serde_json::from_value::<ServerRequest>(request)
            .expect("legacy form elicitations must deserialize without a mode");
        let ServerRequest::ElicitRequest(request) = request else {
            panic!("legacy elicitation/create must dispatch to the typed handler");
        };
        let ElicitRequestParams::FormElicitationParams {
            requested_schema, ..
        } = request.params
        else {
            panic!("an omitted legacy elicitation mode must default to form");
        };

        assert_eq!(
            serde_json::to_value(requested_schema)
                .expect("legacy schema defaults must remain serializable"),
            json!({
                "type": "object",
                "properties": {
                    "name": {"type": "string", "default": "John Doe"},
                    "age": {"type": "integer", "default": 30},
                    "score": {"type": "number", "default": 95.5},
                    "status": {
                        "type": "string",
                        "enum": ["active", "inactive"],
                        "default": "active",
                    },
                    "verified": {"type": "boolean", "default": true},
                },
                "required": [],
            })
        );
    }

    #[test]
    fn parses_legacy_custom_elicitation_without_mode() {
        let request = CustomRequest::new(
            MCP_ELICITATION_CREATE_METHOD,
            Some(json!({
                "message": "Confirm?",
                "requestedSchema": {
                    "type": "object",
                    "properties": {
                        "confirmed": {"type": "boolean"},
                        "age": {"type": "integer", "minimum": 1, "maximum": 99, "default": 30},
                    },
                    "required": ["confirmed"],
                },
            })),
        );
        let Elicitation::Mcp(rmcp::model::ElicitRequestParams::FormElicitationParams {
            requested_schema,
            ..
        }) = custom_mcp_elicitation(request)
            .expect("legacy custom elicitation parameters must deserialize")
        else {
            panic!("omitted legacy elicitation mode must default to form");
        };

        assert_eq!(
            serde_json::to_value(requested_schema).expect("schema must serialize"),
            json!({
                "type": "object",
                "properties": {
                    "confirmed": {"type": "boolean"},
                    "age": {"type": "integer", "minimum": 1, "maximum": 99, "default": 30},
                },
                "required": ["confirmed"],
            })
        );
    }

    #[test]
    fn parses_openai_form_custom_requests() {
        let elicitation = openai_form_elicitation(CustomRequest::new(
            OPENAI_FORM_METHOD,
            Some(json!({
                "message": "Select a template",
                "requestedSchema": {
                    "type": "object",
                    "properties": {
                        "template": {
                            "type": "openai/imagePicker",
                            "items": [{
                                "id": "monthly-review",
                                "title": "Monthly review",
                                "image": "data:image/svg+xml;base64,PHN2ZyB4bWxucz0iaHR0cDovL3d3dy53My5vcmcvMjAwMC9zdmciLz4="
                            }]
                        }
                    }
                }
            })),
        ))
        .expect("valid openai/form request");

        assert_eq!(
            elicitation,
            Elicitation::OpenAiForm {
                meta: None,
                message: "Select a template".to_string(),
                requested_schema: json!({
                    "type": "object",
                    "properties": {
                        "template": {
                            "type": "openai/imagePicker",
                            "items": [{
                                "id": "monthly-review",
                                "title": "Monthly review",
                                "image": "data:image/svg+xml;base64,PHN2ZyB4bWxucz0iaHR0cDovL3d3dy53My5vcmcvMjAwMC9zdmciLz4="
                            }]
                        }
                    }
                }),
            }
        );
    }

    #[test]
    fn elicitation_response_result_serializes_response_meta() {
        let result = rmcp::model::ClientResult::CustomResult(
            elicitation_response_result(ElicitationResponse {
                action: ElicitationAction::Accept,
                content: Some(json!({ "confirmed": true })),
                meta: Some(json!({ "persist": "always" })),
            })
            .expect("elicitation response should serialize"),
        );

        assert_eq!(
            serde_json::to_value(result).expect("client result should serialize"),
            json!({
                "action": "accept",
                "content": { "confirmed": true },
                "_meta": { "persist": "always" },
            })
        );
    }

    #[test]
    fn typed_elicitation_result_preserves_response_meta() {
        let result = typed_elicitation_result(ElicitationResponse {
            action: ElicitationAction::Accept,
            content: Some(json!({ "confirmed": true })),
            meta: Some(json!({ "persist": "always" })),
        })
        .expect("modern elicitation response should serialize");

        assert_eq!(
            serde_json::to_value(result).expect("typed elicitation result should serialize"),
            json!({
                "action": "accept",
                "content": { "confirmed": true },
                "_meta": { "persist": "always" },
            })
        );
    }

    #[test]
    fn typed_elicitation_result_rejects_non_object_meta() {
        let error = typed_elicitation_result(ElicitationResponse {
            action: ElicitationAction::Accept,
            content: None,
            meta: Some(json!(["invalid"])),
        })
        .expect_err("modern elicitation metadata must be an object");

        assert_eq!(error.code, rmcp::model::ErrorCode::INVALID_PARAMS);
    }

    fn form_request(meta: Option<RequestMetaObject>) -> ElicitRequestParams {
        ElicitRequestParams::FormElicitationParams {
            meta,
            message: "Confirm?".to_string(),
            requested_schema: ElicitationSchema::builder()
                .required_property(
                    "confirmed",
                    PrimitiveSchemaDefinition::Boolean(BooleanSchema::new()),
                )
                .build()
                .expect("schema should build"),
        }
    }

    fn meta(value: Value) -> RequestMetaObject {
        let Value::Object(map) = value else {
            panic!("meta must be an object");
        };
        RequestMetaObject::from(map)
    }
}

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