diff --git a/Cargo.lock b/Cargo.lock index bc65a50a..394feaff 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1102,6 +1102,12 @@ version = "0.22.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" +[[package]] +name = "base64" +version = "0.23.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac07cdecf99051d9a5238b80f35af32cdeba5b336e55d957b318b50137e18da5" + [[package]] name = "bincode" version = "1.3.3" @@ -1798,6 +1804,7 @@ dependencies = [ "mcp_client", "md5", "notify", + "open", "pty_session", "regex", "rust-embed", @@ -1812,6 +1819,7 @@ dependencies = [ "tools_core", "tracing", "transmutation", + "urlencoding", "web", ] @@ -6245,7 +6253,7 @@ dependencies = [ "futures", "futures-util", "keyring", - "oauth2", + "oauth2 4.4.2", "rand 0.8.6", "regex", "reqwest 0.11.27", @@ -6595,9 +6603,12 @@ dependencies = [ "axum 0.8.9", "command_executor", "http 1.4.0", + "reqwest 0.13.4", "rmcp", "serde", "serde_json", + "tempfile", + "thiserror 2.0.18", "tokio", "tools_core", "tracing", @@ -7134,6 +7145,25 @@ dependencies = [ "url", ] +[[package]] +name = "oauth2" +version = "5.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "51e219e79014df21a225b1860a479e2dcd7cbd9130f4defd4bd0e191ea31d67d" +dependencies = [ + "base64 0.21.7", + "chrono", + "getrandom 0.2.17", + "http 1.4.0", + "rand 0.8.6", + "serde", + "serde_json", + "serde_path_to_error", + "sha2 0.10.9", + "thiserror 1.0.69", + "url", +] + [[package]] name = "objc" version = "0.2.7" @@ -8257,9 +8287,9 @@ dependencies = [ [[package]] name = "process-wrap" -version = "9.1.0" +version = "10.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2e842efad9119158434d193c6682e2ebee4b44d6ad801d7b349623b3f57cdf55" +checksum = "0e3f4237d0e4741eb50bc5584db701f1299c85fa31ff0274dd6445e79dc42d12" dependencies = [ "futures", "indexmap 2.14.0", @@ -8998,18 +9028,20 @@ dependencies = [ [[package]] name = "rmcp" -version = "2.1.0" +version = "3.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f00a32c3b81b7b254076a65abd5ab2551209146713ba38f73818657e865e9433" +checksum = "b88db56b8ae316560e9e868b6b978ea940f27cb883323fc90e07435e44158f5c" dependencies = [ "async-trait", - "base64 0.22.1", + "base64 0.23.1", "bytes", "chrono", "futures", "http 1.4.0", "http-body 1.0.1", "http-body-util", + "indexmap 2.14.0", + "oauth2 5.0.0", "pastey 0.2.3", "pin-project-lite", "process-wrap", @@ -9025,6 +9057,7 @@ dependencies = [ "tokio-util", "tower-service", "tracing", + "url", "uuid", ] @@ -10219,9 +10252,9 @@ dependencies = [ [[package]] name = "sse-stream" -version = "0.2.1" +version = "0.2.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "eb4dc4d33c68ec1f27d386b5610a351922656e1fdf5c05bbaad930cd1519479a" +checksum = "c25ac7aff0abd1dbc474536e40416e1102c7dd9bfba0b9861c6d357f835dcfb4" dependencies = [ "bytes", "futures-util", diff --git a/crates/code_assistant/src/cli.rs b/crates/code_assistant/src/cli.rs index fc3ae74f..afffa6de 100644 --- a/crates/code_assistant/src/cli.rs +++ b/crates/code_assistant/src/cli.rs @@ -33,6 +33,18 @@ pub enum Mode { /// Show ChatGPT subscription auth status CodexStatus, + /// Log in to an HTTP MCP server that requires OAuth (opens browser) + McpLogin { + /// Name of the server as configured in mcp-servers.json + server: String, + }, + + /// Remove stored OAuth tokens for an HTTP MCP server + McpLogout { + /// Name of the server as configured in mcp-servers.json + server: String, + }, + /// Run as ACP (Agent Client Protocol) agent Acp { /// Enable verbose logging diff --git a/crates/code_assistant/src/main.rs b/crates/code_assistant/src/main.rs index 327f39ff..3799bf16 100644 --- a/crates/code_assistant/src/main.rs +++ b/crates/code_assistant/src/main.rs @@ -2,6 +2,7 @@ mod app; mod cli; mod codex_commands; mod logging; +mod mcp_commands; // The domain layer lives in `code_assistant_core`; re-exported under the // historical module paths so call sites keep using `crate::session::…` etc. @@ -46,6 +47,13 @@ async fn main() -> Result<()> { Some(Mode::CodexStatus) => { return codex_commands::run_codex_status(); } + Some(Mode::McpLogin { server }) => { + setup_logging(1, true); + return mcp_commands::run_mcp_login(&server).await; + } + Some(Mode::McpLogout { server }) => { + return mcp_commands::run_mcp_logout(&server); + } Some(Mode::Server { verbose }) => { #[cfg(feature = "mcp-server")] { diff --git a/crates/code_assistant/src/mcp_commands.rs b/crates/code_assistant/src/mcp_commands.rs new file mode 100644 index 00000000..0af12956 --- /dev/null +++ b/crates/code_assistant/src/mcp_commands.rs @@ -0,0 +1,32 @@ +//! CLI commands for authenticating HTTP MCP servers that require OAuth +//! (MCP's authorization spec — e.g. servers that answer `initialize` with +//! `401 Auth required`). +//! +//! `mcp-login ` runs the browser OAuth flow and stores the resulting +//! tokens under `/mcp-oauth/.json`; subsequent agent runs +//! reuse them silently. This mirrors what Cline's "Authenticate" button does. +//! The flow itself lives in `code_assistant_core::tools::mcp_auth` so the +//! settings UI can drive the same login. + +use anyhow::Result; +use code_assistant_core::tools::{mcp, mcp_auth}; + +/// Run the OAuth browser login for the configured HTTP MCP server `server`. +pub async fn run_mcp_login(server: &str) -> Result<()> { + println!("Starting OAuth login for MCP server '{server}'..."); + mcp_auth::login_mcp_server(server).await?; + println!(); + println!("Login successful. Tokens stored for MCP server '{server}'."); + println!("The next agent run will use them automatically."); + Ok(()) +} + +/// Remove any stored OAuth tokens for the MCP server `server`. +pub fn run_mcp_logout(server: &str) -> Result<()> { + if mcp::forget_mcp_oauth_tokens(server)? { + println!("Logged out. Removed stored OAuth tokens for MCP server '{server}'."); + } else { + println!("No stored OAuth tokens for MCP server '{server}'."); + } + Ok(()) +} diff --git a/crates/code_assistant_core/Cargo.toml b/crates/code_assistant_core/Cargo.toml index d888a9bb..e0dc9525 100644 --- a/crates/code_assistant_core/Cargo.toml +++ b/crates/code_assistant_core/Cargo.toml @@ -61,6 +61,10 @@ dirs = "5.0" md5 = "0.7.0" async-channel = "2.5.0" +# MCP OAuth interactive login: open the browser + decode the loopback redirect +open = "5" +urlencoding = "2" + # Base64 encoding for images base64 = "0.22" diff --git a/crates/code_assistant_core/src/tools/mcp.rs b/crates/code_assistant_core/src/tools/mcp.rs index b9720ea1..1286b12c 100644 --- a/crates/code_assistant_core/src/tools/mcp.rs +++ b/crates/code_assistant_core/src/tools/mcp.rs @@ -10,8 +10,8 @@ use anyhow::{Context, Result}; use std::path::{Path, PathBuf}; pub use mcp_client::{ - DiscoveredTool, McpServerConfig, McpServerStatus, McpServersConfig, McpTransport, - discover_tools, parse_local_mcp_json, + AuthorizationOutcome, AuthorizationRequired, DiscoveredTool, McpServerConfig, McpServerStatus, + McpServersConfig, McpTransport, OAuthAuthorizer, discover_tools, parse_local_mcp_json, }; /// Scope tags every MCP tool carries in code-assistant: offered to the main @@ -24,6 +24,120 @@ pub fn mcp_servers_config_path() -> PathBuf { crate::config_dir::config_dir().join("mcp-servers.json") } +/// Directory holding persisted OAuth tokens for HTTP MCP servers — one JSON +/// file per server, written by the interactive login and reused silently on +/// later connects. +pub fn mcp_oauth_dir() -> PathBuf { + crate::config_dir::config_dir().join("mcp-oauth") +} + +/// The OAuth token file for `server`. The server name is sanitized to a safe +/// file stem so an unusual name cannot escape [`mcp_oauth_dir`]. +pub fn mcp_oauth_token_path(server: &str) -> PathBuf { + let stem: String = server + .chars() + .map(|c| if c.is_ascii_alphanumeric() { c } else { '_' }) + .collect(); + mcp_oauth_dir().join(format!("{stem}.json")) +} + +/// A fingerprint of the persisted OAuth tokens (file names + sizes + modified +/// times, sorted). Included in the tool-registry fingerprint so that +/// completing an interactive login — which only writes a token file, none of +/// the config files — invalidates the cached registry and the next agent run +/// rebuilds it, reconnecting the now-authorized server so its tools appear +/// without restarting the app. +pub fn mcp_oauth_fingerprint() -> String { + let dir = mcp_oauth_dir(); + let mut entries: Vec = match std::fs::read_dir(&dir) { + Ok(read_dir) => read_dir + .filter_map(|entry| entry.ok()) + .map(|entry| { + let name = entry.file_name().to_string_lossy().into_owned(); + let (len, modified) = entry + .metadata() + .map(|meta| { + let modified = meta + .modified() + .ok() + .and_then(|time| time.duration_since(std::time::UNIX_EPOCH).ok()) + .map(|d| d.as_nanos()) + .unwrap_or_default(); + (meta.len(), modified) + }) + .unwrap_or_default(); + format!("{name}:{len}:{modified}") + }) + .collect(), + Err(_) => Vec::new(), + }; + entries.sort(); + entries.join("\u{0}") +} + +/// A persistent credential store for one HTTP MCP server's OAuth tokens, +/// backed by a file under [`mcp_oauth_dir`]. Passed to +/// [`mcp_client::McpServerConnection::connect`] so a stored token is reused +/// without user interaction, and to [`mcp_client::authenticate_http_server`] +/// so the interactive login persists. +pub fn mcp_oauth_credential_store(server: &str) -> std::sync::Arc { + std::sync::Arc::new(mcp_client::FileCredentialStore::new(mcp_oauth_token_path( + server, + ))) +} + +/// Run the interactive OAuth login for a configured HTTP MCP server and +/// persist its tokens, so later connects reuse them without user interaction. +/// `authorizer` performs the browser round-trip (open the URL, capture the +/// redirect). The server is looked up in the global `mcp-servers.json`; +/// errors if it is unknown or not an HTTP server. +pub async fn authenticate_mcp_server(name: &str, authorizer: &dyn OAuthAuthorizer) -> Result<()> { + let config = load_mcp_servers_config()?; + let server = config.servers.get(name).ok_or_else(|| { + anyhow::anyhow!( + "No MCP server named '{name}' in {}", + mcp_servers_config_path().display() + ) + })?; + let store = mcp_oauth_credential_store(name); + mcp_client::authenticate_http_server(name, server, store, authorizer, "code-assistant").await +} + +/// Forget any stored OAuth tokens for `server` (e.g. to force a fresh login). +/// Returns whether a token file was present. +pub fn forget_mcp_oauth_tokens(server: &str) -> Result { + let path = mcp_oauth_token_path(server); + if path.exists() { + std::fs::remove_file(&path) + .with_context(|| format!("Failed to remove {}", path.display()))?; + Ok(true) + } else { + Ok(false) + } +} + +/// Whether a persisted OAuth token file exists for `server`. A coarse "have we +/// logged in" signal for status output and UI; it does not check expiry. +pub fn has_mcp_oauth_tokens(server: &str) -> bool { + mcp_oauth_token_path(server).exists() +} + +/// Discover a server's tools for a configuration UI, reusing stored OAuth +/// tokens for HTTP servers so an authenticated server lists its tools. An +/// HTTP server that still needs authorization fails with +/// [`AuthorizationRequired`] (downcastable), which the UI turns into an +/// "Authenticate" action. +pub async fn discover_server_tools( + server_name: &str, + server: &McpServerConfig, +) -> Result> { + let credentials = server + .transport + .is_http() + .then(|| mcp_oauth_credential_store(server_name)); + mcp_client::discover_tools(server_name, server, credentials).await +} + /// Load the MCP servers configuration, substituting `${ENV_VAR}` patterns in /// server environment values. A server whose variables cannot be resolved is /// skipped with a log warning (it could not connect anyway); a missing file diff --git a/crates/code_assistant_core/src/tools/mcp_auth.rs b/crates/code_assistant_core/src/tools/mcp_auth.rs new file mode 100644 index 00000000..9809c049 --- /dev/null +++ b/crates/code_assistant_core/src/tools/mcp_auth.rs @@ -0,0 +1,177 @@ +//! Interactive OAuth login for HTTP MCP servers (MCP's authorization spec). +//! +//! [`LoopbackAuthorizer`] implements the [`OAuthAuthorizer`] seam with the +//! RFC 8252 native-app pattern: a one-shot loopback HTTP server receives the +//! OAuth redirect while the user's browser is opened to the authorization +//! URL. [`login_mcp_server`] wires it to a configured server and persists the +//! resulting tokens, so later agent runs connect silently. +//! +//! Shared by the CLI (`mcp-login`) and the settings UI so both drive the same +//! flow. + +use super::mcp::{self, AuthorizationOutcome, OAuthAuthorizer}; +use anyhow::{Context, Result}; +use std::sync::Mutex; +use std::time::Duration; +use tokio::net::TcpListener; + +/// An [`OAuthAuthorizer`] that runs a one-shot loopback HTTP server for the +/// OAuth redirect and opens the user's browser to the authorization URL. +pub struct LoopbackAuthorizer { + redirect_uri: String, + /// Bound at construction so [`Self::redirect_uri`] reflects the real port; + /// consumed by the single `authorize` call. + listener: Mutex>, +} + +impl LoopbackAuthorizer { + /// Bind an ephemeral loopback port for the redirect callback. + pub async fn bind() -> Result { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .context("binding the OAuth callback server on loopback")?; + let port = listener.local_addr()?.port(); + Ok(Self { + redirect_uri: format!("http://127.0.0.1:{port}/callback"), + listener: Mutex::new(Some(listener)), + }) + } +} + +#[async_trait::async_trait] +impl OAuthAuthorizer for LoopbackAuthorizer { + fn redirect_uri(&self) -> String { + self.redirect_uri.clone() + } + + async fn authorize(&self, authorization_url: String) -> Result { + let listener = self + .listener + .lock() + .expect("authorizer mutex poisoned") + .take() + .context("the loopback authorizer can only authorize once")?; + + // Log and print the URL so it is reachable even if the browser does + // not open (headless, or a GPUI settings screen without a console). + tracing::info!("MCP OAuth authorization URL: {authorization_url}"); + eprintln!("Authorize this MCP server in your browser:\n {authorization_url}"); + if let Err(e) = open::that(&authorization_url) { + tracing::warn!("Could not open the browser automatically: {e}"); + } + + let (mut stream, _) = tokio::time::timeout(Duration::from_secs(300), listener.accept()) + .await + .context("authorization timed out after 5 minutes")? + .context("accepting the OAuth callback connection")?; + + let params = read_callback_query(&mut stream).await?; + + if let Some(error) = params.get("error") { + let description = params.get("error_description").cloned().unwrap_or_default(); + respond( + &mut stream, + "400 Bad Request", + "Authorization failed", + "You can close this window and return to the app.", + ) + .await; + anyhow::bail!("authorization server returned an error: {error} {description}"); + } + + let code = params + .get("code") + .cloned() + .context("the callback did not include an authorization code")?; + let state = params.get("state").cloned().unwrap_or_default(); + let issuer = params.get("iss").cloned(); + + respond( + &mut stream, + "200 OK", + "Authorization complete", + "You can close this window and return to the app.", + ) + .await; + + Ok(AuthorizationOutcome { + code, + state, + issuer, + }) + } +} + +/// Run the OAuth browser login for the configured HTTP MCP server `name` and +/// persist its tokens. A convenience wrapper binding a [`LoopbackAuthorizer`] +/// to [`mcp::authenticate_mcp_server`]. +pub async fn login_mcp_server(name: &str) -> Result<()> { + let authorizer = LoopbackAuthorizer::bind().await?; + mcp::authenticate_mcp_server(name, &authorizer).await +} + +/// Read the redirect request off `stream` and return its decoded query +/// parameters. +async fn read_callback_query( + stream: &mut tokio::net::TcpStream, +) -> Result> { + use tokio::io::AsyncReadExt; + + let mut buf = vec![0u8; 8192]; + let n = tokio::time::timeout(Duration::from_secs(10), stream.read(&mut buf)) + .await + .context("timeout reading the OAuth callback request")? + .context("reading the OAuth callback request")?; + let request = String::from_utf8_lossy(&buf[..n]); + + let request_line = request.lines().next().unwrap_or(""); + let path = request_line.split_whitespace().nth(1).unwrap_or("/"); + let query = path.split('?').nth(1).unwrap_or(""); + + Ok(query + .split('&') + .filter(|pair| !pair.is_empty()) + .filter_map(|pair| { + let mut parts = pair.splitn(2, '='); + let key = parts.next()?; + let value = parts.next().unwrap_or(""); + Some(( + urlencoding::decode(key).ok()?.into_owned(), + urlencoding::decode(value).ok()?.into_owned(), + )) + }) + .collect()) +} + +/// Write a minimal HTML response and close the connection. +async fn respond(stream: &mut tokio::net::TcpStream, status: &str, title: &str, body: &str) { + use tokio::io::AsyncWriteExt; + + let html = format!( + "
\ +

{title}

{body}

" + ); + let response = format!( + "HTTP/1.1 {status}\r\nContent-Type: text/html\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{html}", + html.len() + ); + let _ = stream.write_all(response.as_bytes()).await; + let _ = stream.shutdown().await; +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn bind_yields_a_loopback_redirect_uri() { + let authorizer = LoopbackAuthorizer::bind().await.unwrap(); + let uri = authorizer.redirect_uri(); + assert!( + uri.starts_with("http://127.0.0.1:") && uri.ends_with("/callback"), + "unexpected redirect uri: {uri}" + ); + } +} diff --git a/crates/code_assistant_core/src/tools/mod.rs b/crates/code_assistant_core/src/tools/mod.rs index f46233fa..7cc34967 100644 --- a/crates/code_assistant_core/src/tools/mod.rs +++ b/crates/code_assistant_core/src/tools/mod.rs @@ -10,6 +10,9 @@ pub mod config; // MCP client mode: mcp-servers.json + registration of MCP server tools pub mod mcp; +// Interactive OAuth login for HTTP MCP servers (loopback callback authorizer) +pub mod mcp_auth; + // Persistent per-project trust for project-local `.mcp.json` files pub mod mcp_trust; diff --git a/crates/code_assistant_core/src/tools/registry_provider.rs b/crates/code_assistant_core/src/tools/registry_provider.rs index 3a72841e..8d6ffaf7 100644 --- a/crates/code_assistant_core/src/tools/registry_provider.rs +++ b/crates/code_assistant_core/src/tools/registry_provider.rs @@ -132,10 +132,15 @@ impl ConfigToolRegistry { let read = |path: std::path::PathBuf| std::fs::read_to_string(path).unwrap_or_default(); let tools = ToolsConfig::config_path().map(read).unwrap_or_default(); let mcp = read(crate::tools::mcp::mcp_servers_config_path()); + // Persisted OAuth tokens are an input to what an HTTP server + // contributes: completing a login writes a token file but touches + // none of the config files, so include it here or the cached + // (pre-auth, tool-less) registry would be reused until restart. + let oauth = crate::tools::mcp::mcp_oauth_fingerprint(); let (dir, local) = local_mcp_dir .map(|dir| (dir.display().to_string(), read(dir.join(".mcp.json")))) .unwrap_or_default(); - format!("{tools}\u{0}{mcp}\u{0}{dir}\u{0}{local}") + format!("{tools}\u{0}{mcp}\u{0}{oauth}\u{0}{dir}\u{0}{local}") } } @@ -164,7 +169,14 @@ impl ConnectionProvider for ConfigToolRegistry { } } // Connect outside the lock (slow: process launch / HTTP handshake). - let connection = Arc::new(McpServerConnection::connect(name, config).await?); + // HTTP servers get a persistent OAuth credential store so a token + // obtained via the interactive login is reused silently; stdio servers + // have no OAuth, so pass none. + let credentials = config + .transport + .is_http() + .then(|| crate::tools::mcp::mcp_oauth_credential_store(name)); + let connection = Arc::new(McpServerConnection::connect(name, config, credentials).await?); // Re-check under the lock: if another caller connected the same server // meanwhile, keep theirs and drop ours (shut down on drop). let mut connections = self.connections.lock().await; @@ -309,6 +321,32 @@ mod tests { .await; } + /// Completing an OAuth login (a new/updated token file under the config + /// dir's `mcp-oauth/`) rebuilds the registry, so the authorized server's + /// tools appear on the next run without restarting the app. + #[tokio::test] + async fn oauth_token_change_rebuilds() { + let config = tempfile::tempdir().unwrap(); + temp_env::async_with_vars( + [("CODE_ASSISTANT_CONFIG_DIR", Some(config.path()))], + async { + let provider = ConfigToolRegistry::new(); + let before = provider.current().await; + // Simulate the interactive login writing a token file. + let oauth_dir = config.path().join("mcp-oauth"); + std::fs::create_dir_all(&oauth_dir).unwrap(); + std::fs::write(oauth_dir.join("sap_knowledge.json"), r#"{"client_id":"x"}"#) + .unwrap(); + let after = provider.current().await; + assert!( + !Arc::ptr_eq(&before, &after), + "a new OAuth token must invalidate the cached registry" + ); + }, + ) + .await; + } + /// The cache holds only weak references: once nothing uses a registry any /// more, it is gone (with its MCP connections) and a later request builds /// afresh instead of resurrecting stale state. diff --git a/crates/mcp_client/Cargo.toml b/crates/mcp_client/Cargo.toml index 2b769ce1..f313ac85 100644 --- a/crates/mcp_client/Cargo.toml +++ b/crates/mcp_client/Cargo.toml @@ -5,7 +5,8 @@ edition = "2024" [dependencies] tools_core = { path = "../tools_core" } -rmcp = { version = "2.1", default-features = false, features = [ +rmcp = { version = "3.3", default-features = false, features = [ + "auth", "base64", "client", "transport-child-process", @@ -15,14 +16,17 @@ rmcp = { version = "2.1", default-features = false, features = [ anyhow = "1.0" async-trait = "0.1" http = "1" +reqwest = { version = "0.13", default-features = false, features = ["rustls"] } serde = { version = "1.0", features = ["derive"] } serde_json = "1.0" +thiserror = "2" tokio = { version = "1.48", features = ["macros", "process", "rt", "io-util", "time", "sync"] } tracing = "0.1" [dev-dependencies] command_executor = { path = "../command_executor" } -rmcp = { version = "2.1", default-features = false, features = [ +tempfile = "3" +rmcp = { version = "3.3", default-features = false, features = [ "base64", "client", "server", diff --git a/crates/mcp_client/src/auth.rs b/crates/mcp_client/src/auth.rs new file mode 100644 index 00000000..765df148 --- /dev/null +++ b/crates/mcp_client/src/auth.rs @@ -0,0 +1,256 @@ +//! OAuth authorization support for HTTP (streamable) MCP servers. +//! +//! MCP's authorization spec is an OAuth 2.1 authorization-code flow with PKCE +//! and discovery (RFC 9728 protected-resource metadata → RFC 8414 +//! authorization-server metadata). A server that requires it answers the +//! `initialize` request with `401` and a `WWW-Authenticate` challenge; the +//! client then discovers the authorization server, runs a browser consent +//! flow, and reconnects with a bearer token. +//! +//! This module keeps the pieces that a client needs but that are not baked +//! into the transport: +//! +//! * [`FileCredentialStore`] — persists one server's tokens to disk so the +//! interactive login survives restarts (implements rmcp's +//! [`CredentialStore`]). +//! * [`OAuthAuthorizer`] — the embedder-provided seam that presents the +//! authorization URL to the user (open a browser) and returns the redirect +//! callback parameters. +//! +//! The interactive login itself lives in [`crate::client`]; token *reuse* is +//! non-interactive and happens transparently when connecting. + +use async_trait::async_trait; +use rmcp::transport::auth::{AuthError, CredentialStore, StoredCredentials}; +use std::path::{Path, PathBuf}; +use std::sync::Arc; + +/// The error a connect attempt fails with when the server demands OAuth +/// authorization we do not (yet) have. Distinct and downcastable so an +/// embedder can tell "you must authenticate" apart from a generic connection +/// failure and offer an *Authenticate* action instead of just a retry. +/// +/// It carries the server's `WWW-Authenticate` challenge, which seeds OAuth +/// discovery in the interactive flow ([`crate::client::authenticate_http_server`]). +#[derive(Debug, Clone, thiserror::Error)] +#[error("MCP server '{server}' requires OAuth authorization")] +pub struct AuthorizationRequired { + /// The configured server name. + pub server: String, + /// The raw `WWW-Authenticate` header from the server's 401/403. + pub challenge: String, +} + +/// A [`CredentialStore`] that persists a single MCP server's OAuth tokens to a +/// JSON file. Each server gets its own file (the embedder picks the path, +/// typically `/mcp-oauth/.json`), so one store instance +/// backs exactly one server's credentials. +/// +/// Writes are atomic (write-to-temp-then-rename) and, on Unix, the file is +/// created with `0600` permissions since it holds bearer/refresh tokens. +#[derive(Debug, Clone)] +pub struct FileCredentialStore { + path: PathBuf, +} + +impl FileCredentialStore { + /// Back this store with the file at `path`. The file and its parent + /// directory are created lazily on the first `save`. + pub fn new(path: impl Into) -> Self { + Self { path: path.into() } + } + + /// The file this store reads and writes. + pub fn path(&self) -> &Path { + &self.path + } +} + +fn store_err(context: &str, error: impl std::fmt::Display) -> AuthError { + AuthError::CredentialStoreError(format!("{context}: {error}")) +} + +#[async_trait] +impl CredentialStore for FileCredentialStore { + async fn load(&self) -> Result, AuthError> { + let bytes = match std::fs::read(&self.path) { + Ok(bytes) => bytes, + Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(None), + Err(error) => { + return Err(store_err( + &format!("reading {}", self.path.display()), + error, + )); + } + }; + let credentials = serde_json::from_slice(&bytes) + .map_err(|error| store_err(&format!("parsing {}", self.path.display()), error))?; + Ok(Some(credentials)) + } + + async fn save(&self, credentials: StoredCredentials) -> Result<(), AuthError> { + if let Some(parent) = self.path.parent() { + std::fs::create_dir_all(parent) + .map_err(|error| store_err(&format!("creating {}", parent.display()), error))?; + } + let json = serde_json::to_vec_pretty(&credentials) + .map_err(|error| store_err("serializing credentials", error))?; + + // Write to a sibling temp file, then rename over the target so a + // reader never sees a half-written file. + let tmp = self.path.with_extension("json.tmp"); + std::fs::write(&tmp, &json) + .map_err(|error| store_err(&format!("writing {}", tmp.display()), error))?; + restrict_permissions(&tmp)?; + std::fs::rename(&tmp, &self.path) + .map_err(|error| store_err(&format!("renaming into {}", self.path.display()), error))?; + Ok(()) + } + + async fn clear(&self) -> Result<(), AuthError> { + match std::fs::remove_file(&self.path) { + Ok(()) => Ok(()), + Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()), + Err(error) => Err(store_err( + &format!("removing {}", self.path.display()), + error, + )), + } + } +} + +#[cfg(unix)] +fn restrict_permissions(path: &Path) -> Result<(), AuthError> { + use std::os::unix::fs::PermissionsExt; + let perms = std::fs::Permissions::from_mode(0o600); + std::fs::set_permissions(path, perms) + .map_err(|error| store_err(&format!("chmod {}", path.display()), error)) +} + +#[cfg(not(unix))] +fn restrict_permissions(_path: &Path) -> Result<(), AuthError> { + Ok(()) +} + +/// Adapts a shared `Arc` into an owned [`CredentialStore`] +/// value. rmcp's `AuthorizationManager::set_credential_store` takes the store +/// *by value*, but we want to hand the same underlying store to both the +/// non-interactive reuse path and the interactive login — so we share one +/// behind an `Arc` and clone this cheap adapter into each manager. +#[derive(Clone)] +pub struct SharedCredentialStore(pub Arc); + +impl std::fmt::Debug for SharedCredentialStore { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str("SharedCredentialStore") + } +} + +#[async_trait] +impl CredentialStore for SharedCredentialStore { + async fn load(&self) -> Result, AuthError> { + self.0.load().await + } + + async fn save(&self, credentials: StoredCredentials) -> Result<(), AuthError> { + self.0.save(credentials).await + } + + async fn clear(&self) -> Result<(), AuthError> { + self.0.clear().await + } +} + +/// The parameters an authorization server returns to the OAuth redirect URI +/// after the user consents. +#[derive(Debug, Clone)] +pub struct AuthorizationOutcome { + /// The authorization code to exchange for tokens. + pub code: String, + /// The CSRF `state` value the server echoed back. + pub state: String, + /// The optional RFC 9207 `iss` (issuer) parameter, validated when present. + pub issuer: Option, +} + +/// The embedder-provided half of the interactive OAuth flow: present the +/// authorization URL to the user (typically by opening a browser and running +/// a loopback redirect server) and return the callback parameters. +/// +/// The MCP client crate drives discovery, PKCE, token exchange and reconnect; +/// it only needs the embedder to handle the human-facing browser round-trip, +/// which is why this is a seam rather than baked in. +#[async_trait] +pub trait OAuthAuthorizer: Send + Sync { + /// The redirect URI the authorization server should send the user back to + /// (e.g. `http://127.0.0.1:8117/callback`). Must match what the loopback + /// server in [`Self::authorize`] listens on. + fn redirect_uri(&self) -> String; + + /// Present `authorization_url` to the user and resolve once the redirect + /// callback has been received. + async fn authorize(&self, authorization_url: String) -> anyhow::Result; +} + +#[cfg(test)] +mod tests { + use super::*; + use rmcp::transport::auth::StoredCredentials; + + fn sample_credentials() -> StoredCredentials { + StoredCredentials::new( + "client-123".to_string(), + None, + vec!["mcp:read".to_string()], + Some(1_700_000_000), + ) + .with_issuer(Some("https://issuer.example.com".to_string())) + } + + #[tokio::test] + async fn load_returns_none_when_file_is_absent() { + let dir = tempfile::tempdir().unwrap(); + let store = FileCredentialStore::new(dir.path().join("nested/server.json")); + assert!(store.load().await.unwrap().is_none()); + } + + #[tokio::test] + async fn save_then_load_round_trips_and_creates_dirs() { + let dir = tempfile::tempdir().unwrap(); + let store = FileCredentialStore::new(dir.path().join("mcp-oauth/sap.json")); + store.save(sample_credentials()).await.unwrap(); + + let loaded = store.load().await.unwrap().expect("credentials present"); + assert_eq!(loaded.client_id, "client-123"); + assert_eq!(loaded.granted_scopes, vec!["mcp:read".to_string()]); + assert_eq!(loaded.issuer.as_deref(), Some("https://issuer.example.com")); + } + + #[tokio::test] + async fn clear_removes_the_file_and_is_idempotent() { + let dir = tempfile::tempdir().unwrap(); + let store = FileCredentialStore::new(dir.path().join("server.json")); + store.save(sample_credentials()).await.unwrap(); + assert!(store.path().exists()); + + store.clear().await.unwrap(); + assert!(!store.path().exists()); + // Clearing an absent file is a no-op, not an error. + store.clear().await.unwrap(); + } + + #[cfg(unix)] + #[tokio::test] + async fn saved_file_is_not_world_readable() { + use std::os::unix::fs::PermissionsExt; + let dir = tempfile::tempdir().unwrap(); + let store = FileCredentialStore::new(dir.path().join("server.json")); + store.save(sample_credentials()).await.unwrap(); + + let mode = std::fs::metadata(store.path()) + .unwrap() + .permissions() + .mode(); + assert_eq!(mode & 0o077, 0, "group/other bits must be clear: {mode:o}"); + } +} diff --git a/crates/mcp_client/src/client.rs b/crates/mcp_client/src/client.rs index 3472851b..73d4e7b0 100644 --- a/crates/mcp_client/src/client.rs +++ b/crates/mcp_client/src/client.rs @@ -4,15 +4,30 @@ //! lives as long as the connection; for an HTTP server it is a streamable HTTP //! session. Wrapped tools hold the connection behind an `Arc`, so a dead //! server degrades to tool errors, never a crashed agent. +//! +//! HTTP servers may require OAuth (MCP's authorization spec): connecting is +//! *reactive* — we try unauthenticated first and, on the server's `401` +//! challenge, either reuse a stored token non-interactively or surface an +//! [`AuthorizationRequired`] error so the embedder can offer an interactive +//! login. That login is [`authenticate_http_server`], driven by an +//! [`OAuthAuthorizer`]; it persists tokens into the given [`CredentialStore`] +//! so later connects reuse them silently. +use crate::auth::{AuthorizationRequired, OAuthAuthorizer, SharedCredentialStore}; use crate::config::{McpServerConfig, McpTransport}; -use anyhow::{Context, Result}; +use anyhow::{Context, Result, anyhow}; use rmcp::ServiceExt; use rmcp::model::{CallToolRequestParams, CallToolResult, JsonObject, Tool as McpToolDescriptor}; use rmcp::service::{RoleClient, RunningService}; use rmcp::transport::IntoTransport; -use rmcp::transport::streamable_http_client::StreamableHttpClientTransportConfig; +use rmcp::transport::auth::{ + AuthClient, AuthorizationManager, AuthorizationRequest, AuthorizationSession, CredentialStore, +}; +use rmcp::transport::streamable_http_client::{ + StreamableHttpClientTransport, StreamableHttpClientTransportConfig, +}; use std::collections::HashMap; +use std::sync::Arc; use std::time::Duration; /// Timeout for the initialize handshake and for tool discovery. @@ -21,6 +36,10 @@ const CONNECT_TIMEOUT: Duration = Duration::from_secs(30); /// real work (searches, API calls), but a hung server must not hang a turn /// forever. const CALL_TIMEOUT: Duration = Duration::from_secs(300); +/// Timeout for OAuth HTTP operations (discovery, registration, token exchange +/// and refresh). Kept off the transport client so it never cuts long-lived +/// SSE streams short. +const OAUTH_HTTP_TIMEOUT: Duration = Duration::from_secs(30); /// A live connection to one MCP server. pub struct McpServerConnection { @@ -30,14 +49,29 @@ pub struct McpServerConnection { impl McpServerConnection { /// Connect to the configured server, running the MCP initialize handshake - /// over its transport: a launched child process (stdio) or an HTTP - /// streamable endpoint. - pub async fn connect(name: &str, config: &McpServerConfig) -> Result { + /// over its transport, reusing OAuth tokens from `credentials` for an HTTP + /// server when present. + /// + /// Pass `None` when there is nothing to authenticate with (stdio servers, + /// which have no OAuth; or a server reached purely via a static + /// `Authorization` header). With `None`, or when the store holds no valid + /// token, an HTTP server that demands OAuth fails with a typed + /// [`AuthorizationRequired`] error — this never opens a browser, so the + /// caller (or a UI) can offer the interactive login + /// ([`authenticate_http_server`]) instead. With a valid stored token the + /// connection is authorized silently (rmcp refreshes it transparently). + pub async fn connect( + name: &str, + config: &McpServerConfig, + credentials: Option>, + ) -> Result { match &config.transport { McpTransport::Stdio { command, args, env } => { Self::connect_stdio(name, command, args, env).await } - McpTransport::Http { url, headers } => Self::connect_http(name, url, headers).await, + McpTransport::Http { url, headers } => { + Self::connect_http(name, url, headers, credentials).await + } } } @@ -57,36 +91,104 @@ impl McpServerConnection { } /// Connect to an HTTP (streamable) MCP server at `url`, sending the given - /// custom headers (e.g. `Authorization`) with every request. + /// custom headers (e.g. a static `Authorization`) with every request, and + /// reusing OAuth tokens from `credentials` when available. async fn connect_http( name: &str, url: &str, headers: &HashMap, + credentials: Option>, ) -> Result { - let mut config = StreamableHttpClientTransportConfig::with_uri(url.to_string()); - if !headers.is_empty() { - let mut header_map = HashMap::with_capacity(headers.len()); - for (key, value) in headers { - let name = http::HeaderName::from_bytes(key.as_bytes()) - .with_context(|| format!("invalid HTTP header name '{key}'"))?; - let value = http::HeaderValue::from_str(value) - .with_context(|| format!("invalid value for HTTP header '{key}'"))?; - header_map.insert(name, value); - } - config = config.custom_headers(header_map); + // 1. If we hold a stored OAuth token, connect authorized without any + // user interaction (rmcp refreshes it transparently when needed). + if let Some(store) = &credentials + && let Some(service) = + Self::connect_http_with_stored_token(name, url, headers, store.clone()).await? + { + return Ok(Self { + name: name.to_string(), + service, + }); } + + // 2. Try unauthenticated (also the path for a static `Authorization` + // header). A 401 surfaces as a typed AuthorizationRequired. let transport = - rmcp::transport::streamable_http_client::StreamableHttpClientTransport::from_config( - config, - ); - Self::connect_transport(name, transport) + StreamableHttpClientTransport::from_config(http_transport_config(url, headers)?); + Self::serve_http(name, transport) .await + .map(|service| Self { + name: name.to_string(), + service, + }) .with_context(|| format!("failed to connect to HTTP MCP server '{name}' ({url})")) } + /// Try to connect with a previously stored OAuth token. Returns `Ok(None)` + /// when the store holds no usable credentials, so the caller falls back to + /// an unauthenticated attempt. + async fn connect_http_with_stored_token( + name: &str, + url: &str, + headers: &HashMap, + store: Arc, + ) -> Result>> { + let mut manager = AuthorizationManager::new(url).await.map_err(|error| { + anyhow!("initializing OAuth manager for MCP server '{name}': {error}") + })?; + manager.set_credential_store(SharedCredentialStore(store)); + manager.with_client(oauth_http_client()?).map_err(|error| { + anyhow!("configuring OAuth client for MCP server '{name}': {error}") + })?; + + let has_credentials = manager + .initialize_from_store() + .await + .map_err(|error| anyhow!("loading stored credentials for '{name}': {error}"))?; + if !has_credentials { + return Ok(None); + } + + let auth_client = AuthClient::new(reqwest::Client::new(), manager); + let transport = StreamableHttpClientTransport::with_client( + auth_client, + http_transport_config(url, headers)?, + ); + Self::serve_http(name, transport) + .await + .map(Some) + .with_context(|| { + format!( + "failed to connect to HTTP MCP server '{name}' ({url}) with stored credentials" + ) + }) + } + + /// Serve the client handler over a streamable HTTP transport, mapping the + /// server's `401` authorization challenge to a typed + /// [`AuthorizationRequired`] error and everything else to a plain failure. + async fn serve_http(name: &str, transport: T) -> Result> + where + T: IntoTransport, + E: std::error::Error + Send + Sync + 'static, + { + match tokio::time::timeout(CONNECT_TIMEOUT, ().serve(transport)).await { + Err(_elapsed) => Err(anyhow!("timeout initializing MCP server '{name}'")), + Ok(Ok(service)) => Ok(service), + Ok(Err(error)) => match error.auth_challenge() { + Some(challenge) => Err(anyhow::Error::new(AuthorizationRequired { + server: name.to_string(), + challenge: challenge.to_string(), + })), + None => Err(anyhow::Error::new(error) + .context(format!("failed to initialize MCP server '{name}'"))), + }, + } + } + /// Run the MCP initialize handshake over an arbitrary transport. Used by - /// tests (in-process duplex streams); embedders normally use - /// [`Self::connect`]. + /// stdio servers and by tests (in-process duplex streams); HTTP servers go + /// through [`Self::connect_http`] so they get OAuth handling. pub async fn connect_transport(name: &str, transport: T) -> Result where T: IntoTransport, @@ -143,3 +245,101 @@ impl McpServerConnection { .map_err(|e| anyhow::anyhow!("failed to shut down MCP server '{}': {e}", self.name)) } } + +/// Build the streamable HTTP transport config for `url`, attaching any custom +/// headers (values already have `${VAR}` substituted by the config layer). +fn http_transport_config( + url: &str, + headers: &HashMap, +) -> Result { + let mut config = StreamableHttpClientTransportConfig::with_uri(url.to_string()); + if !headers.is_empty() { + let mut header_map = HashMap::with_capacity(headers.len()); + for (key, value) in headers { + let name = http::HeaderName::from_bytes(key.as_bytes()) + .with_context(|| format!("invalid HTTP header name '{key}'"))?; + let value = http::HeaderValue::from_str(value) + .with_context(|| format!("invalid value for HTTP header '{key}'"))?; + header_map.insert(name, value); + } + config = config.custom_headers(header_map); + } + Ok(config) +} + +/// A `reqwest` client for OAuth HTTP operations (discovery, registration, +/// token exchange, refresh) with a bounded timeout. +fn oauth_http_client() -> Result { + reqwest::Client::builder() + .timeout(OAUTH_HTTP_TIMEOUT) + .build() + .context("building OAuth HTTP client") +} + +/// Run the interactive OAuth login for an HTTP MCP server and persist the +/// resulting tokens into `credential_store`, so a later +/// [`McpServerConnection::connect`] reuses them without any user interaction. +/// +/// The flow: discover the authorization server from `config`'s URL, register +/// (or select) an OAuth client, hand the authorization URL to `authorizer` +/// (which opens a browser and waits for the redirect callback), then exchange +/// the returned code for tokens. `client_name` labels this client during +/// dynamic client registration. +/// +/// Errors if `config` is not an HTTP server — OAuth does not apply to stdio. +pub async fn authenticate_http_server( + name: &str, + config: &McpServerConfig, + credential_store: Arc, + authorizer: &dyn OAuthAuthorizer, + client_name: &str, +) -> Result<()> { + let McpTransport::Http { url, .. } = &config.transport else { + return Err(anyhow!( + "MCP server '{name}' is not an HTTP server; OAuth authorization only applies to HTTP transports" + )); + }; + + let mut manager = AuthorizationManager::new(url) + .await + .map_err(|error| anyhow!("initializing OAuth manager for MCP server '{name}': {error}"))?; + manager.set_credential_store(SharedCredentialStore(credential_store)); + manager + .with_client(oauth_http_client()?) + .map_err(|error| anyhow!("configuring OAuth client for MCP server '{name}': {error}"))?; + + // Discover the authorization server (RFC 9728 → RFC 8414 / OIDC). + let resolution = manager + .resolve_metadata() + .await + .map_err(|error| anyhow!("discovering OAuth metadata for MCP server '{name}': {error}"))?; + manager.set_metadata(resolution.metadata); + + // Select a client (pre-registered / CIMD / dynamic registration) and build + // the authorization URL. + let request = + AuthorizationRequest::new(authorizer.redirect_uri()).with_client_name(client_name); + let session = + AuthorizationSession::new(manager, request) + .await + .map_err(|(_manager, error)| { + anyhow!("starting OAuth authorization for MCP server '{name}': {error}") + })?; + + // Hand the URL to the embedder: open a browser and await the redirect. + let outcome = authorizer + .authorize(session.get_authorization_url().to_string()) + .await + .with_context(|| format!("browser authorization for MCP server '{name}'"))?; + + // Exchange the code for tokens; this persists StoredCredentials into the + // shared credential store as a side effect. + session + .handle_callback_with_issuer(&outcome.code, &outcome.state, outcome.issuer.as_deref()) + .await + .map_err(|error| { + anyhow!("exchanging OAuth authorization code for MCP server '{name}': {error}") + })?; + + Ok(()) +} diff --git a/crates/mcp_client/src/lib.rs b/crates/mcp_client/src/lib.rs index 934880e5..a1a5dac0 100644 --- a/crates/mcp_client/src/lib.rs +++ b/crates/mcp_client/src/lib.rs @@ -6,6 +6,7 @@ //! //! Built on the official Rust MCP SDK (`rmcp`). +pub mod auth; pub mod client; pub mod config; pub mod naming; @@ -16,7 +17,11 @@ pub mod tool; #[cfg(test)] mod tests; -pub use client::McpServerConnection; +pub use auth::{ + AuthorizationOutcome, AuthorizationRequired, FileCredentialStore, OAuthAuthorizer, + SharedCredentialStore, +}; +pub use client::{McpServerConnection, authenticate_http_server}; pub use config::{ McpServerConfig, McpServersConfig, McpTransport, parse_local_mcp_json, substitute_variables, }; @@ -27,4 +32,8 @@ pub use registry::{ register_connection_tools, register_mcp_tools, register_mcp_tools_pooled, server_scope_capability, }; +/// Re-exported so embedders can name the credential-store trait object type +/// (`Arc`) that [`McpServerConnection::connect`] and +/// [`authenticate_http_server`] take, without depending on `rmcp` directly. +pub use rmcp::transport::auth::CredentialStore; pub use tool::McpTool; diff --git a/crates/mcp_client/src/registry.rs b/crates/mcp_client/src/registry.rs index 7ac22f7a..3d0b5b22 100644 --- a/crates/mcp_client/src/registry.rs +++ b/crates/mcp_client/src/registry.rs @@ -49,6 +49,13 @@ pub struct McpServerStatus { /// tools. Every registered tool carries [`MCP_CAPABILITY`], its server's /// scope tag, and the given extra capability tags (the embedder's scope /// vocabulary, e.g. `scope:agent`). +/// +/// Connections are made without OAuth credentials, so this simple path +/// supports stdio and static-header servers; an HTTP server that requires +/// interactive OAuth reports an [`crate::AuthorizationRequired`] status and +/// contributes no tools. Embedders that persist OAuth tokens should use +/// [`register_mcp_tools_pooled`] with a [`ConnectionProvider`] that threads a +/// credential store. pub async fn register_mcp_tools( registry: &mut ToolRegistry, config: &McpServersConfig, @@ -57,7 +64,7 @@ pub async fn register_mcp_tools( let mut statuses = Vec::new(); for (name, server_config) in config.enabled_servers() { let result = async { - let connection = McpServerConnection::connect(name, server_config).await?; + let connection = McpServerConnection::connect(name, server_config, None).await?; register_connection_tools( registry, Arc::new(connection), @@ -146,11 +153,17 @@ pub struct DiscoveredTool { /// Connect to a server, list everything it offers (ignoring the tool /// filter), and shut the connection down again. For configuration UIs. +/// Reuses OAuth tokens from `credentials` for an HTTP server; pass `None` +/// when there is nothing to authenticate with (stdio, or a static-header +/// server). An HTTP server that still needs authorization fails with +/// [`crate::AuthorizationRequired`], which a UI can turn into an +/// "Authenticate" action. pub async fn discover_tools( server_name: &str, config: &McpServerConfig, + credentials: Option>, ) -> Result> { - let connection = McpServerConnection::connect(server_name, config).await?; + let connection = McpServerConnection::connect(server_name, config, credentials).await?; let descriptors = connection.list_tools().await?; let _ = connection.shutdown().await; Ok(descriptors diff --git a/crates/mcp_client/src/tests.rs b/crates/mcp_client/src/tests.rs index 3558fa9a..aff08701 100644 --- a/crates/mcp_client/src/tests.rs +++ b/crates/mcp_client/src/tests.rs @@ -6,8 +6,8 @@ use crate::client::McpServerConnection; use crate::config::McpServerConfig; use crate::registry::{MCP_CAPABILITY, register_connection_tools, server_scope_capability}; use rmcp::model::{ - CallToolRequestParams, CallToolResult, ContentBlock, ErrorData, ListToolsResult, - PaginatedRequestParams, Tool as McpToolDescriptor, + CallToolRequestParams, CallToolResponse, CallToolResult, ContentBlock, ErrorData, + ListToolsResult, PaginatedRequestParams, Tool as McpToolDescriptor, }; use rmcp::service::{RequestContext, RoleServer}; use rmcp::{ServerHandler, ServiceExt}; @@ -53,7 +53,7 @@ impl ServerHandler for TestServer { &self, request: CallToolRequestParams, _context: RequestContext, - ) -> Result { + ) -> Result { match request.name.as_ref() { "echo" => { let message = request @@ -62,11 +62,12 @@ impl ServerHandler for TestServer { .and_then(|arguments| arguments.get("message")) .and_then(|value| value.as_str()) .unwrap_or_default(); - Ok(CallToolResult::success(vec![ContentBlock::text(format!( - "echo: {message}" - ))])) + Ok( + CallToolResult::success(vec![ContentBlock::text(format!("echo: {message}"))]) + .into(), + ) } - "fail" => Ok(CallToolResult::error(vec![ContentBlock::text("it broke")])), + "fail" => Ok(CallToolResult::error(vec![ContentBlock::text("it broke")]).into()), other => Err(ErrorData::invalid_params( format!("unknown tool: {other}"), None, @@ -241,7 +242,9 @@ async fn connects_to_a_real_stdio_server() { "command": binary.to_string_lossy(), "args": ["server"] })); - let connection = McpServerConnection::connect("self", &config).await.unwrap(); + let connection = McpServerConnection::connect("self", &config, None) + .await + .unwrap(); let tools = connection.list_tools().await.unwrap(); assert!( tools.iter().any(|tool| tool.name == "read_files"), @@ -299,7 +302,7 @@ async fn connects_and_calls_over_http() { let config = server_config(json!({ "url": url })); let connection = Arc::new( - McpServerConnection::connect("http-test", &config) + McpServerConnection::connect("http-test", &config, None) .await .expect("client failed to connect over HTTP"), ); @@ -316,3 +319,92 @@ async fn connects_and_calls_over_http() { let rendered = output.as_render().render(&mut ResourcesTracker::new()); assert_eq!(rendered, "echo: over http"); } + +/// Serve a minimal HTTP endpoint that answers every request with `401` and a +/// `WWW-Authenticate` challenge — the reactive OAuth trigger. Returns its +/// `/mcp` URL and the server task (aborted on drop). +async fn spawn_auth_required_server() -> (String, tokio::task::JoinHandle<()>) { + use axum::http::StatusCode; + use axum::response::IntoResponse; + + async fn unauthorized() -> impl IntoResponse { + ( + StatusCode::UNAUTHORIZED, + [( + "WWW-Authenticate", + "Bearer resource_metadata=\"https://auth.example.com/.well-known/oauth-protected-resource\"", + )], + "authorization required", + ) + } + + let router = axum::Router::new().route("/mcp", axum::routing::any(unauthorized)); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + let _ = axum::serve(listener, router).await; + }); + (format!("http://{addr}/mcp"), server) +} + +/// A server that demands OAuth must fail the connect with a *typed* +/// [`AuthorizationRequired`] (carrying the challenge and server name), not a +/// generic error — that is what lets an embedder offer an "Authenticate" +/// action instead of a plain retry. +#[tokio::test] +async fn http_server_requiring_auth_surfaces_authorization_required() { + let (url, _server) = spawn_auth_required_server().await; + let config = server_config(json!({ "url": url })); + + let error = match McpServerConnection::connect("sap", &config, None).await { + Ok(_) => panic!("connect must fail when the server demands authorization"), + Err(error) => error, + }; + + let auth = error + .downcast_ref::() + .unwrap_or_else(|| panic!("expected AuthorizationRequired, got: {error:#}")); + assert_eq!(auth.server, "sap"); + assert!( + auth.challenge.contains("resource_metadata"), + "the WWW-Authenticate challenge is carried through: {}", + auth.challenge + ); +} + +/// OAuth authorization only applies to HTTP servers; asking to authenticate a +/// stdio server is a configuration error, surfaced without touching the +/// network. +#[tokio::test] +async fn authenticate_rejects_stdio_servers() { + use crate::auth::{AuthorizationOutcome, OAuthAuthorizer}; + + struct NeverCalled; + #[async_trait::async_trait] + impl OAuthAuthorizer for NeverCalled { + fn redirect_uri(&self) -> String { + "http://127.0.0.1:0/callback".to_string() + } + async fn authorize(&self, _url: String) -> anyhow::Result { + panic!("authorizer must not be invoked for a stdio server") + } + } + + let config = server_config(json!({ "command": "npx" })); + let store: Arc = + Arc::new(rmcp::transport::auth::InMemoryCredentialStore::new()); + + let error = crate::client::authenticate_http_server( + "stdio-srv", + &config, + store, + &NeverCalled, + "code-assistant", + ) + .await + .expect_err("authenticating a stdio server must fail"); + assert!( + error.to_string().contains("not an HTTP server"), + "unexpected error: {error:#}" + ); +} diff --git a/crates/ui_gpui/src/settings_screen/mcp_section.rs b/crates/ui_gpui/src/settings_screen/mcp_section.rs index 1cc58d1c..99a00244 100644 --- a/crates/ui_gpui/src/settings_screen/mcp_section.rs +++ b/crates/ui_gpui/src/settings_screen/mcp_section.rs @@ -35,6 +35,12 @@ enum DiscoveryState { Loading, Loaded(Vec), Failed(String), + /// The HTTP server requires OAuth and we have no (valid) token — offer an + /// Authenticate action instead of a bare retry. + NeedsAuth, + /// The interactive OAuth login is running (browser open, awaiting the + /// redirect callback). + Authenticating, } pub struct McpSection { @@ -160,12 +166,17 @@ impl McpSection { let runtime = tokio::runtime::Builder::new_current_thread() .enable_all() .build()?; - runtime.block_on(mcp::discover_tools(&server_name, server)) + runtime.block_on(mcp::discover_server_tools(&server_name, server)) }) .await; this.update(cx, |this, cx| { let state = match result { Ok(tools) => DiscoveryState::Loaded(tools), + // An HTTP server that demands OAuth gets an Authenticate + // action instead of a plain failure. + Err(error) if error.downcast_ref::().is_some() => { + DiscoveryState::NeedsAuth + } Err(error) => DiscoveryState::Failed(format!("{error:#}")), }; this.discovered.insert(name, state); @@ -176,6 +187,40 @@ impl McpSection { .detach(); } + /// Run the interactive OAuth login for `name` on a background thread, then + /// re-run tool discovery so its tools appear on success (or the error + /// shows on failure). + fn start_authentication(&mut self, name: String, cx: &mut Context) { + self.discovered + .insert(name.clone(), DiscoveryState::Authenticating); + cx.notify(); + cx.spawn(async move |this, cx| { + let server_name = name.clone(); + let result = cx + .background_spawn(async move { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build()?; + runtime.block_on(code_assistant_core::tools::mcp_auth::login_mcp_server( + &server_name, + )) + }) + .await; + this.update(cx, |this, cx| match result { + // On success re-discover: the stored token now lets us list + // the server's tools. + Ok(()) => this.start_discovery(name, cx), + Err(error) => { + this.discovered + .insert(name, DiscoveryState::Failed(format!("{error:#}"))); + cx.notify(); + } + }) + .ok(); + }) + .detach(); + } + fn open_add_form(&mut self, window: &mut gpui::Window, cx: &mut Context) { self.form_mode = FormMode::Adding; self.fill_form("", None, window, cx); @@ -407,6 +452,7 @@ impl McpSection { let name_for_edit = name.to_string(); let name_for_delete = name.to_string(); let name_for_retry = name.to_string(); + let name_for_auth = name.to_string(); div() .flex() @@ -475,6 +521,41 @@ impl McpSection { .text_color(cx.theme().muted_foreground) .child("Connecting to server…") .into_any_element(), + Some(DiscoveryState::Authenticating) => div() + .py_2() + .text_xs() + .text_color(cx.theme().muted_foreground) + .child("Authenticating… complete the login in your browser.") + .into_any_element(), + Some(DiscoveryState::NeedsAuth) => div() + .flex() + .flex_col() + .gap_2() + .py_2() + .child( + div() + .text_xs() + .text_color(cx.theme().muted_foreground) + .child("This server requires authorization."), + ) + .child( + div() + .id(SharedString::from(format!("mcp-auth-{name}"))) + .self_start() + .px_3() + .py_1() + .rounded_md() + .cursor_pointer() + .text_xs() + .bg(cx.theme().primary) + .text_color(cx.theme().primary_foreground) + .hover(|s| s.opacity(0.9)) + .child("Authenticate") + .on_click(cx.listener(move |this, _, _window, cx| { + this.start_authentication(name_for_auth.clone(), cx); + })), + ) + .into_any_element(), Some(DiscoveryState::Failed(error)) => div() .flex() .flex_col()