codex/codex-rs/app-server/src/request_processors/mcp_processor.rs

use super::account_processor::EnterpriseLoginCompletion;
use super::account_processor::EnterpriseLoginTarget;
use super::thread_input::ensure_direct_input_allowed;
use super::*;
use codex_config::types::McpServerAuth;
use codex_mcp::McpServerSource;
use codex_mcp::ReadResourceRequestParams;
use codex_mcp::ema_auth_scope;
use codex_mcp::resolve_oauth_callback;
use codex_otel::auth_storage::AuthStorageOriginator;
use codex_rmcp_client::EnterpriseOAuthLoginRequest;
use std::future::Future;

use crate::thread_state::ThreadStateManager;

const MCP_TOOL_THREAD_ID_META_KEY: &str = "threadId";

#[derive(Clone)]
pub(crate) struct McpRequestProcessor {
    pub(super) auth_manager: Arc<AuthManager>,
    thread_manager: Arc<ThreadManager>,
    pub(super) outgoing: Arc<OutgoingMessageSender>,
    config_manager: ConfigManager,
    account_processor: Arc<AccountRequestProcessor>,
    pub(super) thread_state_manager: ThreadStateManager,
}

impl McpRequestProcessor {
    pub(crate) fn new(
        auth_manager: Arc<AuthManager>,
        thread_manager: Arc<ThreadManager>,
        thread_state_manager: ThreadStateManager,
        outgoing: Arc<OutgoingMessageSender>,
        config_manager: ConfigManager,
        account_processor: Arc<AccountRequestProcessor>,
    ) -> Self {
        Self {
            auth_manager,
            thread_manager,
            outgoing,
            config_manager,
            account_processor,
            thread_state_manager,
        }
    }

    pub(crate) async fn mcp_server_oauth_login(
        &self,
        params: McpServerOauthLoginParams,
    ) -> Result<Option<ClientResponsePayload>, JSONRPCErrorError> {
        self.mcp_server_oauth_login_response(params)
            .await
            .map(|response| Some(response.into()))
    }

    pub(crate) fn mcp_server_refresh(
        &self,
        params: Option<()>,
    ) -> impl Future<Output = Result<Option<ClientResponsePayload>, JSONRPCErrorError>> + Send + '_
    {
        // Keep refresh state behind a pointer before returning to the shared request handler.
        Box::pin(async move {
            self.mcp_server_refresh_response(params)
                .await
                .map(|response| Some(response.into()))
        })
    }

    pub(crate) async fn mcp_server_status_list(
        &self,
        request_id: &ConnectionRequestId,
        params: ListMcpServerStatusParams,
    ) -> Result<Option<ClientResponsePayload>, JSONRPCErrorError> {
        self.list_mcp_server_status(request_id, params)
            .await
            .map(|()| None)
    }

    pub(crate) async fn mcp_resource_read(
        &self,
        request_id: &ConnectionRequestId,
        params: McpResourceReadParams,
    ) -> Result<Option<ClientResponsePayload>, JSONRPCErrorError> {
        self.read_mcp_resource(request_id, params)
            .await
            .map(|()| None)
    }

    pub(crate) async fn mcp_server_tool_call(
        &self,
        request_id: &ConnectionRequestId,
        params: McpServerToolCallParams,
    ) -> Result<Option<ClientResponsePayload>, JSONRPCErrorError> {
        self.call_mcp_server_tool(request_id, params)
            .await
            .map(|()| None)
    }

    async fn mcp_server_refresh_response(
        &self,
        _params: Option<()>,
    ) -> Result<McpServerRefreshResponse, JSONRPCErrorError> {
        crate::mcp_refresh::reload_mcp_config(&self.thread_manager, &self.config_manager)
            .await
            .map_err(|err| internal_error(format!("failed to refresh MCP servers: {err}")))?;
        Ok(McpServerRefreshResponse {})
    }

    async fn load_global_config(&self) -> Result<Config, JSONRPCErrorError> {
        self.config_manager
            .load_non_project_config()
            .await
            .map_err(|err| internal_error(format!("failed to reload config: {err}")))
    }

