codex-rs/codex-mcp/src/mcp/auth.rs

use std::collections::HashMap;
use std::sync::Arc;

use anyhow::Result;
use codex_config::McpServerAuth;
use codex_config::McpServerConfig;
use codex_config::McpServerTransportConfig;
use codex_config::types::AuthKeyringBackendKind;
use codex_config::types::OAuthCredentialsStoreMode;
use codex_exec_server::HttpClient;
use codex_login::CodexAuth;
use codex_rmcp_client::McpAuthState;
use codex_rmcp_client::McpOAuthCallbackMode;
use codex_rmcp_client::OAuthDiscoveryTimeout;
use codex_rmcp_client::OAuthProviderError;
use codex_rmcp_client::StreamableHttpRedirectMode;
use codex_rmcp_client::determine_streamable_http_auth_status;
use codex_rmcp_client::determine_streamable_http_auth_status_from_credentials;
use codex_rmcp_client::discover_streamable_http_oauth;
use codex_rmcp_client::resolve_mcp_oauth_callback_url;
use futures::FutureExt;
use futures::future::join_all;
use tracing::warn;

use crate::runtime::McpRuntimeContext;
use crate::server::EffectiveMcpServer;
use crate::server::has_explicit_http_authorization;

#[derive(Debug, Clone)]
pub struct McpOAuthLoginConfig {
    pub url: String,
    pub http_headers: Option<HashMap<String, String>>,
    pub env_http_headers: Option<HashMap<String, String>>,
    pub discovered_scopes: Option<Vec<String>>,
    pub callback_mode: McpOAuthCallbackMode,
}

