ai_agent/services/mcp/
auth.rs

1// Source: /data/home/swei/claudecode/openclaudecode/src/types/generated/events_mono/common/v1/auth.ts
2//! MCP Authentication module
3//!
4//! Provides OAuth flow support for MCP servers that require authentication.
5//! SDK users register an OAuth callback via `register_mcp_oauth_callback()`.
6
7use serde::{Deserialize, Serialize};
8use std::sync::Arc;
9
10/// MCP authentication config
11#[derive(Debug, Clone, Default, Serialize, Deserialize)]
12pub struct AuthConfig {
13    pub enabled: bool,
14    pub auth_type: Option<String>,
15    pub token: Option<String>,
16}
17
18/// MCP OAuth configuration
19#[derive(Debug, Clone, Default, Serialize, Deserialize)]
20pub struct OAuthConfig {
21    pub client_id: Option<String>,
22    pub client_secret: Option<String>,
23    pub redirect_uri: Option<String>,
24    pub scopes: Vec<String>,
25}
26
27/// Get MCP auth headers
28pub fn get_auth_headers(config: &AuthConfig) -> std::collections::HashMap<String, String> {
29    let mut headers = std::collections::HashMap::new();
30
31    if let Some(token) = &config.token {
32        headers.insert("Authorization".to_string(), format!("Bearer {}", token));
33    }
34
35    headers
36}
37
38/// Check if auth is required for an MCP server
39pub fn is_auth_required(config: &AuthConfig) -> bool {
40    config.enabled && config.auth_type.is_some()
41}
42
43// ============================================================================
44// MCP OAuth flow — matches TypeScript performMCPOAuthFlow
45// ============================================================================
46
47/// Result of starting the MCP OAuth flow
48#[derive(Debug, Clone)]
49pub struct McpOAuthResult {
50    /// Status of the auth flow
51    pub status: McpOAuthStatus,
52    /// Human-readable message for the user
53    pub message: String,
54    /// Authorization URL to share with the user (if status is AuthUrl)
55    pub auth_url: Option<String>,
56}
57
58/// Status of the MCP OAuth flow
59#[derive(Debug, Clone, PartialEq)]
60pub enum McpOAuthStatus {
61    /// An authorization URL was returned; the user needs to open it
62    AuthUrl,
63    /// The auth completed silently (e.g., cached token)
64    Authenticated,
65    /// OAuth is not supported for this transport
66    Unsupported,
67    /// An error occurred during the flow
68    Error,
69}
70
71/// Callback type for performing MCP OAuth flow.
72/// Takes (server_name, config_json, on_auth_url_callback).
73pub type McpOAuthCallback = Arc<
74    dyn Fn(
75        String,
76        serde_json::Value,
77        Option<Arc<dyn Fn(String) + Send + Sync>>,
78    ) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<McpOAuthResult, crate::AgentError>> + Send + Sync>>
79        + Send
80        + Sync,
81>;
82
83/// Global OAuth callback registered by the SDK user.
84static MCP_OAUTH_CALLBACK: once_cell::sync::Lazy<parking_lot::RwLock<Option<McpOAuthCallback>>> =
85    once_cell::sync::Lazy::new(Default::default);
86
87/// Register a callback for MCP OAuth flow.
88///
89/// The callback is invoked when a server that requires OAuth is used.
90/// It should start the OAuth flow and return the authorization URL.
91pub fn register_mcp_oauth_callback<F, Fut>(callback: F)
92where
93    F: Fn(
94        String,
95        serde_json::Value,
96        Option<Arc<dyn Fn(String) + Send + Sync>>,
97    ) -> Fut + Send + Sync + 'static,
98    Fut: std::future::Future<Output = Result<McpOAuthResult, crate::AgentError>> + Send + Sync + 'static,
99{
100    let wrapped: McpOAuthCallback = Arc::new(
101        move |server: String, config: serde_json::Value, on_url: Option<Arc<dyn Fn(String) + Send + Sync>>| {
102            Box::pin(callback(server, config, on_url))
103        },
104    );
105    *MCP_OAUTH_CALLBACK.write() = Some(wrapped);
106}
107
108/// Execute the MCP OAuth flow for a server.
109///
110/// Returns an error if no OAuth callback has been registered.
111pub async fn perform_mcp_oauth_flow(
112    server_name: String,
113    config: serde_json::Value,
114    on_auth_url: Option<Arc<dyn Fn(String) + Send + Sync>>,
115) -> Result<McpOAuthResult, crate::AgentError> {
116    let callback = MCP_OAUTH_CALLBACK.read().clone();
117    match callback {
118        Some(cb) => cb(server_name, config, on_auth_url).await,
119        None => Err(crate::AgentError::Tool(
120            "No MCP OAuth callback registered. Call register_mcp_oauth_callback() to enable OAuth.".to_string(),
121        )),
122    }
123}
124
125// Re-export for use by McpAuthTool (accessed via crate::services::mcp::client::clear_mcp_auth_cache)
ai_agent/services/mcp/auth.rs

ai_agent/services/mcp/
auth.rs