    pub(super) async fn load_thread(
        &self,
        thread_id: &str,
    ) -> Result<(ThreadId, Arc<CodexThread>), JSONRPCErrorError> {
        let thread_id = ThreadId::from_string(thread_id)
            .map_err(|err| invalid_request(format!("invalid thread id: {err}")))?;

        let thread = self
            .thread_manager
            .get_thread(thread_id)
            .await
            .map_err(|_| invalid_request(format!("thread not found: {thread_id}")))?;

        Ok((thread_id, thread))
    }

    async fn mcp_server_oauth_login_response(
        &self,
        params: McpServerOauthLoginParams,
    ) -> Result<McpServerOauthLoginResponse, JSONRPCErrorError> {
        let McpServerOauthLoginParams {
            name,
            thread_id,
            client_registration,
            scopes,
            timeout_secs,
        } = params;
        let auth = self.auth_manager.auth().await;
        let (mcp_config, runtime_context) = match thread_id.as_deref() {
            Some(thread_id) => {
                let (_, thread) = self.load_thread(thread_id).await?;
                let (config, runtime_context) =
                    thread.current_mcp_config_and_runtime_context().await;
                ((*config).clone(), runtime_context)
            }
            None => {
                let config = self.load_global_config().await?;
                let mcp_config = self
                    .thread_manager
                    .mcp_manager()
                    .runtime_config(&config)
                    .await;
                let runtime_context = McpRuntimeContext::new(
                    self.thread_manager.environment_manager(),
                    config.cwd.to_path_buf(),
                );
                (mcp_config, runtime_context)
            }
        };
        let effective_servers = codex_mcp::effective_mcp_servers(&mcp_config, auth.as_ref());
        let Some(server) = effective_servers.get(&name) else {
            return Err(invalid_request(format!(
                "No MCP server named '{name}' found."
            )));
        };
        let redirect_mode = if server.is_agent_plugin() {
            StreamableHttpRedirectMode::AgentPluginV1
        } else {
            StreamableHttpRedirectMode::Legacy
        };
        let server = server.config();
        let enterprise = matches!(server.auth, McpServerAuth::EmaAuth);
        let (authorization_url, login_id, completion): (
            String,
            String,
            futures::future::BoxFuture<'static, anyhow::Result<()>>,
        ) = if enterprise {
            let thread_id = thread_id.clone().ok_or_else(|| {
                invalid_request("enterprise MCP login requires a connected thread")
            })?;
            if !mcp_config.xaa_enabled {
                return Err(invalid_request(
                    "enterprise MCP authentication requires `[features].use_xaa = true`",
                ));
            }
            if !server.enabled {
                return Err(invalid_request(format!("MCP server '{name}' is disabled.")));
            }
            server
                .validate_ema_auth_transport()
                .map_err(|reason| invalid_request(format!("MCP server '{name}' {reason}")))?;
            if client_registration.is_some() {
                return Err(invalid_request(
                    "enterprise MCP login uses the configured enterprise client and does not support clientRegistration overrides",
                ));
            }
            // Resource scopes and per-server callbacks do not alter the shared IdP registration.
            let scope = ema_auth_scope(auth.as_ref()).ok_or_else(|| {
                invalid_request(
                    "enterprise MCP authentication requires a signed-in Codex user and workspace",
                )
            })?;
            let idp = server.oauth_idp().ok_or_else(|| {
                invalid_request(format!(
                    "MCP server '{name}' has no configured enterprise IdP"
                ))
            })?;
            let registration = server.ema_registration().cloned().ok_or_else(|| {
                invalid_request(format!(
                    "MCP server '{name}' has no validated enterprise registration"
                ))
            })?;
            let credential_name = idp.credential_name(&scope);
            let http_client =
                runtime_context
                    .resolve_http_client(&name, server)
                    .map_err(|err| {
                        internal_error(format!("failed to resolve MCP server runtime: {err}"))
                    })?;
            let login = self
                .account_processor
                .enterprise_login
                .start(
                    EnterpriseOAuthLoginRequest {
                        credential_name: &credential_name,
                        issuer: &idp.issuer,
                        client_id: &idp.client_id,
                        keyring_backend_kind: mcp_config.auth_keyring_backend_kind,
                        callback_port: mcp_config.mcp_oauth_callback_port,
                        callback_url: mcp_config.mcp_oauth_callback_url.as_deref(),
                        timeout_secs,
                        http_client,
                        redirect_mode,
                    },
                    scope,
                    EnterpriseLoginTarget {
                        thread_id,
                        server_name: name.clone(),
                        registration,
                    },
                )
                .await?;
            let mut completion = login.completion;
            (
                login.authorization_url,
                login.login_id,
                Box::pin(async move {
                    match completion.wait_for(|state| state.is_complete()).await {
                        Ok(result) if *result == EnterpriseLoginCompletion::Succeeded => Ok(()),
                        _ => Err(anyhow::anyhow!("Enterprise sign-in failed.")),
                    }
                }),
            )
        } else {
            let (url, http_headers, env_http_headers) = match &server.transport {
                McpServerTransportConfig::StreamableHttp {
                    url,
                    http_headers,
                    env_http_headers,
                    ..
                } => (url.clone(), http_headers.clone(), env_http_headers.clone()),
                _ => {
                    return Err(invalid_request(
                        "OAuth login is only supported for streamable HTTP servers.",
                    ));
                }
            };

            let http_client =
                runtime_context
                    .resolve_http_client(&name, server)
                    .map_err(|err| {
                        internal_error(format!("failed to resolve MCP server runtime: {err}"))
                    })?;

            let discovered_scopes = if scopes.is_none() && server.scopes.is_none() {
                discover_supported_scopes(
                    &server.transport,
                    Arc::clone(&http_client),
                    codex_rmcp_client::OAuthDiscoveryTimeout::Requested,
                    redirect_mode,
                )
                .await
            } else {
                None
            };
            let resolved_scopes =
                resolve_oauth_scopes(scopes, server.scopes.clone(), discovered_scopes);
            let oauth_credential_name = server.oauth_credential_name(&name);
            let client_registration = match client_registration.unwrap_or_default() {
                McpServerOauthClientRegistration::Auto => McpOAuthClientRegistration::Auto,
                McpServerOauthClientRegistration::Cimd => McpOAuthClientRegistration::Cimd,
                McpServerOauthClientRegistration::Dcr => McpOAuthClientRegistration::Dcr,
            };
            let callback_url =
                resolve_oauth_callback(server, &url, mcp_config.mcp_oauth_callback_url.as_deref())
                    .map_err(|err| {
                        internal_error(format!("failed to resolve MCP OAuth callback: {err}"))
                    })?;

            let handle = perform_oauth_login_return_url(
                oauth_credential_name.as_ref(),
                &url,
                mcp_config.mcp_oauth_credentials_store_mode,
                mcp_config.auth_keyring_backend_kind,
                http_headers,
                env_http_headers,
                &resolved_scopes.scopes,
                server.oauth.as_ref(),
                client_registration,
                server.oauth_resource.as_deref(),
                timeout_secs,
                server.oauth_callback_port(mcp_config.mcp_oauth_callback_port),
                callback_url.as_deref(),
                mcp_config.mcp_oauth_callback_url.as_deref(),
                http_client,
                redirect_mode,
            )
            .await
            .map_err(|err| {
                internal_error(format!("failed to login to MCP server '{name}': {err}"))
            })?;
            (
                handle.authorization_url().to_string(),
                Uuid::now_v7().to_string(),
                Box::pin(handle.wait()),
            )
        };
        let notification_name = name.clone();
        let notification_thread_id = thread_id;
        let notification_login_id = login_id.clone();
        let outgoing = Arc::clone(&self.outgoing);
        let thread_manager = Arc::clone(&self.thread_manager);

        tokio::spawn(async move {
            let (success, error) = match completion.await {
                Ok(()) => (true, None),
                Err(err) => (false, Some(err.to_string())),
            };
            if success && !enterprise {
                thread_manager.invalidate_mcp_runtimes().await;
            }

            let notification = ServerNotification::McpServerOauthLoginCompleted(
                McpServerOauthLoginCompletedNotification {
                    name: notification_name,
                    thread_id: notification_thread_id,
                    login_id: Some(notification_login_id),
                    success,
                    error,
                },
            );
            outgoing.send_server_notification(notification).await;
        });

        Ok(McpServerOauthLoginResponse {
            authorization_url,
            login_id: Some(login_id),
        })
    }