#[derive(Debug)]
pub enum McpOAuthLoginSupport {
    Supported(McpOAuthLoginConfig),
    Unsupported,
    Unknown(anyhow::Error),
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum McpOAuthScopesSource {
    Explicit,
    Configured,
    Discovered,
    Empty,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ResolvedMcpOAuthScopes {
    pub scopes: Vec<String>,
    pub source: McpOAuthScopesSource,
}

/// Keeps registered callbacks tied to their client while preserving legacy redirects.
pub fn resolve_oauth_callback(
    server: &McpServerConfig,
    server_url: &str,
    global_callback_url: Option<&str>,
) -> Result<Option<String>> {
    if let Some(callback_url) = server
        .oauth
        .as_ref()
        .and_then(|oauth| oauth.callback_url.as_deref())
    {
        return Ok(Some(callback_url.to_string()));
    }

    if server
        .oauth_client_id()
        .is_none_or(|client_id| client_id.trim().is_empty())
    {
        return Ok(global_callback_url.map(ToOwned::to_owned));
    }

    resolve_mcp_oauth_callback_url(
        server_url,
        global_callback_url,
        McpOAuthCallbackMode::CallbackSpecific,
    )
    .map(Some)
}

#[derive(Debug, Clone)]
pub struct McpAuthStatusEntry {
    pub config: Option<McpServerConfig>,
    pub auth_state: McpAuthState,
}

pub async fn oauth_login_support(
    transport: &McpServerTransportConfig,
    http_client: Arc<dyn HttpClient>,
    discovery_timeout: OAuthDiscoveryTimeout,
    redirect_mode: StreamableHttpRedirectMode,
) -> McpOAuthLoginSupport {
    let Some(mut config) = oauth_login_candidate(transport) else {
        return McpOAuthLoginSupport::Unsupported;
    };
    match discover_streamable_http_oauth(
        &config.url,
        config.http_headers.clone(),
        config.env_http_headers.clone(),
        http_client,
        discovery_timeout,
        redirect_mode,
    )
    .await
    {
        Ok(Some(discovery)) => {
            config.discovered_scopes = discovery.scopes_supported;
            config.callback_mode = discovery.callback_mode;
            McpOAuthLoginSupport::Supported(config)
        }
        Ok(None) => McpOAuthLoginSupport::Unsupported,
        Err(err) => McpOAuthLoginSupport::Unknown(err),
    }
}

fn oauth_login_candidate(transport: &McpServerTransportConfig) -> Option<McpOAuthLoginConfig> {
    let McpServerTransportConfig::StreamableHttp {
        url,
        bearer_token_env_var,
        http_headers,
        env_http_headers,
        ..
    } = transport
    else {
        return None;
    };
    if bearer_token_env_var.is_some() {
        return None;
    }
    Some(McpOAuthLoginConfig {
        url: url.clone(),
        http_headers: http_headers.clone(),
        env_http_headers: env_http_headers.clone(),
        discovered_scopes: None,
        callback_mode: McpOAuthCallbackMode::CallbackSpecific,
    })
}

pub async fn discover_supported_scopes(
    transport: &McpServerTransportConfig,
    http_client: Arc<dyn HttpClient>,
    discovery_timeout: OAuthDiscoveryTimeout,
    redirect_mode: StreamableHttpRedirectMode,
) -> Option<Vec<String>> {
    match oauth_login_support(transport, http_client, discovery_timeout, redirect_mode).await {
        McpOAuthLoginSupport::Supported(config) => config.discovered_scopes,
        McpOAuthLoginSupport::Unsupported | McpOAuthLoginSupport::Unknown(_) => None,
    }
}

pub fn resolve_oauth_scopes(
    explicit_scopes: Option<Vec<String>>,
    configured_scopes: Option<Vec<String>>,
    discovered_scopes: Option<Vec<String>>,
) -> ResolvedMcpOAuthScopes {
    if let Some(scopes) = explicit_scopes {
        return ResolvedMcpOAuthScopes {
            scopes,
            source: McpOAuthScopesSource::Explicit,
        };
    }

    if let Some(scopes) = configured_scopes {
        return ResolvedMcpOAuthScopes {
            scopes,
            source: McpOAuthScopesSource::Configured,
        };
    }

    if let Some(scopes) = discovered_scopes
        && !scopes.is_empty()
    {
        return ResolvedMcpOAuthScopes {
            scopes,
            source: McpOAuthScopesSource::Discovered,
        };
    }

    ResolvedMcpOAuthScopes {
        scopes: Vec::new(),
        source: McpOAuthScopesSource::Empty,
    }
}

pub fn should_retry_without_scopes(scopes: &ResolvedMcpOAuthScopes, error: &anyhow::Error) -> bool {
    scopes.source == McpOAuthScopesSource::Discovered
        && error.downcast_ref::<OAuthProviderError>().is_some()
}

pub async fn compute_auth_statuses<'a, I>(
    servers: I,
    store_mode: OAuthCredentialsStoreMode,
    keyring_backend_kind: AuthKeyringBackendKind,
    auth: Option<&CodexAuth>,
    runtime_context: &McpRuntimeContext,
) -> HashMap<String, McpAuthStatusEntry>
where
    I: IntoIterator<Item = (&'a String, &'a EffectiveMcpServer)>,
{
    let futures = servers.into_iter().map(|(name, server)| {
        let name = name.clone();
        let redirect_mode = if server.is_agent_plugin() {
            StreamableHttpRedirectMode::AgentPluginV1
        } else {
            StreamableHttpRedirectMode::Legacy
        };
        let config = server.config().clone();
        let runtime_context = runtime_context.clone();
        async move {
            let auth_state = match compute_auth_status(
                &name,
                &config,
                store_mode,
                keyring_backend_kind,
                auth,
                &runtime_context,
                redirect_mode,
            )
            .await
            {
                Ok(status) => status,
                Err(error) => {
                    warn!("failed to determine auth status for MCP server `{name}`: {error:?}");
                    McpAuthState::Unknown
                }
            };
            let entry = McpAuthStatusEntry {
                config: Some(config),
                auth_state,
            };
            (name, entry)
        }
    });

    join_all(futures).await.into_iter().collect()
}

async fn compute_auth_status(
    server_name: &str,
    config: &McpServerConfig,
    store_mode: OAuthCredentialsStoreMode,
    keyring_backend_kind: AuthKeyringBackendKind,
    auth: Option<&CodexAuth>,
    runtime_context: &McpRuntimeContext,
    redirect_mode: StreamableHttpRedirectMode,
) -> Result<McpAuthState> {
    if !config.enabled {
        return Ok(McpAuthState::Unsupported);
    }
    if matches!(config.auth, McpServerAuth::EmaAuth) {
        // EMA connections are not enabled until the runtime stage of the stack.
        return Ok(McpAuthState::Unsupported);
    }
    let has_runtime_auth = matches!(config.auth, McpServerAuth::ChatGpt)
        && auth.is_some_and(CodexAuth::uses_codex_backend)
        && matches!(
            &config.transport,
            McpServerTransportConfig::StreamableHttp {
                bearer_token_env_var: None,
                ..
            }
        );

    if matches!(config.auth, McpServerAuth::ChatGpt) && !config.is_local_environment() {
        return Ok(if has_explicit_http_authorization(config) {
            McpAuthState::BearerToken
        } else {
            McpAuthState::Unsupported
        });
    }

    if has_runtime_auth {
        return Ok(McpAuthState::BearerToken);
    }

    match &config.transport {
        McpServerTransportConfig::Stdio { .. } => Ok(McpAuthState::Unsupported),
        McpServerTransportConfig::StreamableHttp {
            url,
            bearer_token_env_var,
            http_headers,
            env_http_headers,
            http_headers_helper,
        } => {
            if http_headers_helper.is_some() {
                // Status inspection must not execute an arbitrary local helper. Existing
                // credentials remain reportable; otherwise discovery waits for startup/login.
                return Ok(determine_streamable_http_auth_status_from_credentials(
                    config.oauth_credential_name(server_name).as_ref(),
                    url,
                    bearer_token_env_var.as_deref(),
                    http_headers.clone(),
                    env_http_headers.clone(),
                    store_mode,
                    keyring_backend_kind,
                )?
                .unwrap_or(McpAuthState::Unknown));
            }
            let http_client = runtime_context
                .resolve_http_client(server_name, config)
                .map_err(anyhow::Error::msg)?;
            let discovery_timeout = if config.is_local_environment() {
                OAuthDiscoveryTimeout::LOCAL
            } else {
                OAuthDiscoveryTimeout::Requested
            };
            let oauth_credential_name = config.oauth_credential_name(server_name);
            determine_streamable_http_auth_status(
                oauth_credential_name.as_ref(),
                url,
                bearer_token_env_var.as_deref(),
                http_headers.clone(),
                env_http_headers.clone(),
                store_mode,
                keyring_backend_kind,
                http_client,
                discovery_timeout,
                redirect_mode,
            )
            .boxed()
            .await
        }
    }
}

#[cfg(test)]
mod tests {
    use anyhow::anyhow;
    use pretty_assertions::assert_eq;

