codex/codex-rs/mcp-server/src/message_processor.rs

use std::sync::Arc;

use codex_arg0::Arg0DispatchPaths;
use codex_core::StateDbHandle;
use codex_core::ThreadManager;
use codex_core::config::Config;
use codex_exec_server::EnvironmentManager;
use codex_extension_api::ExtensionRegistryBuilder;
use codex_home::CodexHomeUserInstructionsProvider;
use codex_login::AuthManager;
use codex_login::default_client::USER_AGENT_SUFFIX;
use codex_login::default_client::get_codex_user_agent;
use codex_protocol::protocol::SessionSource;
use rmcp::model::CallToolRequestParams;
use rmcp::model::CallToolResult;
use rmcp::model::ClientNotification;
use rmcp::model::ClientRequest;
use rmcp::model::ErrorCode;
use rmcp::model::ErrorData;
use rmcp::model::Implementation;
use rmcp::model::InitializeResult;
use rmcp::model::JsonRpcError;
use rmcp::model::JsonRpcNotification;
use rmcp::model::JsonRpcRequest;
use rmcp::model::JsonRpcResponse;
use rmcp::model::RequestId;
use rmcp::model::ServerCapabilities;
use serde_json::json;
use tokio::task;

use crate::active_turn_registry::ActiveTurnRegistry;
use crate::codex_tool_config::CodexToolCallParam;
use crate::codex_tool_config::CodexToolCallReplyParam;
use crate::codex_tool_config::create_tool_for_codex_tool_call_param;
use crate::codex_tool_config::create_tool_for_codex_tool_call_reply_param;
use crate::extension_event_sink::extension_event_sink;
use crate::outgoing_message::OutgoingMessageSender;

pub(crate) struct MessageProcessor {
    outgoing: Arc<OutgoingMessageSender>,
    initialized: bool,
    arg0_paths: Arg0DispatchPaths,
    thread_manager: Arc<ThreadManager>,
    active_turns: Arc<ActiveTurnRegistry>,
}

impl MessageProcessor {
    /// Create a new `MessageProcessor`, retaining a handle to the outgoing
    /// `Sender` so handlers can enqueue messages to be written to stdout.
    pub(crate) async fn new(
        outgoing: OutgoingMessageSender,
        arg0_paths: Arg0DispatchPaths,
        config: Arc<Config>,
        environment_manager: Arc<EnvironmentManager>,
        state_db: Option<StateDbHandle>,
        installation_id: String,
    ) -> std::io::Result<Self> {
        let outgoing = Arc::new(outgoing);
        let auth_manager = AuthManager::shared_from_config(
            config.as_ref(),
            /*enable_codex_api_key_env*/ false,
        )
        .await
        .map_err(std::io::Error::other)?;
        let user_instructions_provider = Arc::new(CodexHomeUserInstructionsProvider::new(
            config.codex_home.clone(),
        ));
        let active_turns = Arc::new(ActiveTurnRegistry::default());
        let mut extensions = ExtensionRegistryBuilder::<Config>::with_event_sink(
            extension_event_sink(Arc::clone(&outgoing), Arc::clone(&active_turns)),
        );
        codex_git_attribution::install(
            &mut extensions,
            auth_manager.clone(),
            config.chatgpt_base_url.clone(),
            config.http_client_factory(),
        );
        codex_image_generation_extension::install(
            &mut extensions,
            auth_manager.clone(),
            |config: &Config| Some(config.codex_home.clone()),
        );
        let skill_providers = codex_skills_extension::SkillProviders::new()
            .with_host_provider(Arc::new(codex_skills_extension::HostSkillProvider::new()));
        codex_skills_extension::install_with_providers_and_metrics(
            &mut extensions,
            skill_providers,
            codex_otel::global(),
            |config: &Config| codex_skills_extension::SkillsExtensionConfig {
                include_instructions: config.include_skill_instructions,
                max_context_tokens: config.skill_max_context_tokens,
                bundled_skills_enabled: config.bundled_skills_enabled(),
                orchestrator_skills_enabled: config.orchestrator_skills_enabled,
                shadow_selection_enabled: config
                    .features
                    .enabled(codex_features::Feature::SkillSearch),
            },
        );
        let thread_manager = Arc::new(ThreadManager::new(
            config.as_ref(),
            Arc::clone(&auth_manager),
            codex_core::build_models_manager(config.as_ref(), auth_manager),
            codex_core::CodexAppsToolsCache::default(),
            SessionSource::Mcp,
            environment_manager,
            Arc::new(extensions.build()),
            user_instructions_provider,
            /*analytics_events_client*/ None,
            codex_core::thread_store_from_config(config.as_ref(), state_db.clone()),
            codex_core::local_agent_graph_store_from_state_db(state_db.as_ref()),
            installation_id,
            /*attestation_provider*/ None,
            /*external_time_provider*/ None,
        ));
        Ok(Self {
            outgoing,
            initialized: false,
            arg0_paths,
            thread_manager,
            active_turns,
        })
    }