    async fn list_mcp_server_status(
        &self,
        request_id: &ConnectionRequestId,
        params: ListMcpServerStatusParams,
    ) -> Result<(), JSONRPCErrorError> {
        let request = request_id.clone();

        let outgoing = Arc::clone(&self.outgoing);
        let thread = match params.thread_id.as_deref() {
            Some(thread_id) => Some(self.load_thread(thread_id).await?.1),
            None => None,
        };
        let processor = self.clone();
        let task = async move {
            let result = processor
                .list_mcp_server_status_response(params, thread)
                .await;
            outgoing.send_result(request, result).await;
        };
        tokio::spawn(AuthStorageOriginator::current().scope(task));
        Ok(())
    }

    async fn list_mcp_server_status_response(
        &self,
        params: ListMcpServerStatusParams,
        thread: Option<Arc<codex_core::CodexThread>>,
    ) -> Result<ListMcpServerStatusResponse, JSONRPCErrorError> {
        let detail = match params.detail.unwrap_or(McpServerStatusDetail::Full) {
            McpServerStatusDetail::Full => McpSnapshotDetail::Full,
            McpServerStatusDetail::ToolsAndAuthOnly => McpSnapshotDetail::ToolsAndAuthOnly,
        };

        let (mcp_config, snapshot) = if let (Some(thread), Some(server)) =
            (thread.as_ref(), params.server_name.as_deref())
        {
            thread
                .mcp_server_status_snapshot(server, detail)
                .await
                .map_err(|error| {
                    internal_error(format!("failed to read MCP server status: {error:#}"))
                })?
        } else {
            let config = match thread.as_ref() {
                Some(thread) => {
                    let thread_config = thread.config().await;
                    self.config_manager
                        .load_latest_config_with_session_layers(
                            &thread_config.config_layer_stack,
                            &thread_config.cwd,
                        )
                        .await
                        .map_err(|err| internal_error(format!("failed to reload config: {err}")))?
                }
                None => self.load_global_config().await?,
            };
            let mcp_manager = self.thread_manager.mcp_manager();
            let auth = self.auth_manager.auth().await;
            let (mcp_config, runtime_context) = match thread.as_ref() {
                Some(thread) => thread.runtime_mcp_config_and_context(&config).await,
                None => (
                    mcp_manager.runtime_config(&config).await,
                    McpRuntimeContext::new(
                        self.thread_manager.environment_manager(),
                        config.cwd.to_path_buf(),
                    ),
                ),
            };
            let snapshot = collect_mcp_server_status_snapshot_with_detail(
                &mcp_config,
                auth.as_ref(),
                runtime_context,
                mcp_manager.codex_apps_tools_cache(),
                mcp_manager.tool_catalog_cache(),
                detail,
                params.server_name.as_deref(),
            )
            .await;
            (Arc::new(mcp_config), snapshot)
        };

        let runtime_statuses = match thread {
            Some(thread) => thread.mcp_connection_statuses(&mcp_config).await,
            None => HashMap::new(),
        };
        let McpServerStatusSnapshot {
            mut server_infos,
            mut server_capabilities,
            mut tools_by_server,
            mut tools_errors,
            mut resources,
            mut resource_templates,
            mut auth_statuses,
            mut server_names,
        } = snapshot;
        server_names.extend(runtime_statuses.keys().cloned());
        server_names.extend(
            auth_statuses
                .keys()
                .cloned()
                .chain(resources.keys().cloned())
                .chain(resource_templates.keys().cloned()),
        );
        server_names.sort();
        server_names.dedup();
        server_names.retain(|name| {
            params
                .server_name
                .as_deref()
                .is_none_or(|server| name == server)
        });

        let total = server_names.len();
        let limit = params.limit.unwrap_or(total as u32).max(1) as usize;
        let effective_limit = limit.min(total);
        let start = match params.cursor {
            Some(cursor) => match cursor.parse::<usize>() {
                Ok(idx) => idx,
                Err(_) => return Err(invalid_request(format!("invalid cursor: {cursor}"))),
            },
            None => 0,
        };

        if start > total {
            return Err(invalid_request(format!(
                "cursor {start} exceeds total MCP servers {total}"
            )));
        }

        let end = start.saturating_add(effective_limit).min(total);

        let data: Vec<McpServerStatus> = server_names[start..end]
            .iter()
            .map(|name| McpServerStatus {
                name: name.clone(),
                runtime_status: runtime_statuses.get(name).copied().map(Into::into),
                plugin_id: mcp_config.mcp_server_catalog.server(name).and_then(
                    |server| match server.source() {
                        McpServerSource::Plugin(plugin)
                        | McpServerSource::SelectedPlugin(plugin) => {
                            Some(plugin.plugin_id().to_owned())
                        }
                        McpServerSource::Config
                        | McpServerSource::Compatibility { .. }
                        | McpServerSource::Extension { .. } => None,
                    },
                ),
                http_origin: mcp_config
                    .mcp_server_catalog
                    .server(name)
                    .and_then(|server| match &server.config().transport {
                        McpServerTransportConfig::StreamableHttp { url, .. } => {
                            url::Url::parse(url).ok().and_then(|url| {
                                matches!(url.scheme(), "http" | "https")
                                    .then(|| url.origin().ascii_serialization())
                            })
                        }
                        McpServerTransportConfig::Stdio { .. } => None,
                    }),
                server_info: server_infos.remove(name),
                server_capabilities: server_capabilities.remove(name),
                tools: tools_by_server.remove(name).unwrap_or_default(),
                tools_error: tools_errors.remove(name),
                resources: resources.remove(name).unwrap_or_default(),
                resource_templates: resource_templates.remove(name).unwrap_or_default(),
                auth_status: auth_statuses
                    .remove(name)
                    .unwrap_or(CoreMcpAuthStatus::Unsupported)
                    .into(),
            })
            .collect();

        let next_cursor = if end < total {
            Some(end.to_string())
        } else {
            None
        };

        Ok(ListMcpServerStatusResponse { data, next_cursor })
    }