    use super::McpOAuthScopesSource;
    use super::OAuthProviderError;
    use super::ResolvedMcpOAuthScopes;
    use super::resolve_oauth_callback;
    use super::resolve_oauth_scopes;
    use super::should_retry_without_scopes;

    #[test]
    fn callback_resolution_preserves_registered_and_legacy_clients() -> anyhow::Result<()> {
        for (client_id, saved_callback, global_callback, expected_callback) in [
            (
                Some("registered-client"),
                Some("http://127.0.0.1/callback"),
                Some("https://override.example/callback"),
                Some("http://127.0.0.1/callback"),
            ),
            (
                Some("legacy-client"),
                None,
                None,
                Some("http://127.0.0.1/callback/epMNJ6P1xGQ9"),
            ),
            (
                None,
                Some("https://plugin.example/callback"),
                Some("https://override.example/callback"),
                Some("https://plugin.example/callback"),
            ),
            (
                None,
                None,
                Some("https://override.example/callback"),
                Some("https://override.example/callback"),
            ),
        ] {
            let server = serde_json::from_value(serde_json::json!({
                "url": "https://mcp.example.com/mcp",
                "oauth": {
                    "client_id": client_id,
                    "callback_url": saved_callback,
                },
            }))?;
            assert_eq!(
                resolve_oauth_callback(&server, "https://mcp.example.com/mcp", global_callback)?
                    .as_deref(),
                expected_callback
            );
        }

        Ok(())
    }