    pub(crate) async fn process_request(&mut self, request: JsonRpcRequest<ClientRequest>) {
        let request_id = request.id.clone();
        let client_request = request.request;

        match client_request {
            ClientRequest::InitializeRequest(params) => {
                self.handle_initialize(request_id, params.params).await;
            }
            ClientRequest::PingRequest(_params) => {
                self.handle_ping(request_id).await;
            }
            ClientRequest::ListResourcesRequest(params) => {
                self.handle_list_resources(params.params);
            }
            ClientRequest::ListResourceTemplatesRequest(params) => {
                self.handle_list_resource_templates(params.params);
            }
            ClientRequest::ReadResourceRequest(params) => {
                self.handle_read_resource(params.params);
            }
            ClientRequest::SubscribeRequest(params) => {
                self.handle_subscribe(params.params);
            }
            ClientRequest::UnsubscribeRequest(params) => {
                self.handle_unsubscribe(params.params);
            }
            ClientRequest::ListPromptsRequest(params) => {
                self.handle_list_prompts(params.params);
            }
            ClientRequest::GetPromptRequest(params) => {
                self.handle_get_prompt(params.params);
            }
            ClientRequest::ListToolsRequest(params) => {
                self.handle_list_tools(request_id, params.params).await;
            }
            ClientRequest::CallToolRequest(params) => {
                self.handle_call_tool(request_id, params.params).await;
            }
            ClientRequest::SetLevelRequest(params) => {
                self.handle_set_level(params.params);
            }
            ClientRequest::CompleteRequest(params) => {
                self.handle_complete(params.params);
            }
            ClientRequest::GetTaskRequest(_) => {
                self.handle_unsupported_request(request_id, "tasks/get")
                    .await;
            }
            ClientRequest::CancelTaskRequest(_) => {
                self.handle_unsupported_request(request_id, "tasks/cancel")
                    .await;
            }
            ClientRequest::CustomRequest(custom) => {
                let method = custom.method;
                self.outgoing.send_error(
                    request_id,
                    ErrorData::new(
                        ErrorCode::METHOD_NOT_FOUND,
                        format!("method not found: {method}"),
                        Some(json!({ "method": method })),
                    ),
                );
            }
            request => {
                self.handle_unsupported_request(request_id, request.method())
                    .await;
            }
        }
    }

    pub(crate) async fn process_response(&mut self, response: JsonRpcResponse<serde_json::Value>) {
        tracing::info!("<- response: {:?}", response);
        let JsonRpcResponse { id, result, .. } = response;
        self.outgoing.notify_client_response(id, result).await
    }

    pub(crate) async fn process_notification(
        &mut self,
        notification: JsonRpcNotification<ClientNotification>,
    ) {
        match notification.notification {
            ClientNotification::CancelledNotification(params) => {
                self.handle_cancelled_notification(params.params).await;
            }
            ClientNotification::ProgressNotification(params) => {
                self.handle_progress_notification(params.params);
            }
            ClientNotification::RootsListChangedNotification(_params) => {
                self.handle_roots_list_changed();
            }
            ClientNotification::InitializedNotification(_) => {
                self.handle_initialized_notification();
            }
            ClientNotification::CustomNotification(_) => {
                tracing::warn!("ignoring custom client notification");
            }
            _ => {
                tracing::warn!("ignoring unsupported client notification");
            }
        }
    }