    async fn read_mcp_resource(
        &self,
        request_id: &ConnectionRequestId,
        params: McpResourceReadParams,
    ) -> Result<(), JSONRPCErrorError> {
        let outgoing = Arc::clone(&self.outgoing);
        let McpResourceReadParams {
            thread_id,
            origin_call_id,
            server,
            uri,
            connector_id,
            target,
        } = params;
        let mut resource_params = ReadResourceRequestParams::new(uri);
        let mut meta = serde_json::Map::new();
        if let Some(connector_id) = connector_id {
            meta.insert(
                "x-codex-turn-metadata".to_string(),
                serde_json::json!({
                    "mcp_request_meta": {"selected_connector_ids": [connector_id]},
                }),
            );
        }
        if origin_call_id.is_none()
            && let Some(target) = target
        {
            if server != codex_mcp::CODEX_APPS_MCP_SERVER_NAME
                || target.connector_id.trim().is_empty()
                || target
                    .link_id
                    .as_ref()
                    .is_some_and(|id| id.trim().is_empty() || id.starts_with("synthetic_link::"))
            {
                return Err(invalid_request(
                    "target requires codex_apps, a connectorId, and a real linkId or null",
                ));
            }
            meta.insert(
                "connector_id".to_string(),
                serde_json::json!(target.connector_id),
            );
            meta.insert("link_id".to_string(), serde_json::json!(target.link_id));
        }
        if !meta.is_empty() {
            resource_params.meta = Some(meta.into());
        }

        if let Some(thread_id) = thread_id {
            let (_, thread) = self.load_thread(&thread_id).await?;
            let request_id = request_id.clone();

            let task = async move {
                let origin_call_id =
                    origin_call_id.filter(|_| server == codex_mcp::CODEX_APPS_MCP_SERVER_NAME);
                let result = match origin_call_id.as_deref() {
                    Some(call_id) => {
                        thread
                            .read_mcp_resource_for_call(call_id, &resource_params.uri)
                            .await
                    }
                    None => thread.read_mcp_resource(&server, resource_params).await,
                };
                Self::send_mcp_resource_read_response(outgoing, request_id, result, origin_call_id)
                    .await;
            };
            tokio::spawn(AuthStorageOriginator::current().scope(task));
            return Ok(());
        }

        if origin_call_id.is_some() {
            return Err(invalid_request("originCallId requires threadId"));
        }

        let config = self.load_global_config().await?;
        let mcp_manager = self.thread_manager.mcp_manager();
        let mcp_config = mcp_manager.runtime_config(&config).await;
        let codex_apps_tools_cache = mcp_manager.codex_apps_tools_cache();
        let tool_catalog_cache = mcp_manager.tool_catalog_cache();
        let auth = self.auth_manager.auth().await;
        let environment_manager = self.thread_manager.environment_manager();
        // This threadless resource-read path has no turn cwd or turn-selected
        // environment. Use config cwd only as the local stdio fallback; named
        // environment stdio MCPs must declare their own absolute cwd.
        let runtime_context =
            McpRuntimeContext::new(Arc::clone(&environment_manager), config.cwd.to_path_buf());
        let request_id = request_id.clone();

        let task = async move {
            let result = read_mcp_resource_without_thread(
                &mcp_config,
                auth.as_ref(),
                runtime_context,
                codex_apps_tools_cache,
                tool_catalog_cache,
                &server,
                resource_params,
            )
            .await
            .and_then(|result| serde_json::to_value(result).map_err(anyhow::Error::from));
            Self::send_mcp_resource_read_response(
                outgoing, request_id, result, /*origin_call_id*/ None,
            )
            .await;
        };
        tokio::spawn(AuthStorageOriginator::current().scope(task));
        Ok(())
    }