    #[test]
    fn resolve_oauth_scopes_prefers_explicit() {
        let resolved = resolve_oauth_scopes(
            Some(vec!["explicit".to_string()]),
            Some(vec!["configured".to_string()]),
            Some(vec!["discovered".to_string()]),
        );

        assert_eq!(
            resolved,
            ResolvedMcpOAuthScopes {
                scopes: vec!["explicit".to_string()],
                source: McpOAuthScopesSource::Explicit,
            }
        );
    }

    #[test]
    fn resolve_oauth_scopes_prefers_configured_over_discovered() {
        let resolved = resolve_oauth_scopes(
            /*explicit_scopes*/ None,
            Some(vec!["configured".to_string()]),
            Some(vec!["discovered".to_string()]),
        );

        assert_eq!(
            resolved,
            ResolvedMcpOAuthScopes {
                scopes: vec!["configured".to_string()],
                source: McpOAuthScopesSource::Configured,
            }
        );
    }

    #[test]
    fn resolve_oauth_scopes_uses_discovered_when_needed() {
        let resolved = resolve_oauth_scopes(
            /*explicit_scopes*/ None,
            /*configured_scopes*/ None,
            Some(vec!["discovered".to_string()]),
        );

        assert_eq!(
            resolved,
            ResolvedMcpOAuthScopes {
                scopes: vec!["discovered".to_string()],
                source: McpOAuthScopesSource::Discovered,
            }
        );
    }

    #[test]
    fn resolve_oauth_scopes_preserves_explicitly_empty_configured_scopes() {
        let resolved = resolve_oauth_scopes(
            /*explicit_scopes*/ None,
            Some(Vec::new()),
            Some(vec!["ignored".into()]),
        );

        assert_eq!(
            resolved,
            ResolvedMcpOAuthScopes {
                scopes: Vec::new(),
                source: McpOAuthScopesSource::Configured,
            }
        );
    }

    #[test]
    fn resolve_oauth_scopes_falls_back_to_empty() {
        let resolved = resolve_oauth_scopes(
            /*explicit_scopes*/ None, /*configured_scopes*/ None,
            /*discovered_scopes*/ None,
        );

        assert_eq!(
            resolved,
            ResolvedMcpOAuthScopes {
                scopes: Vec::new(),
                source: McpOAuthScopesSource::Empty,
            }
        );
    }

    #[test]
    fn should_retry_without_scopes_only_for_discovered_provider_errors() {
        let discovered = ResolvedMcpOAuthScopes {
            scopes: vec!["scope".to_string()],
            source: McpOAuthScopesSource::Discovered,
        };
        let provider_error = anyhow!(OAuthProviderError::new(
            Some("invalid_scope".to_string()),
            Some("scope rejected".to_string()),
        ));

        assert!(should_retry_without_scopes(&discovered, &provider_error));

        let configured = ResolvedMcpOAuthScopes {
            scopes: vec!["scope".to_string()],
            source: McpOAuthScopesSource::Configured,
        };
        assert!(!should_retry_without_scopes(&configured, &provider_error));
        assert!(!should_retry_without_scopes(
            &discovered,
            &anyhow!("timed out waiting for OAuth callback"),
        ));
    }
}