    pub(crate) fn process_error(&mut self, err: JsonRpcError) {
        tracing::error!("<- error: {:?}", err);
    }

    async fn handle_initialize(
        &mut self,
        id: RequestId,
        params: rmcp::model::InitializeRequestParams,
    ) {
        tracing::info!("initialize -> params: {:?}", params);

        if self.initialized {
            self.outgoing.send_error(
                id,
                ErrorData::invalid_request("initialize called more than once", None),
            );
            return;
        }

        let client_info = params.client_info;
        let name = client_info.name;
        let version = client_info.version;
        let user_agent_suffix = format!("{name}; {version}");
        if let Ok(mut suffix) = USER_AGENT_SUFFIX.lock() {
            *suffix = Some(user_agent_suffix);
        }

        let server_info =
            Implementation::new("codex-mcp-server", env!("CARGO_PKG_VERSION")).with_title("Codex");

        // Preserve Codex's existing non-spec `serverInfo.user_agent` field.
        let mut server_info_value = match serde_json::to_value(&server_info) {
            Ok(value) => value,
            Err(err) => {
                self.outgoing.send_error(
                    id,
                    ErrorData::internal_error(
                        format!("failed to serialize server info: {err}"),
                        None,
                    ),
                );
                return;
            }
        };
        if let serde_json::Value::Object(ref mut obj) = server_info_value {
            obj.insert("user_agent".to_string(), json!(get_codex_user_agent()));
        }

        let capabilities = ServerCapabilities::builder()
            .enable_tools()
            .enable_tool_list_changed()
            .build();
        let result = InitializeResult::new(capabilities)
            .with_protocol_version(params.protocol_version)
            .with_server_info(server_info);
        let mut result_value = match serde_json::to_value(result) {
            Ok(value) => value,
            Err(err) => {
                self.outgoing.send_error(
                    id,
                    ErrorData::internal_error(
                        format!("failed to serialize initialize response: {err}"),
                        None,
                    ),
                );
                return;
            }
        };

        if let serde_json::Value::Object(ref mut obj) = result_value {
            obj.insert("serverInfo".to_string(), server_info_value);
        }

        self.initialized = true;
        self.outgoing.send_response(id, result_value);
    }

    async fn handle_ping(&self, id: RequestId) {
        tracing::info!("ping");
        self.outgoing.send_response(id, json!({}));
    }

    fn handle_list_resources(&self, params: Option<rmcp::model::PaginatedRequestParams>) {
        tracing::info!("resources/list -> params: {:?}", params);
    }

    fn handle_list_resource_templates(&self, params: Option<rmcp::model::PaginatedRequestParams>) {
        tracing::info!("resources/templates/list -> params: {:?}", params);
    }

    fn handle_read_resource(&self, params: rmcp::model::ReadResourceRequestParams) {
        tracing::info!("resources/read -> params: {:?}", params);
    }

    fn handle_subscribe(&self, params: rmcp::model::SubscribeRequestParams) {
        tracing::info!("resources/subscribe -> params: {:?}", params);
    }

    fn handle_unsubscribe(&self, params: rmcp::model::UnsubscribeRequestParams) {
        tracing::info!("resources/unsubscribe -> params: {:?}", params);
    }

    fn handle_list_prompts(&self, params: Option<rmcp::model::PaginatedRequestParams>) {
        tracing::info!("prompts/list -> params: {:?}", params);
    }

    fn handle_get_prompt(&self, params: rmcp::model::GetPromptRequestParams) {
        tracing::info!("prompts/get -> params: {:?}", params);
    }