    async fn send_mcp_resource_read_response(
        outgoing: Arc<OutgoingMessageSender>,
        request_id: ConnectionRequestId,
        result: anyhow::Result<serde_json::Value>,
        origin_call_id: Option<String>,
    ) {
        let result = result
            .map_err(mcp_operation_error)
            .and_then(|result| {
                serde_json::from_value::<McpResourceReadResponse>(result).map_err(|error| {
                    internal_error(format!(
                        "failed to deserialize MCP resource read response: {error}"
                    ))
                })
            })
            .map(|mut response| {
                response.origin_call_id = origin_call_id;
                response
            });
        outgoing.send_result(request_id, result).await;
    }

    async fn call_mcp_server_tool(
        &self,
        request_id: &ConnectionRequestId,
        params: McpServerToolCallParams,
    ) -> Result<(), JSONRPCErrorError> {
        let outgoing = Arc::clone(&self.outgoing);
        let thread_id = params.thread_id.clone();
        let (_, thread) = self.load_thread(&thread_id).await?;
        ensure_direct_input_allowed(thread.as_ref()).await?;
        let meta = with_mcp_tool_call_thread_id_meta(params.meta, &thread_id);
        let request_id = request_id.clone();

        let task = async move {
            let result = thread
                .call_mcp_tool(&params.server, &params.tool, params.arguments, meta)
                .await
                .map(McpServerToolCallResponse::from)
                .map_err(mcp_operation_error);
            outgoing.send_result(request_id, result).await;
        };
        tokio::spawn(AuthStorageOriginator::current().scope(task));
        Ok(())
    }
}

fn mcp_operation_error(error: anyhow::Error) -> JSONRPCErrorError {
    match codex_rmcp_client::mcp_error(&error) {
        Some(error) => JSONRPCErrorError {
            code: i64::from(error.code.0),
            message: error.message.to_string(),
            data: error.data.clone(),
        },
        None => internal_error(format!("{error:#}")),
    }
}

fn with_mcp_tool_call_thread_id_meta(
    meta: Option<serde_json::Value>,
    thread_id: &str,
) -> Option<serde_json::Value> {
    match meta {
        Some(serde_json::Value::Object(mut map)) => {
            map.insert(
                MCP_TOOL_THREAD_ID_META_KEY.to_string(),
                serde_json::Value::String(thread_id.to_string()),
            );
            Some(serde_json::Value::Object(map))
        }
        None => {
            let mut map = serde_json::Map::new();
            map.insert(
                MCP_TOOL_THREAD_ID_META_KEY.to_string(),
                serde_json::Value::String(thread_id.to_string()),
            );
            Some(serde_json::Value::Object(map))
        }
        other => other,
    }
}