    async fn handle_list_tools(
        &self,
        id: RequestId,
        params: Option<rmcp::model::PaginatedRequestParams>,
    ) {
        tracing::trace!("tools/list -> {params:?}");
        let result = rmcp::model::ListToolsResult::with_all_items(vec![
            create_tool_for_codex_tool_call_param(),
            create_tool_for_codex_tool_call_reply_param(),
        ]);

        self.outgoing.send_response(id, result);
    }

    async fn handle_call_tool(&self, id: RequestId, params: CallToolRequestParams) {
        tracing::info!("tools/call -> params: {:?}", params);
        let CallToolRequestParams {
            name, arguments, ..
        } = params;

        match name.as_ref() {
            "codex" => self.handle_tool_call_codex(id, arguments).await,
            "codex-reply" => {
                self.handle_tool_call_codex_session_reply(id, arguments)
                    .await
            }
            _ => {
                let result = CallToolResult::error(vec![rmcp::model::ContentBlock::text(format!(
                    "Unknown tool '{name}'"
                ))]);
                self.outgoing.send_response(id, result);
            }
        }
    }

    async fn handle_tool_call_codex(
        &self,
        id: RequestId,
        arguments: Option<rmcp::model::JsonObject>,
    ) {
        let arguments = arguments.map(serde_json::Value::Object);
        let (initial_prompt, config): (String, Config) = match arguments {
            Some(json_val) => match serde_json::from_value::<CodexToolCallParam>(json_val) {
                Ok(tool_cfg) => match tool_cfg.into_config(self.arg0_paths.clone()).await {
                    Ok(cfg) => cfg,
                    Err(e) => {
                        let result = CallToolResult::error(vec![rmcp::model::ContentBlock::text(
                            format!("Failed to load Codex configuration from overrides: {e}"),
                        )]);
                        self.outgoing.send_response(id, result);
                        return;
                    }
                },
                Err(e) => {
                    let result = CallToolResult::error(vec![rmcp::model::ContentBlock::text(
                        format!("Failed to parse configuration for Codex tool: {e}"),
                    )]);
                    self.outgoing.send_response(id, result);
                    return;
                }
            },
            None => {
                let result = CallToolResult::error(vec![rmcp::model::ContentBlock::text(
                    "Missing arguments for codex tool-call; the `prompt` field is required.",
                )]);
                self.outgoing.send_response(id, result);
                return;
            }
        };

        // Clone outgoing and server to move into async task.
        let outgoing = self.outgoing.clone();
        let thread_manager = self.thread_manager.clone();
        let active_turns = Arc::clone(&self.active_turns);

        // Spawn an async task to handle the Codex session so that we do not
        // block the synchronous message-processing loop.
        task::spawn(async move {
            // Run the Codex session and stream events back to the client.
            crate::codex_tool_runner::run_codex_tool_session(
                id,
                initial_prompt,
                config,
                outgoing,
                thread_manager,
                active_turns,
            )
            .await;
        });
    }

    async fn handle_tool_call_codex_session_reply(
        &self,
        request_id: RequestId,
        arguments: Option<rmcp::model::JsonObject>,
    ) {
        let arguments = arguments.map(serde_json::Value::Object);
        tracing::info!("tools/call -> params: {:?}", arguments);

        // parse arguments
        let codex_tool_call_reply_param: CodexToolCallReplyParam = match arguments {
            Some(json_val) => match serde_json::from_value::<CodexToolCallReplyParam>(json_val) {
                Ok(params) => params,
                Err(e) => {
                    tracing::error!("Failed to parse Codex tool call reply parameters: {e}");
                    let result = CallToolResult::error(vec![rmcp::model::ContentBlock::text(
                        format!("Failed to parse configuration for Codex tool: {e}"),
                    )]);
                    self.outgoing.send_response(request_id, result);
                    return;
                }
            },
            None => {
                tracing::error!(
                    "Missing arguments for codex-reply tool-call; the `thread_id` and `prompt` fields are required."
                );
                let result = CallToolResult::error(vec![rmcp::model::ContentBlock::text(
                    "Missing arguments for codex-reply tool-call; the `thread_id` and `prompt` fields are required.",
                )]);
                self.outgoing.send_response(request_id, result);
                return;
            }
        };

        let thread_id = match codex_tool_call_reply_param.get_thread_id() {
            Ok(id) => id,
            Err(e) => {
                tracing::error!("Failed to parse thread_id: {e}");
                let result = CallToolResult::error(vec![rmcp::model::ContentBlock::text(format!(
                    "Failed to parse thread_id: {e}"
                ))]);
                self.outgoing.send_response(request_id, result);
                return;
            }
        };

        // Clone outgoing to move into async task.
        let outgoing = self.outgoing.clone();
        let active_turns = Arc::clone(&self.active_turns);

        let codex = match self.thread_manager.get_thread(thread_id).await {
            Ok(c) => c,
            Err(_) => {
                tracing::warn!("Session not found for thread_id: {thread_id}");
                let result = crate::codex_tool_runner::create_call_tool_result_with_thread_id(
                    thread_id,
                    format!("Session not found for thread_id: {thread_id}"),
                    Some(true),
                );
                outgoing.send_response(request_id, result);
                return;
            }
        };

        // Spawn the long-running reply handler.
        let prompt = codex_tool_call_reply_param.prompt.clone();
        tokio::spawn({
            let outgoing = outgoing.clone();
            let active_turns = Arc::clone(&active_turns);

            async move {
                crate::codex_tool_runner::run_codex_tool_session_reply(
                    thread_id,
                    codex,
                    outgoing,
                    request_id,
                    prompt,
                    active_turns,
                )
                .await;
            }
        });
    }

    #[allow(deprecated)]
    fn handle_set_level(&self, params: rmcp::model::SetLevelRequestParams) {
        tracing::info!("logging/setLevel -> params: {:?}", params);
    }

    fn handle_complete(&self, params: rmcp::model::CompleteRequestParams) {
        tracing::info!("completion/complete -> params: {:?}", params);
    }

    async fn handle_unsupported_request(&self, id: RequestId, method: &str) {
        self.outgoing.send_error(
            id,
            ErrorData::new(
                ErrorCode::METHOD_NOT_FOUND,
                format!("method not found: {method}"),
                Some(json!({ "method": method })),
            ),
        );
    }

    // ---------------------------------------------------------------------
    // Notification handlers
    // ---------------------------------------------------------------------

    async fn handle_cancelled_notification(&self, params: rmcp::model::CancelledNotificationParam) {
        let Some(request_id) = params.request_id else {
            tracing::warn!("ignoring cancellation without a request id");
            return;
        };
        // Create a stable string form early for logging.
        let request_id_string = request_id.to_string();

        // Resolve the thread for the active MCP request.
        let thread_id = match self.active_turns.thread_id(&request_id) {
            Some(id) => id,
            None => {
                tracing::warn!("Session not found for request_id: {request_id_string}");
                return;
            }
        };
        tracing::info!("thread_id: {thread_id}");

        // Obtain the Codex thread from the server.
        let codex_arc = match self.thread_manager.get_thread(thread_id).await {
            Ok(c) => c,
            Err(_) => {
                tracing::warn!("Session not found for thread_id: {thread_id}");
                return;
            }
        };

        // Submit interrupt to Codex.
        if let Err(e) = codex_arc
            .submit(codex_protocol::protocol::Op::Interrupt)
            .await
        {
            tracing::error!("Failed to submit interrupt to Codex: {e}");
            return;
        }
        // Stop routing extension events to the cancelled turn.
        self.active_turns.unregister(&request_id);
    }

    fn handle_progress_notification(&self, params: rmcp::model::ProgressNotificationParam) {
        tracing::info!("notifications/progress -> params: {:?}", params);
    }

    fn handle_roots_list_changed(&self) {
        tracing::info!("notifications/roots/list_changed");
    }

    fn handle_initialized_notification(&self) {
        tracing::info!("notifications/initialized");
    }
}