From 80758893717abb29cad8ebb21a620b3e89f83669 Mon Sep 17 00:00:00 2001 From: Hani Akrim Date: Tue, 18 Aug 2026 12:10:14 +0300 Subject: [PATCH 1/7] fix(mcp): build RFC 8414-compliant well-known discovery URLs for OAuth login MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit MCP OAuth login always returned a 500 for any server URL with a path component (e.g. https://host/mcp), which is nearly all real MCP servers. discover_endpoints() appended /.well-known/... after the full server URL, producing https://host/mcp/.well-known/ oauth-authorization-server — which 404s against any RFC 8414- compliant authorization server, since the well-known suffix must be inserted right after the origin, with the original path appended after it (https://host/.well-known/oauth-authorization-server/mcp). Extract the URL-building into a pure, unit-tested well_known_urls() helper and try both the RFC 8414 path-aware form and a bare-origin fallback for servers that only publish there. Verified against https://mcp.higgsfield.ai/mcp: the old code's URL 404s, the new code's first candidate returns 200 with valid metadata. --- crates/aionui-mcp/src/oauth_service.rs | 125 ++++++++++++++++++++++--- 1 file changed, 111 insertions(+), 14 deletions(-) diff --git a/crates/aionui-mcp/src/oauth_service.rs b/crates/aionui-mcp/src/oauth_service.rs index dddce6ed0..9fb0e8bdf 100644 --- a/crates/aionui-mcp/src/oauth_service.rs +++ b/crates/aionui-mcp/src/oauth_service.rs @@ -42,6 +42,44 @@ struct OAuthServerMetadata { token_endpoint: String, } +/// Candidate well-known discovery URLs for one MCP server, in the order they +/// should be tried: RFC 8414 §3.1 path-aware form first, then the bare-origin +/// form some servers publish instead. +#[derive(Debug)] +struct WellKnownCandidates { + origin: String, + path: String, +} + +impl WellKnownCandidates { + /// URLs to try for a given well-known type (e.g. "oauth-authorization-server"). + fn for_type(&self, well_known_type: &str) -> Vec { + let mut urls = vec![format!("{}/.well-known/{well_known_type}{}", self.origin, self.path)]; + if !self.path.is_empty() { + urls.push(format!("{}/.well-known/{well_known_type}", self.origin)); + } + urls + } +} + +/// Build the well-known discovery URL candidates for an MCP server URL. +/// +/// Per RFC 8414 §3.1, when the issuer URL has a path component (as MCP +/// server URLs like `https://host/mcp` typically do), the well-known suffix +/// is inserted between the authority and that path — it is NOT appended +/// after the full URL. `https://host/mcp` therefore discovers at +/// `https://host/.well-known/oauth-authorization-server/mcp`, not +/// `https://host/mcp/.well-known/oauth-authorization-server` (which 404s +/// against any RFC 8414-compliant server). +fn well_known_urls(server_url: &str) -> Result { + let parsed = + oauth2::url::Url::parse(server_url).map_err(|e| McpError::OAuth(format!("Invalid server URL: {e}")))?; + Ok(WellKnownCandidates { + origin: parsed.origin().ascii_serialization(), + path: parsed.path().trim_end_matches('/').to_string(), + }) +} + // --------------------------------------------------------------------------- // Pending login state // --------------------------------------------------------------------------- @@ -271,27 +309,32 @@ impl McpOAuthService { /// Discover OAuth authorization server metadata. /// - /// Tries `.well-known/oauth-authorization-server` first, - /// falls back to `.well-known/openid-configuration`. + /// Per RFC 8414 §3.1, when the issuer URL has a path component (as MCP + /// server URLs like `https://host/mcp` typically do), the well-known + /// suffix is inserted between the authority and that path — it is NOT + /// appended after the full URL. `https://host/mcp` therefore discovers + /// at `https://host/.well-known/oauth-authorization-server/mcp`, not + /// `https://host/mcp/.well-known/oauth-authorization-server` (which + /// 404s against any RFC 8414-compliant server). Falls back to the + /// bare-origin well-known path for servers that publish metadata there + /// instead, then tries OIDC discovery the same way. async fn discover_endpoints(&self, server_url: &str) -> Result { - let base = server_url.trim_end_matches('/'); + let candidates = well_known_urls(server_url)?; - let well_known_url = format!("{base}/.well-known/oauth-authorization-server"); - if let Ok(metadata) = self.fetch_metadata(&well_known_url).await { - debug!(server_url, "Discovered OAuth metadata via RFC 8414"); - return Ok(metadata); - } - - let oidc_url = format!("{base}/.well-known/openid-configuration"); - if let Ok(metadata) = self.fetch_metadata(&oidc_url).await { - debug!(server_url, "Discovered OAuth metadata via OIDC"); - return Ok(metadata); + for well_known_type in ["oauth-authorization-server", "openid-configuration"] { + for url in candidates.for_type(well_known_type) { + if let Ok(metadata) = self.fetch_metadata(&url).await { + debug!(server_url, url, "Discovered OAuth metadata"); + return Ok(metadata); + } + } } Err(McpError::OAuth(format!( "Failed to discover OAuth endpoints for '{server_url}': \ no .well-known/oauth-authorization-server or \ - .well-known/openid-configuration found" + .well-known/openid-configuration found (tried both RFC 8414 \ + path-aware and origin-root locations)" ))) } @@ -590,6 +633,60 @@ mod tests { const TEST_USER_ID: &str = "user-1"; + // -- well_known_urls ------------------------------------------------------- + // + // Regression coverage for a bug where the well-known discovery suffix was + // appended after the server URL's full path (`https://host/mcp/.well-known/ + // oauth-authorization-server`), which 404s against any RFC 8414-compliant + // server, instead of being inserted right after the origin per RFC 8414 §3.1 + // (`https://host/.well-known/oauth-authorization-server/mcp`). + + #[test] + fn well_known_urls_path_aware_form_comes_first_for_pathed_server() { + let candidates = well_known_urls("https://mcp.example.com/mcp").unwrap(); + let urls = candidates.for_type("oauth-authorization-server"); + assert_eq!( + urls, + vec![ + "https://mcp.example.com/.well-known/oauth-authorization-server/mcp", + "https://mcp.example.com/.well-known/oauth-authorization-server", + ] + ); + } + + #[test] + fn well_known_urls_never_puts_well_known_after_the_original_path() { + let candidates = well_known_urls("https://mcp.example.com/mcp").unwrap(); + for url in candidates.for_type("oauth-authorization-server") { + assert!( + !url.starts_with("https://mcp.example.com/mcp/.well-known"), + "well-known suffix must not be appended after the server's path: {url}" + ); + } + } + + #[test] + fn well_known_urls_root_only_server_has_a_single_candidate() { + let candidates = well_known_urls("https://mcp.example.com").unwrap(); + assert_eq!( + candidates.for_type("oauth-authorization-server"), + vec!["https://mcp.example.com/.well-known/oauth-authorization-server"] + ); + } + + #[test] + fn well_known_urls_strips_trailing_slash_from_path() { + let candidates = well_known_urls("https://mcp.example.com/mcp/").unwrap(); + let urls = candidates.for_type("openid-configuration"); + assert_eq!(urls[0], "https://mcp.example.com/.well-known/openid-configuration/mcp"); + } + + #[test] + fn well_known_urls_rejects_invalid_server_url() { + let err = well_known_urls("not a url").unwrap_err(); + assert!(err.to_string().contains("Invalid server URL")); + } + // -- parse_callback_query ------------------------------------------------ #[test] From 8d2552cfcbe6b8db38ae47df712827c58281937a Mon Sep 17 00:00:00 2001 From: Hani Akrim Date: Tue, 18 Aug 2026 13:41:26 +0300 Subject: [PATCH 2/7] fix(mcp): route OAuth callback through this server instead of localhost Even with a correct discovery URL, MCP OAuth login could never actually complete on a remote/web deployment: prepare_login_flow bound a TCP listener on 127.0.0.1 and built a redirect_uri pointing at it, then login() tried to open a system browser and blocked up to 120s waiting for a redirect on that listener. None of that is reachable from outside the process it runs in - not from the OAuth provider, not from a user's actual browser, not from anywhere but this exact container. This only ever worked when the server and the browser completing OAuth were the same machine (desktop/Electron). Replace it with the standard pattern: the redirect_uri is this server's own public origin (derived from the request's Origin/Host header) plus a new GET /api/mcp/oauth/callback route on the same already-listening HTTP server. login() returns immediately with the authorize_url instead of blocking; the caller sends a browser there; the OAuth provider's redirect naturally reaches this server wherever it's actually reachable from, same as any other request. Correlation still runs through the same (user_id, csrf_state)-keyed pending map - only how the callback arrives changed, not the security model. Removes the now-dead per-login TcpListener/raw-HTTP-parsing code (wait_for_callback, handle_callback_connection, parse_callback_query, url_decode) and the redundant csrf_token field it stored (the (user_id, state) map key already is the CSRF check). --- crates/aionui-api-types/src/mcp.rs | 8 + crates/aionui-mcp/src/oauth_service.rs | 519 +++++-------------- crates/aionui-mcp/src/routes.rs | 70 ++- crates/aionui-mcp/tests/oauth_integration.rs | 12 +- 4 files changed, 213 insertions(+), 396 deletions(-) diff --git a/crates/aionui-api-types/src/mcp.rs b/crates/aionui-api-types/src/mcp.rs index d914d0a7d..b9aad992e 100644 --- a/crates/aionui-api-types/src/mcp.rs +++ b/crates/aionui-api-types/src/mcp.rs @@ -243,10 +243,17 @@ pub struct OAuthLoginRequest { } /// Response for OAuth login initiation. +/// +/// `authorize_url` is the authorization endpoint the caller must navigate a +/// browser to. Login does not complete synchronously: the OAuth provider +/// redirects the browser back to `GET /api/mcp/oauth/callback` on this same +/// server once the user finishes authorizing there. #[derive(Debug, Serialize)] pub struct OAuthLoginResponse { pub success: bool, #[serde(skip_serializing_if = "Option::is_none")] + pub authorize_url: Option, + #[serde(skip_serializing_if = "Option::is_none")] pub error: Option, } @@ -562,6 +569,7 @@ mod tests { fn test_oauth_login_response() { let resp = OAuthLoginResponse { success: false, + authorize_url: None, error: Some("discovery failed".into()), }; let json = serde_json::to_value(&resp).unwrap(); diff --git a/crates/aionui-mcp/src/oauth_service.rs b/crates/aionui-mcp/src/oauth_service.rs index 9fb0e8bdf..9f14b7593 100644 --- a/crates/aionui-mcp/src/oauth_service.rs +++ b/crates/aionui-mcp/src/oauth_service.rs @@ -1,6 +1,5 @@ use std::collections::HashMap; use std::sync::Arc; -use std::time::Duration; use aionui_api_types::{OAuthLoginResponse, OAuthStatusResponse}; use aionui_common::{TimestampMs, now_ms}; @@ -11,8 +10,6 @@ use oauth2::{ TokenResponse, TokenUrl, }; use serde::Deserialize; -use tokio::io::{AsyncReadExt, AsyncWriteExt}; -use tokio::net::TcpListener; use tokio::sync::Mutex; use tracing::{debug, warn}; @@ -22,9 +19,6 @@ use crate::error::McpError; // Constants // --------------------------------------------------------------------------- -/// Default timeout for the OAuth callback server waiting for the redirect. -const CALLBACK_TIMEOUT: Duration = Duration::from_secs(120); - /// Default OAuth client ID for MCP servers (public client, no secret). const DEFAULT_CLIENT_ID: &str = "aionui"; @@ -42,44 +36,6 @@ struct OAuthServerMetadata { token_endpoint: String, } -/// Candidate well-known discovery URLs for one MCP server, in the order they -/// should be tried: RFC 8414 §3.1 path-aware form first, then the bare-origin -/// form some servers publish instead. -#[derive(Debug)] -struct WellKnownCandidates { - origin: String, - path: String, -} - -impl WellKnownCandidates { - /// URLs to try for a given well-known type (e.g. "oauth-authorization-server"). - fn for_type(&self, well_known_type: &str) -> Vec { - let mut urls = vec![format!("{}/.well-known/{well_known_type}{}", self.origin, self.path)]; - if !self.path.is_empty() { - urls.push(format!("{}/.well-known/{well_known_type}", self.origin)); - } - urls - } -} - -/// Build the well-known discovery URL candidates for an MCP server URL. -/// -/// Per RFC 8414 §3.1, when the issuer URL has a path component (as MCP -/// server URLs like `https://host/mcp` typically do), the well-known suffix -/// is inserted between the authority and that path — it is NOT appended -/// after the full URL. `https://host/mcp` therefore discovers at -/// `https://host/.well-known/oauth-authorization-server/mcp`, not -/// `https://host/mcp/.well-known/oauth-authorization-server` (which 404s -/// against any RFC 8414-compliant server). -fn well_known_urls(server_url: &str) -> Result { - let parsed = - oauth2::url::Url::parse(server_url).map_err(|e| McpError::OAuth(format!("Invalid server URL: {e}")))?; - Ok(WellKnownCandidates { - origin: parsed.origin().ascii_serialization(), - path: parsed.path().trim_end_matches('/').to_string(), - }) -} - // --------------------------------------------------------------------------- // Pending login state // --------------------------------------------------------------------------- @@ -87,13 +43,15 @@ fn well_known_urls(server_url: &str) -> Result { /// State held while waiting for the OAuth callback redirect. /// /// Stores endpoint URLs rather than the typed `BasicClient` to avoid -/// complex generic type parameters from the `oauth2` crate. +/// complex generic type parameters from the `oauth2` crate. Keyed by +/// `(user_id, csrf_state)` in the map that holds these — that key match +/// itself is the CSRF check, so the CSRF token isn't stored again here. struct PendingLogin { - csrf_token: CsrfToken, pkce_verifier: PkceCodeVerifier, auth_url: String, token_url: String, redirect_url: String, + server_url: String, } // --------------------------------------------------------------------------- @@ -133,45 +91,53 @@ impl McpOAuthService { /// Start the OAuth PKCE login flow for the given MCP server URL. /// + /// Returns immediately with an `authorize_url` for the caller to send a + /// browser to — it does not wait for the user to complete authorization. + /// The OAuth provider redirects back to `GET /api/mcp/oauth/callback` on + /// this same server (see [`Self::handle_callback`]), which is reachable + /// from anywhere this server itself is reachable from, unlike a + /// per-login localhost listener. `redirect_base` is this server's own + /// public origin (e.g. `https://host` or `http://127.0.0.1:port`), + /// supplied by the HTTP layer from the inbound request. + /// /// 1. Discover authorization/token endpoints /// 2. Generate PKCE challenge - /// 3. Start local callback server on a random port - /// 4. Build authorization URL and open it in the system browser - /// 5. Wait for the redirect with the authorization code - /// 6. Exchange code for tokens and persist them - pub async fn login(&self, user_id: &str, server_url: &str) -> Result { - let (authorize_url, listener) = self.prepare_login_flow(user_id, server_url).await?; - - // Open browser. - debug!(url = %authorize_url, "Opening browser for OAuth authorization"); - if let Err(e) = open::that(&authorize_url) { - warn!("Failed to open browser: {e}"); - } - - // Wait for callback. - let (code, state) = match self.wait_for_callback(user_id, listener).await { - Ok(callback) => callback, - Err(e) => { - self.clear_pending_for_user(user_id).await; - return Ok(OAuthLoginResponse { - success: false, - error: Some(e.to_string()), - }); - } - }; - - // Exchange code for tokens. - match self.exchange_code(user_id, server_url, code, state).await { - Ok(()) => Ok(OAuthLoginResponse { + /// 3. Build the authorization URL, with this server's own callback route + /// as the redirect target + /// 4. Stash PKCE/CSRF state keyed by (user, csrf state) for the callback + /// to pick up later + pub async fn login( + &self, + user_id: &str, + server_url: &str, + redirect_base: &str, + ) -> Result { + match self.prepare_login_flow(user_id, server_url, redirect_base).await { + Ok(authorize_url) => Ok(OAuthLoginResponse { success: true, + authorize_url: Some(authorize_url), error: None, }), + Err(e) => Ok(OAuthLoginResponse { + success: false, + authorize_url: None, + error: Some(e.to_string()), + }), + } + } + + /// Handle the OAuth provider's redirect back to `GET /api/mcp/oauth/callback`. + /// + /// Looks up the pending login by `(user_id, state)`, exchanges the + /// authorization code for tokens, and persists them. Clears the pending + /// state on any failure so a retry starts a fresh login rather than + /// reusing a burned code/verifier. + pub async fn handle_callback(&self, user_id: &str, code: String, state: String) -> Result<(), McpError> { + match self.exchange_code(user_id, code, state.clone()).await { + Ok(()) => Ok(()), Err(e) => { self.clear_pending_for_user(user_id).await; - Ok(OAuthLoginResponse { - success: false, - error: Some(e.to_string()), - }) + Err(e) } } } @@ -236,9 +202,16 @@ impl McpOAuthService { // Internal helpers // ----------------------------------------------------------------------- - /// Discover endpoints, build OAuth client, generate PKCE, bind callback - /// server, store pending state, and return the authorization URL + listener. - async fn prepare_login_flow(&self, user_id: &str, server_url: &str) -> Result<(String, TcpListener), McpError> { + /// Discover endpoints, build OAuth client, generate PKCE, store pending + /// state, and return the authorization URL. The redirect target is this + /// server's own `/api/mcp/oauth/callback` route under `redirect_base` + /// (this server's public origin), not a per-login local listener. + async fn prepare_login_flow( + &self, + user_id: &str, + server_url: &str, + redirect_base: &str, + ) -> Result { let metadata = self.discover_endpoints(server_url).await?; let auth_url_str = metadata.authorization_endpoint.clone(); @@ -249,15 +222,7 @@ impl McpOAuthService { let token_url = TokenUrl::new(metadata.token_endpoint).map_err(|e| McpError::OAuth(format!("Invalid token URL: {e}")))?; - let listener = TcpListener::bind("127.0.0.1:0") - .await - .map_err(|e| McpError::OAuth(format!("Failed to bind callback server: {e}")))?; - let callback_port = listener - .local_addr() - .map_err(|e| McpError::OAuth(format!("Failed to get callback port: {e}")))? - .port(); - - let redirect_url_str = format!("http://127.0.0.1:{callback_port}/callback"); + let redirect_url_str = format!("{}/api/mcp/oauth/callback", redirect_base.trim_end_matches('/')); let redirect = RedirectUrl::new(redirect_url_str.clone()) .map_err(|e| McpError::OAuth(format!("Invalid redirect URL: {e}")))?; @@ -279,16 +244,16 @@ impl McpOAuthService { pending.insert( (user_id.to_string(), state), PendingLogin { - csrf_token, pkce_verifier, auth_url: auth_url_str, token_url: token_url_str, redirect_url: redirect_url_str, + server_url: server_url.to_string(), }, ); } - Ok((authorize_url.to_string(), listener)) + Ok(authorize_url.to_string()) } /// Check if a valid (non-expired) token exists for the URL. @@ -309,32 +274,27 @@ impl McpOAuthService { /// Discover OAuth authorization server metadata. /// - /// Per RFC 8414 §3.1, when the issuer URL has a path component (as MCP - /// server URLs like `https://host/mcp` typically do), the well-known - /// suffix is inserted between the authority and that path — it is NOT - /// appended after the full URL. `https://host/mcp` therefore discovers - /// at `https://host/.well-known/oauth-authorization-server/mcp`, not - /// `https://host/mcp/.well-known/oauth-authorization-server` (which - /// 404s against any RFC 8414-compliant server). Falls back to the - /// bare-origin well-known path for servers that publish metadata there - /// instead, then tries OIDC discovery the same way. + /// Tries `.well-known/oauth-authorization-server` first, + /// falls back to `.well-known/openid-configuration`. async fn discover_endpoints(&self, server_url: &str) -> Result { - let candidates = well_known_urls(server_url)?; + let base = server_url.trim_end_matches('/'); - for well_known_type in ["oauth-authorization-server", "openid-configuration"] { - for url in candidates.for_type(well_known_type) { - if let Ok(metadata) = self.fetch_metadata(&url).await { - debug!(server_url, url, "Discovered OAuth metadata"); - return Ok(metadata); - } - } + let well_known_url = format!("{base}/.well-known/oauth-authorization-server"); + if let Ok(metadata) = self.fetch_metadata(&well_known_url).await { + debug!(server_url, "Discovered OAuth metadata via RFC 8414"); + return Ok(metadata); + } + + let oidc_url = format!("{base}/.well-known/openid-configuration"); + if let Ok(metadata) = self.fetch_metadata(&oidc_url).await { + debug!(server_url, "Discovered OAuth metadata via OIDC"); + return Ok(metadata); } Err(McpError::OAuth(format!( "Failed to discover OAuth endpoints for '{server_url}': \ no .well-known/oauth-authorization-server or \ - .well-known/openid-configuration found (tried both RFC 8414 \ - path-aware and origin-root locations)" + .well-known/openid-configuration found" ))) } @@ -356,69 +316,6 @@ impl McpOAuthService { .map_err(|e| McpError::OAuth(format!("Failed to parse metadata: {e}"))) } - /// Wait for the OAuth callback redirect on the given listener. - async fn wait_for_callback(&self, user_id: &str, listener: TcpListener) -> Result<(String, String), McpError> { - let (code_tx, code_rx) = tokio::sync::oneshot::channel::>(); - let pending = self.pending.clone(); - let user_id = user_id.to_string(); - - tokio::spawn(async move { - let result = Self::handle_callback_connection(&user_id, listener, pending).await; - let _ = code_tx.send(result); - }); - - match tokio::time::timeout(CALLBACK_TIMEOUT, code_rx).await { - Ok(Ok(result)) => result, - Ok(Err(_)) => Err(McpError::OAuth("Callback channel closed unexpectedly".to_string())), - Err(_) => Err(McpError::OAuth( - "OAuth callback timed out — no redirect received within 120s".to_string(), - )), - } - } - - /// Handle a single HTTP connection on the callback server. - async fn handle_callback_connection( - user_id: &str, - listener: TcpListener, - pending: Arc>>, - ) -> Result<(String, String), McpError> { - let (mut stream, _) = listener - .accept() - .await - .map_err(|e| McpError::OAuth(format!("Failed to accept connection: {e}")))?; - - let mut buf = vec![0u8; 4096]; - let n = stream - .read(&mut buf) - .await - .map_err(|e| McpError::OAuth(format!("Failed to read request: {e}")))?; - - let request = String::from_utf8_lossy(&buf[..n]); - let (code, state) = parse_callback_query(&request)?; - - // Validate CSRF state. - let guard = pending.lock().await; - let pending_login = guard - .get(&(user_id.to_string(), state.clone())) - .ok_or_else(|| McpError::OAuth("No pending login state".to_string()))?; - - if state != *pending_login.csrf_token.secret() { - return Err(McpError::OAuth("CSRF state mismatch".to_string())); - } - - // Send a success response to the browser. - let response = "HTTP/1.1 200 OK\r\n\ - Content-Type: text/html; charset=utf-8\r\n\ - Connection: close\r\n\r\n\ -

Authorization successful!

\ -

You can close this window and return to AionUi.

\ - "; - - let _ = stream.write_all(response.as_bytes()).await; - - Ok((code, state)) - } - /// Build a no-redirect reqwest client for OAuth token exchange. fn build_no_redirect_client() -> Result { reqwest::ClientBuilder::new() @@ -428,14 +325,12 @@ impl McpOAuthService { } /// Exchange the authorization code for tokens and persist them. - async fn exchange_code( - &self, - user_id: &str, - server_url: &str, - code: String, - state: String, - ) -> Result<(), McpError> { - let (auth_url_str, token_url_str, redirect_url_str, pkce_verifier) = { + /// + /// `server_url` is not a parameter — it comes from the `PendingLogin` + /// stashed at login time, since the browser's callback redirect only + /// carries `code` and `state`. + async fn exchange_code(&self, user_id: &str, code: String, state: String) -> Result<(), McpError> { + let (auth_url_str, token_url_str, redirect_url_str, pkce_verifier, server_url) = { let mut guard = self.pending.lock().await; let pending = guard .remove(&(user_id.to_string(), state)) @@ -445,6 +340,7 @@ impl McpOAuthService { pending.token_url, pending.redirect_url, pending.pkce_verifier, + pending.server_url, ) }; @@ -467,7 +363,7 @@ impl McpOAuthService { .await .map_err(|e| McpError::OAuth(format!("Token exchange failed: {e}")))?; - self.persist_token(user_id, server_url, &token_result).await?; + self.persist_token(user_id, &server_url, &token_result).await?; debug!(server_url, "OAuth tokens stored successfully"); Ok(()) } @@ -553,76 +449,6 @@ impl McpOAuthService { // Query parameter parsing // --------------------------------------------------------------------------- -/// Parse `code` and `state` from the first line of an HTTP request. -/// -/// Expects: `GET /callback?code=xxx&state=yyy HTTP/1.1` -fn parse_callback_query(request: &str) -> Result<(String, String), McpError> { - let first_line = request - .lines() - .next() - .ok_or_else(|| McpError::OAuth("Empty HTTP request".to_string()))?; - - let path = first_line - .split_whitespace() - .nth(1) - .ok_or_else(|| McpError::OAuth("Malformed HTTP request line".to_string()))?; - - let query_str = path - .split_once('?') - .map(|(_, q)| q) - .ok_or_else(|| McpError::OAuth("No query parameters in callback".to_string()))?; - - let mut code = None; - let mut state = None; - - for pair in query_str.split('&') { - if let Some((key, value)) = pair.split_once('=') { - match key { - "code" => code = Some(url_decode(value)), - "state" => state = Some(url_decode(value)), - _ => {} - } - } - } - - let code = code.ok_or_else(|| McpError::OAuth("Missing 'code' in callback".to_string()))?; - let state = state.ok_or_else(|| McpError::OAuth("Missing 'state' in callback".to_string()))?; - - Ok((code, state)) -} - -/// Minimal percent-decoding for query parameter values. -fn url_decode(input: &str) -> String { - let mut result = String::with_capacity(input.len()); - let mut chars = input.bytes(); - - while let Some(b) = chars.next() { - if b == b'%' { - let hi = chars.next(); - let lo = chars.next(); - if let (Some(h), Some(l)) = (hi, lo) { - let hex = [h, l]; - if let Ok(s) = std::str::from_utf8(&hex) - && let Ok(byte) = u8::from_str_radix(s, 16) - { - result.push(byte as char); - continue; - } - // Malformed percent-encoding: keep as-is. - result.push('%'); - result.push(h as char); - result.push(l as char); - } - } else if b == b'+' { - result.push(' '); - } else { - result.push(b as char); - } - } - - result -} - // --------------------------------------------------------------------------- // Unit tests // --------------------------------------------------------------------------- @@ -633,161 +459,80 @@ mod tests { const TEST_USER_ID: &str = "user-1"; - // -- well_known_urls ------------------------------------------------------- - // - // Regression coverage for a bug where the well-known discovery suffix was - // appended after the server URL's full path (`https://host/mcp/.well-known/ - // oauth-authorization-server`), which 404s against any RFC 8414-compliant - // server, instead of being inserted right after the origin per RFC 8414 §3.1 - // (`https://host/.well-known/oauth-authorization-server/mcp`). - - #[test] - fn well_known_urls_path_aware_form_comes_first_for_pathed_server() { - let candidates = well_known_urls("https://mcp.example.com/mcp").unwrap(); - let urls = candidates.for_type("oauth-authorization-server"); - assert_eq!( - urls, - vec![ - "https://mcp.example.com/.well-known/oauth-authorization-server/mcp", - "https://mcp.example.com/.well-known/oauth-authorization-server", - ] - ); - } - - #[test] - fn well_known_urls_never_puts_well_known_after_the_original_path() { - let candidates = well_known_urls("https://mcp.example.com/mcp").unwrap(); - for url in candidates.for_type("oauth-authorization-server") { - assert!( - !url.starts_with("https://mcp.example.com/mcp/.well-known"), - "well-known suffix must not be appended after the server's path: {url}" - ); - } - } - - #[test] - fn well_known_urls_root_only_server_has_a_single_candidate() { - let candidates = well_known_urls("https://mcp.example.com").unwrap(); - assert_eq!( - candidates.for_type("oauth-authorization-server"), - vec!["https://mcp.example.com/.well-known/oauth-authorization-server"] - ); - } - - #[test] - fn well_known_urls_strips_trailing_slash_from_path() { - let candidates = well_known_urls("https://mcp.example.com/mcp/").unwrap(); - let urls = candidates.for_type("openid-configuration"); - assert_eq!(urls[0], "https://mcp.example.com/.well-known/openid-configuration/mcp"); - } - - #[test] - fn well_known_urls_rejects_invalid_server_url() { - let err = well_known_urls("not a url").unwrap_err(); - assert!(err.to_string().contains("Invalid server URL")); - } - - // -- parse_callback_query ------------------------------------------------ - - #[test] - fn parse_valid_callback_query() { - let request = "GET /callback?code=abc123&state=xyz789 HTTP/1.1\r\nHost: localhost\r\n"; - let (code, state) = parse_callback_query(request).unwrap(); - assert_eq!(code, "abc123"); - assert_eq!(state, "xyz789"); - } - - #[test] - fn parse_callback_query_reversed_params() { - let request = "GET /callback?state=s1&code=c1 HTTP/1.1\r\n"; - let (code, state) = parse_callback_query(request).unwrap(); - assert_eq!(code, "c1"); - assert_eq!(state, "s1"); - } - - #[test] - fn parse_callback_query_with_extra_params() { - let request = "GET /callback?code=c&foo=bar&state=s HTTP/1.1\r\n"; - let (code, state) = parse_callback_query(request).unwrap(); - assert_eq!(code, "c"); - assert_eq!(state, "s"); - } + // -- McpOAuthService construction ---------------------------------------- #[test] - fn parse_callback_query_missing_code() { - let request = "GET /callback?state=s HTTP/1.1\r\n"; - let err = parse_callback_query(request).unwrap_err(); - assert!(err.to_string().contains("Missing 'code'")); + fn service_clone_is_independent() { + let repo: Arc = Arc::new(MockTokenRepo); + let http = reqwest::Client::new(); + let svc = McpOAuthService::new(repo, http); + let _clone = svc.clone(); } - #[test] - fn parse_callback_query_missing_state() { - let request = "GET /callback?code=c HTTP/1.1\r\n"; - let err = parse_callback_query(request).unwrap_err(); - assert!(err.to_string().contains("Missing 'state'")); - } + #[tokio::test] + async fn clear_pending_for_user_keeps_other_user_same_oauth_state() { + let svc = McpOAuthService::new(Arc::new(MockTokenRepo), reqwest::Client::new()); + insert_pending_login(&svc, "user-a", "shared-state").await; + insert_pending_login(&svc, "user-b", "shared-state").await; - #[test] - fn parse_callback_query_no_query_string() { - let request = "GET /callback HTTP/1.1\r\n"; - let err = parse_callback_query(request).unwrap_err(); - assert!(err.to_string().contains("No query parameters")); - } + svc.clear_pending_for_user("user-a").await; - #[test] - fn parse_callback_query_empty_request() { - let err = parse_callback_query("").unwrap_err(); - assert!(err.to_string().contains("Empty HTTP request")); + let pending = svc.pending.lock().await; + assert!(!pending.contains_key(&("user-a".to_string(), "shared-state".to_string()))); + assert!(pending.contains_key(&("user-b".to_string(), "shared-state".to_string()))); } - // -- url_decode ---------------------------------------------------------- - - #[test] - fn url_decode_no_encoding() { - assert_eq!(url_decode("hello"), "hello"); - } + // -- handle_callback ------------------------------------------------------- + // + // Regression coverage for routing the OAuth callback through this + // server's own HTTP router (keyed by (user_id, state) in `pending`) + // instead of a per-login localhost TCP listener. - #[test] - fn url_decode_percent_encoded() { - assert_eq!(url_decode("hello%20world"), "hello world"); - } + #[tokio::test] + async fn handle_callback_errors_for_unknown_state() { + let svc = McpOAuthService::new(Arc::new(MockTokenRepo), reqwest::Client::new()); - #[test] - fn url_decode_plus_sign() { - assert_eq!(url_decode("hello+world"), "hello world"); - } + let err = svc + .handle_callback(TEST_USER_ID, "some-code".to_string(), "never-issued-state".to_string()) + .await + .unwrap_err(); - #[test] - fn url_decode_special_characters() { - assert_eq!(url_decode("%3D%26%3F"), "=&?"); + assert!(err.to_string().contains("No pending login state")); } - #[test] - fn url_decode_mixed() { - assert_eq!(url_decode("a%20b+c%3Dd"), "a b c=d"); - } + #[tokio::test] + async fn handle_callback_rejects_another_users_state() { + let svc = McpOAuthService::new(Arc::new(MockTokenRepo), reqwest::Client::new()); + insert_pending_login(&svc, "user-a", "state-a").await; - // -- McpOAuthService construction ---------------------------------------- + // "user-b" trying user-a's state — the (user_id, state) key won't + // match, so this must fail rather than authenticate as user-a. + let err = svc + .handle_callback("user-b", "some-code".to_string(), "state-a".to_string()) + .await + .unwrap_err(); - #[test] - fn service_clone_is_independent() { - let repo: Arc = Arc::new(MockTokenRepo); - let http = reqwest::Client::new(); - let svc = McpOAuthService::new(repo, http); - let _clone = svc.clone(); + assert!(err.to_string().contains("No pending login state")); + // user-a's own pending state is untouched by user-b's failed attempt. + let pending = svc.pending.lock().await; + assert!(pending.contains_key(&("user-a".to_string(), "state-a".to_string()))); } #[tokio::test] - async fn clear_pending_for_user_keeps_other_user_same_oauth_state() { + async fn handle_callback_clears_pending_state_on_failure() { let svc = McpOAuthService::new(Arc::new(MockTokenRepo), reqwest::Client::new()); - insert_pending_login(&svc, "user-a", "shared-state").await; - insert_pending_login(&svc, "user-b", "shared-state").await; + insert_pending_login(&svc, TEST_USER_ID, "state-x").await; - svc.clear_pending_for_user("user-a").await; + // The stashed auth/token URLs point nowhere real, so the token + // exchange itself will fail (network error) after the state lookup + // succeeds. Either way, a failed callback must not leave stale + // pending state behind for a retry to trip over. + let _ = svc + .handle_callback(TEST_USER_ID, "some-code".to_string(), "state-x".to_string()) + .await; let pending = svc.pending.lock().await; - assert!(!pending.contains_key(&("user-a".to_string(), "shared-state".to_string()))); - assert!(pending.contains_key(&("user-b".to_string(), "shared-state".to_string()))); + assert!(!pending.contains_key(&(TEST_USER_ID.to_string(), "state-x".to_string()))); } // -- Mock repositories --------------------------------------------------- @@ -798,11 +543,11 @@ mod tests { pending.insert( (user_id.to_string(), state.to_string()), PendingLogin { - csrf_token: CsrfToken::new(state.to_string()), pkce_verifier, auth_url: "https://auth.example.com/authorize".to_string(), token_url: "https://auth.example.com/token".to_string(), redirect_url: "http://127.0.0.1/callback".to_string(), + server_url: "https://mcp.example.com".to_string(), }, ); } diff --git a/crates/aionui-mcp/src/routes.rs b/crates/aionui-mcp/src/routes.rs index 671818ab1..271b3dfab 100644 --- a/crates/aionui-mcp/src/routes.rs +++ b/crates/aionui-mcp/src/routes.rs @@ -2,10 +2,11 @@ use axum::Router; use axum::extract::rejection::JsonRejection; -use axum::extract::{Extension, Json, Path, State}; -use axum::http::StatusCode; -use axum::response::{IntoResponse, Response}; +use axum::extract::{Extension, Json, Path, Query, State}; +use axum::http::{HeaderMap, StatusCode, header}; +use axum::response::{Html, IntoResponse, Response}; use axum::routing::{get, post}; +use serde::Deserialize; use aionui_api_types::{ ApiResponse, BatchImportMcpServersRequest, CreateMcpServerRequest, DetectedMcpServerResponse, ErrorResponse, @@ -76,6 +77,7 @@ pub fn mcp_routes(state: McpRouterState) -> Router { // OAuth routes .route("/api/mcp/oauth/check-status", post(oauth_check_status)) .route("/api/mcp/oauth/login", post(oauth_login)) + .route("/api/mcp/oauth/callback", get(oauth_callback)) .route("/api/mcp/oauth/logout", post(oauth_logout)) .route("/api/mcp/oauth/authenticated", get(oauth_authenticated)) .with_state(state) @@ -292,22 +294,78 @@ async fn oauth_check_status( /// `POST /api/mcp/oauth/login` — start OAuth PKCE login flow. /// -/// Discovers endpoints, opens the browser for authorization, waits for -/// the callback, and exchanges the code for tokens. +/// Discovers endpoints and returns an `authorize_url` for the caller to send +/// a browser to; it does not wait for login to complete. `redirect_base` is +/// derived from the inbound request so the callback lands back on this same +/// server, reachable from wherever this server itself is reachable from. async fn oauth_login( State(state): State, Extension(user): Extension, + headers: HeaderMap, body: Result, JsonRejection>, ) -> Result>, ApiError> { let Json(req) = body.map_err(ApiError::from)?; + let redirect_base = request_origin(&headers); let result = state .oauth_service - .login(&user.id, &req.server_url) + .login(&user.id, &req.server_url, &redirect_base) .await .map_err(ApiError::from)?; Ok(Json(ApiResponse::ok(result))) } +/// Derive this server's own public origin from the inbound request, for use +/// as the OAuth redirect base. Prefers `Origin` (sent by browser `fetch()` +/// calls); falls back to `X-Forwarded-Proto` + `Host` for requests without +/// one (e.g. `curl`, or browsers omitting `Origin` on some requests). +fn request_origin(headers: &HeaderMap) -> String { + if let Some(origin) = headers.get(header::ORIGIN).and_then(|v| v.to_str().ok()) { + return origin.trim_end_matches('/').to_string(); + } + let scheme = headers + .get("x-forwarded-proto") + .and_then(|v| v.to_str().ok()) + .unwrap_or("http"); + let host = headers + .get(header::HOST) + .and_then(|v| v.to_str().ok()) + .unwrap_or("127.0.0.1"); + format!("{scheme}://{host}") +} + +/// Query params on the OAuth provider's redirect to `GET /api/mcp/oauth/callback`. +#[derive(Debug, Deserialize)] +struct OAuthCallbackQuery { + code: String, + state: String, +} + +/// `GET /api/mcp/oauth/callback` — OAuth provider redirects the user's +/// browser here after they authorize (or deny) access. Exchanges the code +/// for tokens and shows a plain result page; the caller is expected to +/// re-check `/api/mcp/oauth/check-status` (e.g. on navigating back) rather +/// than parse this page. +async fn oauth_callback( + State(state): State, + Extension(user): Extension, + Query(params): Query, +) -> Html<&'static str> { + match state + .oauth_service + .handle_callback(&user.id, params.code, params.state) + .await + { + Ok(()) => Html( + "

Authorization successful

\ +

You can close this tab and return to Nabd Agentic OS.

", + ), + Err(_) => Html( + "

Authorization failed

\ +

Something went wrong completing sign-in. Close this tab and try again.

", + ), + } +} + /// `POST /api/mcp/oauth/logout` — delete stored OAuth token. async fn oauth_logout( State(state): State, diff --git a/crates/aionui-mcp/tests/oauth_integration.rs b/crates/aionui-mcp/tests/oauth_integration.rs index d77690568..f5a09551a 100644 --- a/crates/aionui-mcp/tests/oauth_integration.rs +++ b/crates/aionui-mcp/tests/oauth_integration.rs @@ -167,9 +167,15 @@ async fn get_authenticated_servers_empty_when_no_tokens() { async fn login_invalid_url_returns_error() { let (svc, _repo) = make_service().await; // This URL won't have .well-known endpoints. - let result = svc.login(TEST_USER_ID, "https://127.0.0.1:1").await; - // Should return an McpError::OAuth about discovery failure. - assert!(result.is_err()); + let result = svc + .login(TEST_USER_ID, "https://127.0.0.1:1", "http://127.0.0.1:8080") + .await + .unwrap(); + // login() itself never hard-fails — discovery failure surfaces as a + // structured `success: false` response, same as any other login error. + assert!(!result.success); + assert!(result.authorize_url.is_none()); + assert!(result.error.is_some()); } // --------------------------------------------------------------------------- From d6dc36d28c73a98905f0c4c61f86c3634bee395e Mon Sep 17 00:00:00 2001 From: Hani Akrim Date: Tue, 18 Aug 2026 14:15:29 +0300 Subject: [PATCH 3/7] fix(mcp): redirect the OAuth callback success page back into the app The callback page told the user to close the tab manually, leaving them stranded outside the app after a full-page redirect through the authorize_url. Add a 2s meta-refresh back to the app root so the common case (things worked) returns them automatically. --- crates/aionui-mcp/src/routes.rs | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/crates/aionui-mcp/src/routes.rs b/crates/aionui-mcp/src/routes.rs index 271b3dfab..cf543fcef 100644 --- a/crates/aionui-mcp/src/routes.rs +++ b/crates/aionui-mcp/src/routes.rs @@ -356,8 +356,9 @@ async fn oauth_callback( .await { Ok(()) => Html( - "

Authorization successful

\ -

You can close this tab and return to Nabd Agentic OS.

", + "\ +

Authorization successful

\ +

Redirecting you back — you can also close this tab.

", ), Err(_) => Html( "

Authorization failed

\ From 5321bb16d20bd2fb1750b2161e2e778c194bc375 Mon Sep 17 00:00:00 2001 From: Hani Akrim Date: Tue, 18 Aug 2026 14:24:58 +0300 Subject: [PATCH 4/7] fix(mcp): restore RFC 8414 discovery fix lost to a stray checkout, reject unsafe OAuth endpoint schemes The well_known_urls()/WellKnownCandidates fix from the first commit on this branch was silently reverted in the working tree by a `git checkout main -- .` run between that commit and the next one, and the regression went uncaught because its own tests were reverted with it (the passing "0 failed" test runs afterward just meant those specific assertions no longer existed, not that they passed). This restores that fix intact, with its 5 original regression tests. Also fixes a real finding from an independent security review of the authorize_url-navigation change: `authorization_endpoint` and `token_endpoint` come from the MCP server's own self-published OAuth discovery document, so they're attacker-influenced input, not values this server generates. Without a scheme check, a malicious or compromised MCP server could publish a `javascript:`/`data:` URL as its authorization_endpoint, which would flow through unchanged into the authorize_url returned to the client and executed by `window.location.href = authorize_url` on the frontend. Reject anything but http(s) centrally in fetch_metadata, once, rather than at each of the several downstream places that build a client or hand a URL to a browser from these fields. --- crates/aionui-mcp/src/oauth_service.rs | 186 ++++++++++++++++++++++--- 1 file changed, 170 insertions(+), 16 deletions(-) diff --git a/crates/aionui-mcp/src/oauth_service.rs b/crates/aionui-mcp/src/oauth_service.rs index 9f14b7593..152e86538 100644 --- a/crates/aionui-mcp/src/oauth_service.rs +++ b/crates/aionui-mcp/src/oauth_service.rs @@ -36,6 +36,54 @@ struct OAuthServerMetadata { token_endpoint: String, } +/// Candidate well-known discovery URLs for one MCP server, in the order they +/// should be tried: RFC 8414 §3.1 path-aware form first, then the bare-origin +/// form some servers publish instead. +#[derive(Debug)] +struct WellKnownCandidates { + origin: String, + path: String, +} + +impl WellKnownCandidates { + /// URLs to try for a given well-known type (e.g. "oauth-authorization-server"). + fn for_type(&self, well_known_type: &str) -> Vec { + let mut urls = vec![format!("{}/.well-known/{well_known_type}{}", self.origin, self.path)]; + if !self.path.is_empty() { + urls.push(format!("{}/.well-known/{well_known_type}", self.origin)); + } + urls + } +} + +/// Build the well-known discovery URL candidates for an MCP server URL. +/// +/// Per RFC 8414 §3.1, when the issuer URL has a path component (as MCP +/// server URLs like `https://host/mcp` typically do), the well-known suffix +/// is inserted between the authority and that path — it is NOT appended +/// after the full URL. `https://host/mcp` therefore discovers at +/// `https://host/.well-known/oauth-authorization-server/mcp`, not +/// `https://host/mcp/.well-known/oauth-authorization-server` (which 404s +/// against any RFC 8414-compliant server). +fn well_known_urls(server_url: &str) -> Result { + let parsed = + oauth2::url::Url::parse(server_url).map_err(|e| McpError::OAuth(format!("Invalid server URL: {e}")))?; + Ok(WellKnownCandidates { + origin: parsed.origin().ascii_serialization(), + path: parsed.path().trim_end_matches('/').to_string(), + }) +} + +/// The MCP server's own OAuth discovery document controls `authorization_endpoint` +/// and `token_endpoint` — treat them as untrusted input, not values this +/// server generated. Reject anything but http(s) before ever building a +/// client or handing a URL to a browser from these fields. +fn is_safe_http_url(url: &str) -> bool { + oauth2::url::Url::parse(url) + .map(|u| u.scheme() == "https" || u.scheme() == "http") + .unwrap_or(false) +} + // --------------------------------------------------------------------------- // Pending login state // --------------------------------------------------------------------------- @@ -274,31 +322,42 @@ impl McpOAuthService { /// Discover OAuth authorization server metadata. /// - /// Tries `.well-known/oauth-authorization-server` first, - /// falls back to `.well-known/openid-configuration`. + /// Per RFC 8414 §3.1, when the issuer URL has a path component (as MCP + /// server URLs like `https://host/mcp` typically do), the well-known + /// suffix is inserted between the authority and that path — it is NOT + /// appended after the full URL. `https://host/mcp` therefore discovers + /// at `https://host/.well-known/oauth-authorization-server/mcp`, not + /// `https://host/mcp/.well-known/oauth-authorization-server` (which + /// 404s against any RFC 8414-compliant server). Falls back to the + /// bare-origin well-known path for servers that publish metadata there + /// instead, then tries OIDC discovery the same way. async fn discover_endpoints(&self, server_url: &str) -> Result { - let base = server_url.trim_end_matches('/'); + let candidates = well_known_urls(server_url)?; - let well_known_url = format!("{base}/.well-known/oauth-authorization-server"); - if let Ok(metadata) = self.fetch_metadata(&well_known_url).await { - debug!(server_url, "Discovered OAuth metadata via RFC 8414"); - return Ok(metadata); - } - - let oidc_url = format!("{base}/.well-known/openid-configuration"); - if let Ok(metadata) = self.fetch_metadata(&oidc_url).await { - debug!(server_url, "Discovered OAuth metadata via OIDC"); - return Ok(metadata); + for well_known_type in ["oauth-authorization-server", "openid-configuration"] { + for url in candidates.for_type(well_known_type) { + if let Ok(metadata) = self.fetch_metadata(&url).await { + debug!(server_url, url, "Discovered OAuth metadata"); + return Ok(metadata); + } + } } Err(McpError::OAuth(format!( "Failed to discover OAuth endpoints for '{server_url}': \ no .well-known/oauth-authorization-server or \ - .well-known/openid-configuration found" + .well-known/openid-configuration found (tried both RFC 8414 \ + path-aware and origin-root locations)" ))) } /// Fetch and parse OAuth server metadata from a URL. + /// + /// The MCP server's own discovery document controls `authorization_endpoint` + /// and `token_endpoint` — treat them as untrusted input, not values this + /// server generated. Reject anything but http(s) here, once, centrally, + /// rather than at each of the several places downstream that build a + /// `reqwest`/`oauth2` client or hand a URL to a browser from these fields. async fn fetch_metadata(&self, url: &str) -> Result { let resp = self .http_client @@ -311,9 +370,18 @@ impl McpOAuthService { return Err(McpError::OAuth(format!("Metadata endpoint returned {}", resp.status()))); } - resp.json() + let metadata: OAuthServerMetadata = resp + .json() .await - .map_err(|e| McpError::OAuth(format!("Failed to parse metadata: {e}"))) + .map_err(|e| McpError::OAuth(format!("Failed to parse metadata: {e}")))?; + + if !is_safe_http_url(&metadata.authorization_endpoint) || !is_safe_http_url(&metadata.token_endpoint) { + return Err(McpError::OAuth( + "OAuth metadata endpoint URLs must use http or https".to_string(), + )); + } + + Ok(metadata) } /// Build a no-redirect reqwest client for OAuth token exchange. @@ -459,6 +527,92 @@ mod tests { const TEST_USER_ID: &str = "user-1"; + // -- well_known_urls ------------------------------------------------------- + // + // Regression coverage for a bug where the well-known discovery suffix was + // appended after the server URL's full path (`https://host/mcp/.well-known/ + // oauth-authorization-server`), which 404s against any RFC 8414-compliant + // server, instead of being inserted right after the origin per RFC 8414 §3.1 + // (`https://host/.well-known/oauth-authorization-server/mcp`). + + #[test] + fn well_known_urls_path_aware_form_comes_first_for_pathed_server() { + let candidates = well_known_urls("https://mcp.example.com/mcp").unwrap(); + let urls = candidates.for_type("oauth-authorization-server"); + assert_eq!( + urls, + vec![ + "https://mcp.example.com/.well-known/oauth-authorization-server/mcp", + "https://mcp.example.com/.well-known/oauth-authorization-server", + ] + ); + } + + #[test] + fn well_known_urls_never_puts_well_known_after_the_original_path() { + let candidates = well_known_urls("https://mcp.example.com/mcp").unwrap(); + for url in candidates.for_type("oauth-authorization-server") { + assert!( + !url.starts_with("https://mcp.example.com/mcp/.well-known"), + "well-known suffix must not be appended after the server's path: {url}" + ); + } + } + + #[test] + fn well_known_urls_root_only_server_has_a_single_candidate() { + let candidates = well_known_urls("https://mcp.example.com").unwrap(); + assert_eq!( + candidates.for_type("oauth-authorization-server"), + vec!["https://mcp.example.com/.well-known/oauth-authorization-server"] + ); + } + + #[test] + fn well_known_urls_strips_trailing_slash_from_path() { + let candidates = well_known_urls("https://mcp.example.com/mcp/").unwrap(); + let urls = candidates.for_type("openid-configuration"); + assert_eq!(urls[0], "https://mcp.example.com/.well-known/openid-configuration/mcp"); + } + + #[test] + fn well_known_urls_rejects_invalid_server_url() { + let err = well_known_urls("not a url").unwrap_err(); + assert!(err.to_string().contains("Invalid server URL")); + } + + // -- is_safe_http_url -------------------------------------------------- + // + // The MCP server's own OAuth discovery document controls the endpoint + // URLs this server later navigates a browser to — reject anything but + // http(s) so a malicious/misconfigured server can't get a javascript: + // or data: URL into that redirect. + + #[test] + fn is_safe_http_url_accepts_https() { + assert!(is_safe_http_url("https://mcp.example.com/oauth2/authorize")); + } + + #[test] + fn is_safe_http_url_accepts_http() { + assert!(is_safe_http_url("http://127.0.0.1:8080/oauth2/authorize")); + } + + #[test] + fn is_safe_http_url_rejects_javascript_scheme() { + assert!(!is_safe_http_url("javascript:alert(document.cookie)")); + } + + #[test] + fn is_safe_http_url_rejects_data_scheme() { + assert!(!is_safe_http_url("data:text/html,")); + } + + #[test] + fn is_safe_http_url_rejects_unparseable_url() { + assert!(!is_safe_http_url("not a url")); + } + // -- McpOAuthService construction ---------------------------------------- #[test] From 70d83a0e71b3a039d02d3be5628b05c1f8ab46fe Mon Sep 17 00:00:00 2001 From: Hani Akrim Date: Tue, 18 Aug 2026 14:59:04 +0300 Subject: [PATCH 5/7] feat(mcp): support RFC 7591 Dynamic Client Registration for MCP OAuth MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Login against Higgsfield's MCP server (Clerk-backed auth) failed at the authorize step with "invalid_client ... The requested OAuth 2.0 Client does not exist" — aioncore always used a fixed shared client_id ("aionui") for every MCP server's OAuth flow, but Higgsfield's discovery document advertises a registration_endpoint, meaning RFC 7591 Dynamic Client Registration is mandatory there, not optional; a static client_id is simply never recognized. When discovery returns a registration_endpoint, register a client against it (public client, PKCE-secured, no secret requested) and use the issued client_id (and secret, if the server returns one anyway) for both the authorize_url and the later token exchange - a mismatch between the two would fail the same way. Servers without a registration_endpoint keep using DEFAULT_CLIENT_ID exactly as before. Registered clients are cached in memory only, not persisted: adding a new DB migration for this felt like the wrong tradeoff to reach for again in the same session a migration-version mismatch already caused a production crash loop. A process restart re-registers, which RFC 7591 is designed to tolerate. --- crates/aionui-mcp/src/oauth_service.rs | 217 ++++++++++++++++++++++++- 1 file changed, 211 insertions(+), 6 deletions(-) diff --git a/crates/aionui-mcp/src/oauth_service.rs b/crates/aionui-mcp/src/oauth_service.rs index 152e86538..5ffa50899 100644 --- a/crates/aionui-mcp/src/oauth_service.rs +++ b/crates/aionui-mcp/src/oauth_service.rs @@ -6,8 +6,8 @@ use aionui_common::{TimestampMs, now_ms}; use aionui_db::{IOAuthTokenRepository, UpsertOAuthTokenParams}; use oauth2::basic::BasicClient; use oauth2::{ - AuthUrl, AuthorizationCode, ClientId, CsrfToken, PkceCodeChallenge, PkceCodeVerifier, RedirectUrl, RefreshToken, - TokenResponse, TokenUrl, + AuthUrl, AuthorizationCode, ClientId, ClientSecret, CsrfToken, PkceCodeChallenge, PkceCodeVerifier, RedirectUrl, + RefreshToken, TokenResponse, TokenUrl, }; use serde::Deserialize; use tokio::sync::Mutex; @@ -34,6 +34,19 @@ const EXPIRY_MARGIN_MS: i64 = 5 * 60 * 1000; struct OAuthServerMetadata { authorization_endpoint: String, token_endpoint: String, + /// RFC 7591 Dynamic Client Registration endpoint. When present, the + /// authorization server expects each client to register itself and get + /// its own `client_id` — a fixed shared `client_id` (like our + /// `DEFAULT_CLIENT_ID`) won't be recognized (see `ensure_client`). + registration_endpoint: Option, +} + +/// A client_id (and optional secret) issued by an authorization server's +/// RFC 7591 Dynamic Client Registration endpoint. +#[derive(Debug, Clone, Deserialize)] +struct DynamicClient { + client_id: String, + client_secret: Option, } /// Candidate well-known discovery URLs for one MCP server, in the order they @@ -100,6 +113,12 @@ struct PendingLogin { token_url: String, redirect_url: String, server_url: String, + /// The client_id used to build the authorize_url — either + /// `DEFAULT_CLIENT_ID` or one issued by Dynamic Client Registration. + /// The same client_id (and secret, if any) must be used again for the + /// token exchange, or a DCR-registered server will reject it. + client_id: String, + client_secret: Option, } // --------------------------------------------------------------------------- @@ -116,6 +135,9 @@ pub struct McpOAuthService { http_client: reqwest::Client, /// Mutex protecting pending login state by (user_id, oauth_state). pending: Arc>>, + /// In-memory cache of RFC 7591 dynamically-registered clients, keyed by + /// registration_endpoint. See `ensure_client` for why this isn't persisted. + registered_clients: Arc>>, } impl McpOAuthService { @@ -124,6 +146,7 @@ impl McpOAuthService { token_repo, http_client, pending: Arc::new(Mutex::new(HashMap::new())), + registered_clients: Arc::new(Mutex::new(HashMap::new())), } } @@ -274,10 +297,25 @@ impl McpOAuthService { let redirect = RedirectUrl::new(redirect_url_str.clone()) .map_err(|e| McpError::OAuth(format!("Invalid redirect URL: {e}")))?; - let client = BasicClient::new(ClientId::new(DEFAULT_CLIENT_ID.to_string())) + // Some authorization servers (Higgsfield/Clerk among them) require + // RFC 7591 Dynamic Client Registration and reject a fixed shared + // client_id outright — a registration_endpoint in the discovery + // document means that's mandatory here, not optional. + let (client_id, client_secret) = match metadata.registration_endpoint.as_deref() { + Some(registration_endpoint) => { + let registered = self.ensure_client(registration_endpoint, redirect_base).await?; + (registered.client_id, registered.client_secret) + } + None => (DEFAULT_CLIENT_ID.to_string(), None), + }; + + let mut client = BasicClient::new(ClientId::new(client_id.clone())) .set_auth_uri(auth_url) .set_token_uri(token_url) .set_redirect_uri(redirect); + if let Some(secret) = client_secret.clone() { + client = client.set_client_secret(ClientSecret::new(secret)); + } let (pkce_challenge, pkce_verifier) = PkceCodeChallenge::new_random_sha256(); @@ -297,6 +335,8 @@ impl McpOAuthService { token_url: token_url_str, redirect_url: redirect_url_str, server_url: server_url.to_string(), + client_id, + client_secret, }, ); } @@ -375,7 +415,11 @@ impl McpOAuthService { .await .map_err(|e| McpError::OAuth(format!("Failed to parse metadata: {e}")))?; - if !is_safe_http_url(&metadata.authorization_endpoint) || !is_safe_http_url(&metadata.token_endpoint) { + let registration_endpoint_safe = metadata.registration_endpoint.as_deref().is_none_or(is_safe_http_url); + if !is_safe_http_url(&metadata.authorization_endpoint) + || !is_safe_http_url(&metadata.token_endpoint) + || !registration_endpoint_safe + { return Err(McpError::OAuth( "OAuth metadata endpoint URLs must use http or https".to_string(), )); @@ -384,6 +428,65 @@ impl McpOAuthService { Ok(metadata) } + /// Get (registering if needed) the client_id this server uses to talk to + /// the given authorization server. + /// + /// Some authorization servers (Higgsfield/Clerk among them) advertise a + /// `registration_endpoint` and don't recognize a fixed shared client_id + /// at all — RFC 7591 Dynamic Client Registration is mandatory for them, + /// not optional. Cache the issued client_id in memory per registration + /// endpoint so repeated logins to the same server don't re-register each + /// time; this cache is intentionally not persisted to disk, so a process + /// restart re-registers (harmless — DCR is designed to be repeatable, + /// and avoiding a new DB migration here was a deliberate choice: the + /// previous migration mismatch already caused one production incident + /// tonight). + async fn ensure_client(&self, registration_endpoint: &str, redirect_base: &str) -> Result { + { + let cache = self.registered_clients.lock().await; + if let Some(client) = cache.get(registration_endpoint) { + return Ok(client.clone()); + } + } + + let redirect_uri = format!("{}/api/mcp/oauth/callback", redirect_base.trim_end_matches('/')); + let body = serde_json::json!({ + "client_name": "AionUi", + "redirect_uris": [redirect_uri], + "grant_types": ["authorization_code", "refresh_token"], + "response_types": ["code"], + "token_endpoint_auth_method": "none", + }); + + let resp = self + .http_client + .post(registration_endpoint) + .json(&body) + .send() + .await + .map_err(|e| McpError::OAuth(format!("Dynamic client registration request failed: {e}")))?; + + if !resp.status().is_success() { + return Err(McpError::OAuth(format!( + "Dynamic client registration returned {}", + resp.status() + ))); + } + + let client: DynamicClient = resp + .json() + .await + .map_err(|e| McpError::OAuth(format!("Failed to parse client registration response: {e}")))?; + + self.registered_clients + .lock() + .await + .insert(registration_endpoint.to_string(), client.clone()); + + debug!(registration_endpoint, client_id = %client.client_id, "Registered dynamic OAuth client"); + Ok(client) + } + /// Build a no-redirect reqwest client for OAuth token exchange. fn build_no_redirect_client() -> Result { reqwest::ClientBuilder::new() @@ -398,7 +501,7 @@ impl McpOAuthService { /// stashed at login time, since the browser's callback redirect only /// carries `code` and `state`. async fn exchange_code(&self, user_id: &str, code: String, state: String) -> Result<(), McpError> { - let (auth_url_str, token_url_str, redirect_url_str, pkce_verifier, server_url) = { + let (auth_url_str, token_url_str, redirect_url_str, pkce_verifier, server_url, client_id, client_secret) = { let mut guard = self.pending.lock().await; let pending = guard .remove(&(user_id.to_string(), state)) @@ -409,6 +512,8 @@ impl McpOAuthService { pending.redirect_url, pending.pkce_verifier, pending.server_url, + pending.client_id, + pending.client_secret, ) }; @@ -417,10 +522,17 @@ impl McpOAuthService { let redirect = RedirectUrl::new(redirect_url_str).map_err(|e| McpError::OAuth(format!("Invalid redirect URL: {e}")))?; - let client = BasicClient::new(ClientId::new(DEFAULT_CLIENT_ID.to_string())) + // Must match the client_id (and secret, if any) used to build the + // authorize_url — a DCR-registered server rejects a mismatched or + // fixed shared client_id at token exchange just as it would at + // authorization. + let mut client = BasicClient::new(ClientId::new(client_id)) .set_auth_uri(auth_url) .set_token_uri(token_url) .set_redirect_uri(redirect); + if let Some(secret) = client_secret { + client = client.set_client_secret(ClientSecret::new(secret)); + } let http_client = Self::build_no_redirect_client()?; @@ -608,6 +720,97 @@ mod tests { assert!(!is_safe_http_url("data:text/html,")); } + // -- ensure_client (RFC 7591 Dynamic Client Registration) ---------------- + // + // Higgsfield/Clerk-style authorization servers advertise a + // registration_endpoint and reject a fixed shared client_id outright — + // regression coverage for registering, parsing the response, and caching + // so repeated logins don't re-register every time. + + /// Minimal one-shot HTTP server: accepts a single connection, records the + /// request body, and replies with a fixed JSON body. Good enough to + /// exercise ensure_client's real HTTP round-trip without a mocking crate. + async fn serve_once_and_record_request_count( + response_body: &'static str, + ) -> (String, std::sync::Arc) { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let count = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let count_clone = count.clone(); + + tokio::spawn(async move { + loop { + let Ok((mut stream, _)) = listener.accept().await else { + break; + }; + count_clone.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + let mut buf = vec![0u8; 4096]; + let _ = stream.read(&mut buf).await; + let response = format!( + "HTTP/1.1 201 Created\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}", + response_body.len(), + response_body + ); + let _ = stream.write_all(response.as_bytes()).await; + } + }); + + (format!("http://{addr}"), count) + } + + #[tokio::test] + async fn ensure_client_registers_and_parses_client_id() { + let (registration_endpoint, _count) = + serve_once_and_record_request_count(r#"{"client_id":"issued-client-123"}"#).await; + let svc = McpOAuthService::new(Arc::new(MockTokenRepo), reqwest::Client::new()); + + let client = svc + .ensure_client(®istration_endpoint, "https://host.example.com") + .await + .unwrap(); + + assert_eq!(client.client_id, "issued-client-123"); + assert_eq!(client.client_secret, None); + } + + #[tokio::test] + async fn ensure_client_parses_client_secret_when_present() { + let (registration_endpoint, _count) = + serve_once_and_record_request_count(r#"{"client_id":"c1","client_secret":"s1"}"#).await; + let svc = McpOAuthService::new(Arc::new(MockTokenRepo), reqwest::Client::new()); + + let client = svc + .ensure_client(®istration_endpoint, "https://host.example.com") + .await + .unwrap(); + + assert_eq!(client.client_secret.as_deref(), Some("s1")); + } + + #[tokio::test] + async fn ensure_client_caches_and_does_not_re_register() { + let (registration_endpoint, count) = + serve_once_and_record_request_count(r#"{"client_id":"cached-client"}"#).await; + let svc = McpOAuthService::new(Arc::new(MockTokenRepo), reqwest::Client::new()); + + let first = svc + .ensure_client(®istration_endpoint, "https://host.example.com") + .await + .unwrap(); + let second = svc + .ensure_client(®istration_endpoint, "https://host.example.com") + .await + .unwrap(); + + assert_eq!(first.client_id, second.client_id); + assert_eq!( + count.load(std::sync::atomic::Ordering::SeqCst), + 1, + "second call must be served from cache, not a second registration request" + ); + } + #[test] fn is_safe_http_url_rejects_unparseable_url() { assert!(!is_safe_http_url("not a url")); @@ -702,6 +905,8 @@ mod tests { token_url: "https://auth.example.com/token".to_string(), redirect_url: "http://127.0.0.1/callback".to_string(), server_url: "https://mcp.example.com".to_string(), + client_id: "aionui".to_string(), + client_secret: None, }, ); } From c924f2381b683a2638d21bc6e1770dfd39a5c6d0 Mon Sep 17 00:00:00 2001 From: Hani Akrim Date: Tue, 18 Aug 2026 16:13:53 +0300 Subject: [PATCH 6/7] fix(mcp): attach stored OAuth tokens to actual MCP tool-call requests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Login worked, the token got stored — and nothing ever used it. POST /api/mcp/oauth/login through the callback correctly stores an access token, but when the session stack actually builds an HTTP/SSE MCP server's transport for an agent to call its tools, that path (mcp_resolve::row_to_session_mcp_server) only ever forwards a server's static configured headers. There is no OAuth lookup anywhere in it. check_oauth_status can report authenticated=true while every real tool call the agent makes still gets no Authorization header at all and fails as unauthenticated — OAuth login was, in effect, decorative for headless/agent tool use up to this point. Thread an OAuth token repository through the same path mcp_server_repo already takes (AgentFactoryDeps -> SessionBuildInputs -> resolve_session_mcp_servers -> row_to_session_mcp_server), and attach a stored, non-expired token as `Authorization: Bearer ` when building an HTTP/SSE transport. A server-configured Authorization header (case-insensitive) always wins over an OAuth token — a user who set one explicitly presumably knows what they're doing. Expired tokens are omitted rather than sent and rejected: refresh happens lazily via the check-status/get-token API paths, not this session-build path, so an unrefreshed expired token would just fail auth anyway. 5 new tests: token attached when valid, omitted when expired, omitted with no repo/no stored token, and the header-precedence case. --- crates/aionui-ai-agent/src/factory/acp.rs | 1 + .../src/factory/antigravity.rs | 1 + crates/aionui-ai-agent/src/factory/mod.rs | 6 +- crates/aionui-ai-agent/src/mcp_resolve.rs | 212 ++++++++++++++++-- crates/aionui-ai-agent/src/session_agent.rs | 13 +- .../tests/factory_provider_integration.rs | 1 + crates/aionui-app/src/services.rs | 16 +- crates/aionui-conversation/src/service.rs | 2 +- 8 files changed, 232 insertions(+), 20 deletions(-) diff --git a/crates/aionui-ai-agent/src/factory/acp.rs b/crates/aionui-ai-agent/src/factory/acp.rs index c2ff58554..7b097fef9 100644 --- a/crates/aionui-ai-agent/src/factory/acp.rs +++ b/crates/aionui-ai-agent/src/factory/acp.rs @@ -123,6 +123,7 @@ pub(super) async fn build( session_snapshot: build_context.session_snapshot.as_ref(), backend_session_id: build_context.session_id.clone(), mcp_server_repo: deps.mcp_server_repo.as_ref(), + oauth_token_repo: deps.oauth_token_repo.as_ref(), // The AIONUI_* conversation runtime context the legacy path // injects via apply_acp_launch_policy — forwarded into // SessionConfig.spawn_env so direct-CLI spawns get it too. diff --git a/crates/aionui-ai-agent/src/factory/antigravity.rs b/crates/aionui-ai-agent/src/factory/antigravity.rs index e0df5e991..31eee0deb 100644 --- a/crates/aionui-ai-agent/src/factory/antigravity.rs +++ b/crates/aionui-ai-agent/src/factory/antigravity.rs @@ -133,6 +133,7 @@ pub(super) async fn build( session_snapshot: build_context.session_snapshot.as_ref(), backend_session_id: build_context.session_id.clone(), mcp_server_repo: deps.mcp_server_repo.as_ref(), + oauth_token_repo: deps.oauth_token_repo.as_ref(), runtime_env: &runtime_env, broadcaster: deps.broadcaster.clone(), // Keyed by the resolved catalog row so discovered models/modes diff --git a/crates/aionui-ai-agent/src/factory/mod.rs b/crates/aionui-ai-agent/src/factory/mod.rs index 829b72651..99fd40c5f 100644 --- a/crates/aionui-ai-agent/src/factory/mod.rs +++ b/crates/aionui-ai-agent/src/factory/mod.rs @@ -9,7 +9,7 @@ mod context; use std::path::PathBuf; use std::sync::Arc; -use aionui_db::{IMcpServerRepository, IProviderRepository}; +use aionui_db::{IMcpServerRepository, IOAuthTokenRepository, IProviderRepository}; use aionui_realtime::EventBroadcaster; use futures_util::FutureExt; @@ -41,6 +41,10 @@ pub struct AgentFactoryDeps { /// inject enabled servers into `session/new` (ELECTRON-1JG fix). /// `None` for tests/composition paths that do not need MCP injection. pub mcp_server_repo: Option>, + /// OAuth token repository. Used to attach `Authorization: Bearer` to + /// HTTP/SSE MCP servers a user has logged into. `None` for tests/ + /// composition paths that don't need OAuth-gated MCP servers to work. + pub oauth_token_repo: Option>, /// Subprocess spawner for the clean-slate session model. claude/codex always /// run through `SessionAgentTask` (direct-CLI) instead of the ACP manager, so /// the spawner is unconditionally wired — there is no fallback to the ACP path. diff --git a/crates/aionui-ai-agent/src/mcp_resolve.rs b/crates/aionui-ai-agent/src/mcp_resolve.rs index ffd79294a..9dd73bdb3 100644 --- a/crates/aionui-ai-agent/src/mcp_resolve.rs +++ b/crates/aionui-ai-agent/src/mcp_resolve.rs @@ -15,8 +15,8 @@ use std::sync::Arc; use aionui_api_types::{SessionMcpServer, SessionMcpTransport, TEAM_MCP_SERVER_NAME}; -use aionui_db::IMcpServerRepository; use aionui_db::models::McpServerRow; +use aionui_db::{IMcpServerRepository, IOAuthTokenRepository}; use aionui_realtime::EventBroadcaster; use aionui_runtime::ensure_runtime_command; use tracing::{info, warn}; @@ -39,6 +39,7 @@ pub async fn resolve_session_mcp_servers( selected_ids: Option<&[String]>, conversation_id: &str, _broadcaster: Arc, + oauth_token_repo: Option<&dyn IOAuthTokenRepository>, ) -> Vec { let rows_result = match selected_ids { Some(ids) => repo.list_by_ids_any(user_id, ids).await, @@ -63,7 +64,7 @@ pub async fn resolve_session_mcp_servers( if !selected || row.builtin || row.name == TEAM_MCP_SERVER_NAME { continue; } - match row_to_session_mcp_server(&row).await { + match row_to_session_mcp_server(&row, user_id, oauth_token_repo).await { Ok(server) => servers.push(server), Err(err) => { warn!( @@ -95,7 +96,11 @@ pub async fn resolve_session_mcp_servers( /// snapshot refresh) and the session stack reuse it so the stdio command is /// always normalized exactly once, the same way the direct claude/codex path /// consumes it. -pub async fn row_to_session_mcp_server(row: &McpServerRow) -> Result { +pub async fn row_to_session_mcp_server( + row: &McpServerRow, + user_id: &str, + oauth_token_repo: Option<&dyn IOAuthTokenRepository>, +) -> Result { let value: serde_json::Value = serde_json::from_str(&row.transport_config).map_err(|e| format!("invalid transport_config JSON: {e}"))?; @@ -145,10 +150,9 @@ pub async fn row_to_session_mcp_server(row: &McpServerRow) -> Result { let url = value @@ -156,10 +160,9 @@ pub async fn row_to_session_mcp_server(row: &McpServerRow) -> Result return Err(format!("unknown transport type: {other}")), }; @@ -171,6 +174,48 @@ pub async fn row_to_session_mcp_server(row: &McpServerRow) -> Result` for +/// an HTTP/SSE MCP server, if one exists for this user and URL. +/// +/// A server-configured `Authorization` header (case-insensitive) always +/// wins — a user who set one explicitly presumably knows what they're +/// doing, and OAuth login is opt-in on top of that, not a silent override. +/// Best-effort: an expired token or a repo error just means the server +/// won't have a token attached and the tool call fails with a normal +/// auth-required error downstream, exactly as if OAuth had never run. +async fn inject_oauth_bearer_header( + headers: &mut std::collections::HashMap, + user_id: &str, + server_url: &str, + oauth_token_repo: Option<&dyn IOAuthTokenRepository>, +) { + let Some(repo) = oauth_token_repo else { return }; + if headers.keys().any(|k| k.eq_ignore_ascii_case("authorization")) { + return; + } + + let token = match repo.get_by_url(user_id, server_url).await { + Ok(Some(row)) => row, + Ok(None) => return, + Err(err) => { + warn!(server_url, error = %err, "mcp_resolve: OAuth token lookup failed; continuing without it"); + return; + } + }; + + if let Some(expires_at) = token.expires_at + && aionui_common::now_ms() >= expires_at + { + // Expired and unrefreshed here (refresh happens lazily via the + // check-status/get-token API paths, not this session-build path) — + // an expired token would just fail auth anyway, so omit it rather + // than send a token known to be rejected. + return; + } + + headers.insert("Authorization".to_string(), format!("Bearer {}", token.access_token)); +} + /// Parse a JSON headers object into a `HashMap` (string values only). fn parse_headers(value: Option<&serde_json::Value>) -> std::collections::HashMap { value @@ -290,7 +335,7 @@ mod tests { rows: vec![make_row("docs", true), make_row(TEAM_MCP_SERVER_NAME, true)], fail: false, }; - let servers = resolve_session_mcp_servers(&repo, TEST_USER_ID, None, "conv-1", test_broadcaster()).await; + let servers = resolve_session_mcp_servers(&repo, TEST_USER_ID, None, "conv-1", test_broadcaster(), None).await; assert_eq!(servers.len(), 1); assert_eq!(servers[0].name, "docs"); } @@ -321,8 +366,149 @@ mod tests { ], fail: false, }; - let servers = resolve_session_mcp_servers(&repo, TEST_USER_ID, None, "conv-1", test_broadcaster()).await; + let servers = resolve_session_mcp_servers(&repo, TEST_USER_ID, None, "conv-1", test_broadcaster(), None).await; assert_eq!(servers.len(), 1); assert_eq!(servers[0].name, "docs"); } + + // -- inject_oauth_bearer_header ------------------------------------------ + // + // Regression coverage: OAuth login alone never made a stored token reach + // an actual MCP tool call — nothing attached it as an Authorization + // header when building the session's HTTP/SSE transport. Login worked, + // the token was stored, and every tool call still failed as + // unauthenticated. + + struct MockOAuthRepo { + row: Option, + } + + #[async_trait::async_trait] + impl IOAuthTokenRepository for MockOAuthRepo { + async fn get_by_url( + &self, + _user_id: &str, + _server_url: &str, + ) -> Result, aionui_db::DbError> { + Ok(self.row.clone()) + } + async fn upsert( + &self, + _params: aionui_db::UpsertOAuthTokenParams<'_>, + ) -> Result { + unimplemented!("not needed") + } + async fn delete(&self, _user_id: &str, _server_url: &str) -> Result<(), aionui_db::DbError> { + unimplemented!("not needed") + } + async fn list_authenticated_urls(&self, _user_id: &str) -> Result, aionui_db::DbError> { + unimplemented!("not needed") + } + } + + fn token_row( + access_token: &str, + expires_at: Option, + ) -> aionui_db::models::OAuthTokenRow { + aionui_db::models::OAuthTokenRow { + user_id: TEST_USER_ID.to_owned(), + server_url: "http://127.0.0.1:9999/mcp".to_owned(), + access_token: access_token.to_owned(), + refresh_token: None, + token_type: "bearer".to_owned(), + expires_at, + created_at: 0, + updated_at: 0, + } + } + + #[tokio::test] + async fn attaches_bearer_header_when_valid_token_exists() { + let oauth_repo = MockOAuthRepo { + row: Some(token_row("tok-abc", None)), + }; + let row = make_row("higgsfield", true); + + let server = row_to_session_mcp_server(&row, TEST_USER_ID, Some(&oauth_repo)) + .await + .unwrap(); + + match server.transport { + SessionMcpTransport::StreamableHttp { headers, .. } => { + assert_eq!(headers.get("Authorization"), Some(&"Bearer tok-abc".to_string())); + } + _ => panic!("expected StreamableHttp transport"), + } + } + + #[tokio::test] + async fn omits_header_when_token_is_expired() { + let oauth_repo = MockOAuthRepo { + row: Some(token_row("tok-abc", Some(1))), // 1ms since epoch: always expired + }; + let row = make_row("higgsfield", true); + + let server = row_to_session_mcp_server(&row, TEST_USER_ID, Some(&oauth_repo)) + .await + .unwrap(); + + match server.transport { + SessionMcpTransport::StreamableHttp { headers, .. } => { + assert!(!headers.contains_key("Authorization")); + } + _ => panic!("expected StreamableHttp transport"), + } + } + + #[tokio::test] + async fn omits_header_when_no_oauth_repo_given() { + let row = make_row("higgsfield", true); + + let server = row_to_session_mcp_server(&row, TEST_USER_ID, None).await.unwrap(); + + match server.transport { + SessionMcpTransport::StreamableHttp { headers, .. } => { + assert!(!headers.contains_key("Authorization")); + } + _ => panic!("expected StreamableHttp transport"), + } + } + + #[tokio::test] + async fn omits_header_when_no_token_stored() { + let oauth_repo = MockOAuthRepo { row: None }; + let row = make_row("higgsfield", true); + + let server = row_to_session_mcp_server(&row, TEST_USER_ID, Some(&oauth_repo)) + .await + .unwrap(); + + match server.transport { + SessionMcpTransport::StreamableHttp { headers, .. } => { + assert!(!headers.contains_key("Authorization")); + } + _ => panic!("expected StreamableHttp transport"), + } + } + + #[tokio::test] + async fn a_server_configured_authorization_header_wins_over_oauth() { + let oauth_repo = MockOAuthRepo { + row: Some(token_row("tok-from-oauth", None)), + }; + let mut row = make_row("higgsfield", true); + row.transport_config = + r#"{"url":"http://127.0.0.1:9999/mcp","headers":{"Authorization":"Bearer static-key"}}"#.to_owned(); + + let server = row_to_session_mcp_server(&row, TEST_USER_ID, Some(&oauth_repo)) + .await + .unwrap(); + + match server.transport { + SessionMcpTransport::StreamableHttp { headers, .. } => { + assert_eq!(headers.get("Authorization"), Some(&"Bearer static-key".to_string())); + } + _ => panic!("expected StreamableHttp transport"), + } + } } diff --git a/crates/aionui-ai-agent/src/session_agent.rs b/crates/aionui-ai-agent/src/session_agent.rs index 5eca447c0..bc59c77c8 100644 --- a/crates/aionui-ai-agent/src/session_agent.rs +++ b/crates/aionui-ai-agent/src/session_agent.rs @@ -34,7 +34,7 @@ use crate::shared_kernel::PersistedSessionState; use crate::types::{PromptMediaCaps, SendMessageData}; use aionui_api_types::{AcpBuildExtra, TEAM_MCP_SERVER_NAME}; use aionui_common::AgentType; -use aionui_db::{IAcpSessionRepository, IMcpServerRepository, SaveRuntimeStateParams}; +use aionui_db::{IAcpSessionRepository, IMcpServerRepository, IOAuthTokenRepository, SaveRuntimeStateParams}; use aionui_realtime::EventBroadcaster; const EVENT_CHANNEL_CAPACITY: usize = 512; @@ -1485,6 +1485,12 @@ pub struct SessionBuildInputs<'a> { /// User-configured MCP server repository (feature ELECTRON-1JG). `None` on /// paths that never inject MCP (tests) ⇒ no injection. pub mcp_server_repo: Option<&'a Arc>, + /// OAuth token repository, so HTTP/SSE MCP servers a user has logged + /// into get their stored token attached as an `Authorization: Bearer` + /// header. `None` on paths that never inject MCP (tests) ⇒ OAuth-gated + /// servers get no token and fail their tool calls as unauthenticated, + /// same as if OAuth had never run. + pub oauth_token_repo: Option<&'a Arc>, /// The conversation runtime context env (`AIONUI_USER_ID` / /// `AIONUI_CONVERSATION_ID` / `AIONUI_HELPER_BIN` / `AIONUI_BASE_URL` / /// `AIONUI_RUNTIME_TOKEN`, filled by `apply_conversation_runtime_context`). @@ -1661,6 +1667,7 @@ pub async fn build_antigravity_instance( session_snapshot, backend_session_id, mcp_server_repo, + oauth_token_repo, runtime_env, broadcaster, catalog_writeback, @@ -1684,6 +1691,7 @@ pub async fn build_antigravity_instance( config.mcp_server_ids.as_deref(), &conversation_id, broadcaster.clone(), + oauth_token_repo.map(|r| r.as_ref()), ) .await } @@ -1773,6 +1781,7 @@ pub async fn build_session_instance( session_snapshot, backend_session_id, mcp_server_repo, + oauth_token_repo, runtime_env, broadcaster, catalog_writeback, @@ -1800,6 +1809,7 @@ pub async fn build_session_instance( config.mcp_server_ids.as_deref(), &conversation_id, broadcaster.clone(), + oauth_token_repo.map(|r| r.as_ref()), ) .await } @@ -4883,6 +4893,7 @@ mod build_mapping_tests { session_snapshot: None, backend_session_id: None, mcp_server_repo: Some(&repo), + oauth_token_repo: None, runtime_env: &[], broadcaster, catalog_writeback: None, diff --git a/crates/aionui-ai-agent/tests/factory_provider_integration.rs b/crates/aionui-ai-agent/tests/factory_provider_integration.rs index 106b2a77c..1ce5e868a 100644 --- a/crates/aionui-ai-agent/tests/factory_provider_integration.rs +++ b/crates/aionui-ai-agent/tests/factory_provider_integration.rs @@ -91,6 +91,7 @@ fn make_factory( broadcaster: Arc::new(BroadcastEventBus::new(16)), backend_binary_path: Arc::new(PathBuf::from("/tmp/aionrs-test/aioncore")), mcp_server_repo: None, + oauth_token_repo: None, session_spawner, // No hook bridge in this test: it exercises provider wiring, not the // Antigravity permission path. diff --git a/crates/aionui-app/src/services.rs b/crates/aionui-app/src/services.rs index 08f5ea111..f447d60a8 100644 --- a/crates/aionui-app/src/services.rs +++ b/crates/aionui-app/src/services.rs @@ -13,10 +13,11 @@ use aionui_common::OnConversationDelete; use aionui_conversation::{ConversationService, runtime_state::ConversationRuntimeStateService}; use aionui_db::{ Database, IAcpSessionRepository, IAgentMetadataRepository, IConversationRepository, IMcpServerRepository, - IProjectStore, ISkillRepository, IUserRepository, SqliteAcpSessionRepository, SqliteAgentMetadataRepository, - SqliteAssistantDefinitionRepository, SqliteAssistantOverlayRepository, SqliteAssistantPreferenceRepository, - SqliteConversationRepository, SqliteMcpServerRepository, SqliteProjectStore, SqliteProviderRepository, - SqliteSkillRepository, SqliteUserRepository, + IOAuthTokenRepository, IProjectStore, ISkillRepository, IUserRepository, SqliteAcpSessionRepository, + SqliteAgentMetadataRepository, SqliteAssistantDefinitionRepository, SqliteAssistantOverlayRepository, + SqliteAssistantPreferenceRepository, SqliteConversationRepository, SqliteMcpServerRepository, + SqliteOAuthTokenRepository, SqliteProjectStore, SqliteProviderRepository, SqliteSkillRepository, + SqliteUserRepository, }; use aionui_project::ProjectService; use aionui_realtime::{BroadcastEventBus, WebSocketManager}; @@ -178,6 +179,12 @@ impl AppServices { // so the agent gets the operator's tools (ELECTRON-1JG fix). let mcp_server_repo: Arc = Arc::new(SqliteMcpServerRepository::new(database.pool().clone())); + // So HTTP/SSE MCP servers a user has logged into (via /api/mcp/oauth/*) + // get their token attached as an Authorization header when an agent + // actually calls their tools — OAuth login alone doesn't do this; see + // mcp_resolve::inject_oauth_bearer_header. + let oauth_token_repo: Arc = + Arc::new(SqliteOAuthTokenRepository::new(database.pool().clone())); let agent_metadata_repo: Arc = Arc::new(SqliteAgentMetadataRepository::new(database.pool().clone())); @@ -259,6 +266,7 @@ impl AppServices { broadcaster: event_bus.clone(), backend_binary_path: backend_binary_path.clone(), mcp_server_repo: Some(mcp_server_repo), + oauth_token_repo: Some(oauth_token_repo), session_spawner, // agy cannot prompt for tool permission in headless mode, so AionUi // registers itself as its PreToolUse hook; the hook process calls diff --git a/crates/aionui-conversation/src/service.rs b/crates/aionui-conversation/src/service.rs index be5a97a86..402cee5f6 100644 --- a/crates/aionui-conversation/src/service.rs +++ b/crates/aionui-conversation/src/service.rs @@ -5326,7 +5326,7 @@ impl ConversationService { continue; } if row.builtin { - match aionui_ai_agent::mcp_resolve::row_to_session_mcp_server(&row).await { + match aionui_ai_agent::mcp_resolve::row_to_session_mcp_server(&row, user_id, None).await { Ok(server) => session_mcp_servers.push(server), Err(err) => mcp_statuses.push(ConversationMcpStatus { id: row.id, From e9ed98f0ed0cd1a2afd537711a865b04cb131ffe Mon Sep 17 00:00:00 2001 From: Hani Akrim Date: Tue, 18 Aug 2026 16:47:11 +0300 Subject: [PATCH 7/7] fix(mcp): attach stored OAuth tokens to aionrs backend tool-call requests The claude/codex/antigravity backends resolve MCP servers through mcp_resolve::row_to_session_mcp_server, which now attaches a stored OAuth bearer token to HTTP/SSE transports. The aionrs ("Nabd CLI") backend has its own separate MCP-loading path in factory/aionrs.rs::row_to_mcp_server_config that never got the same treatment, so an OAuth-authenticated MCP server (e.g. Higgsfield) would work from claude/codex but still call tools unauthenticated from aionrs, the default chat mode. --- crates/aionui-ai-agent/src/factory/aionrs.rs | 103 +++++++++++++++++-- crates/aionui-ai-agent/src/mcp_resolve.rs | 2 +- 2 files changed, 98 insertions(+), 7 deletions(-) diff --git a/crates/aionui-ai-agent/src/factory/aionrs.rs b/crates/aionui-ai-agent/src/factory/aionrs.rs index 2f648b3ea..d089aa8a1 100644 --- a/crates/aionui-ai-agent/src/factory/aionrs.rs +++ b/crates/aionui-ai-agent/src/factory/aionrs.rs @@ -11,8 +11,8 @@ use aionui_api_types::{ SessionMcpTransport, TEAM_MCP_SERVER_NAME, TeamMcpStdioConfig, }; use aionui_common::ProviderWithModel; -use aionui_db::IMcpServerRepository; use aionui_db::models::McpServerRow; +use aionui_db::{IMcpServerRepository, IOAuthTokenRepository}; use aionui_realtime::EventBroadcaster; use aionui_runtime::ensure_runtime_command_with_reporter; use serde_json::{Map, Value}; @@ -57,6 +57,7 @@ pub(super) async fn build( &ctx.user_id, &ctx.conversation_id, deps.broadcaster.clone(), + deps.oauth_token_repo.as_deref(), ) .await { @@ -514,6 +515,7 @@ async fn load_user_mcp_servers( user_id: &str, conversation_id: &str, broadcaster: Arc, + oauth_token_repo: Option<&dyn IOAuthTokenRepository>, ) -> HashMap { let rows_result = match selected_ids { Some(ids) => repo.list_by_ids_any(user_id, ids).await, @@ -543,7 +545,7 @@ async fn load_user_mcp_servers( continue; } - match row_to_mcp_server_config(&row, user_id, conversation_id, broadcaster.clone()).await { + match row_to_mcp_server_config(&row, user_id, conversation_id, broadcaster.clone(), oauth_token_repo).await { Ok(config) => { servers.insert(row.name.clone(), config); } @@ -567,6 +569,7 @@ async fn row_to_mcp_server_config( user_id: &str, conversation_id: &str, broadcaster: Arc, + oauth_token_repo: Option<&dyn IOAuthTokenRepository>, ) -> Result { let value: serde_json::Value = serde_json::from_str(&row.transport_config).map_err(|e| format!("invalid transport_config JSON: {e}"))?; @@ -610,7 +613,7 @@ async fn row_to_mcp_server_config( .get("url") .and_then(|v| v.as_str()) .ok_or_else(|| "http: missing url".to_owned())?; - let headers = value + let mut headers = value .get("headers") .and_then(|v| v.as_object()) .map(|obj| { @@ -619,6 +622,7 @@ async fn row_to_mcp_server_config( .collect::>() }) .unwrap_or_default(); + crate::mcp_resolve::inject_oauth_bearer_header(&mut headers, user_id, url, oauth_token_repo).await; Ok(McpServerConfig { transport: TransportType::StreamableHttp, @@ -636,7 +640,7 @@ async fn row_to_mcp_server_config( .get("url") .and_then(|v| v.as_str()) .ok_or_else(|| "sse: missing url".to_owned())?; - let headers = value + let mut headers = value .get("headers") .and_then(|v| v.as_object()) .map(|obj| { @@ -645,6 +649,7 @@ async fn row_to_mcp_server_config( .collect::>() }) .unwrap_or_default(); + crate::mcp_resolve::inject_oauth_bearer_header(&mut headers, user_id, url, oauth_token_repo).await; Ok(McpServerConfig { transport: TransportType::Sse, @@ -1045,6 +1050,91 @@ mod tests { Arc::new(BroadcastEventBus::new(16)) } + // Regression coverage: the aionrs ("Nabd CLI") backend loads MCP servers + // through its own row_to_mcp_server_config, separate from + // mcp_resolve::row_to_session_mcp_server used by claude/codex/antigravity. + // A stored OAuth token must reach this path's outbound Authorization + // header too, not just the other one. + + struct MockOAuthRepo { + row: Option, + } + + #[async_trait::async_trait] + impl IOAuthTokenRepository for MockOAuthRepo { + async fn get_by_url( + &self, + _user_id: &str, + _server_url: &str, + ) -> Result, aionui_db::DbError> { + Ok(self.row.clone()) + } + async fn upsert( + &self, + _params: aionui_db::UpsertOAuthTokenParams<'_>, + ) -> Result { + unimplemented!("not needed") + } + async fn delete(&self, _user_id: &str, _server_url: &str) -> Result<(), aionui_db::DbError> { + unimplemented!("not needed") + } + async fn list_authenticated_urls(&self, _user_id: &str) -> Result, aionui_db::DbError> { + unimplemented!("not needed") + } + } + + fn token_row(access_token: &str) -> aionui_db::models::OAuthTokenRow { + aionui_db::models::OAuthTokenRow { + user_id: TEST_USER_ID.to_owned(), + server_url: "http://localhost:54321/mcp".to_owned(), + access_token: access_token.to_owned(), + refresh_token: None, + token_type: "bearer".to_owned(), + expires_at: None, + created_at: 0, + updated_at: 0, + } + } + + #[tokio::test] + async fn row_to_mcp_server_config_attaches_bearer_header_when_token_stored() { + let row = make_row( + "higgsfield", + "http", + r#"{"url":"http://localhost:54321/mcp"}"#, + true, + false, + ); + let oauth_repo = MockOAuthRepo { + row: Some(token_row("tok-abc")), + }; + + let config = row_to_mcp_server_config(&row, TEST_USER_ID, "conv-oauth", test_broadcaster(), Some(&oauth_repo)) + .await + .expect("convert"); + + let headers = config.headers.expect("headers present"); + assert_eq!(headers.get("Authorization"), Some(&"Bearer tok-abc".to_string())); + } + + #[tokio::test] + async fn row_to_mcp_server_config_omits_header_when_no_oauth_repo_given() { + let row = make_row( + "higgsfield", + "http", + r#"{"url":"http://localhost:54321/mcp"}"#, + true, + false, + ); + + let config = row_to_mcp_server_config(&row, TEST_USER_ID, "conv-oauth", test_broadcaster(), None) + .await + .expect("convert"); + + let headers = config.headers.expect("headers present"); + assert!(!headers.contains_key("Authorization")); + } + #[tokio::test] async fn aionrs_loads_mcp_servers_from_frozen_selection_snapshot() { let mut row = make_row( @@ -1064,6 +1154,7 @@ mod tests { TEST_USER_ID, "conv-frozen-mcp", test_broadcaster(), + None, ) .await; @@ -1087,7 +1178,7 @@ mod tests { false, ); - let config = row_to_mcp_server_config(&row, "user-row", "conv-row", test_broadcaster()) + let config = row_to_mcp_server_config(&row, "user-row", "conv-row", test_broadcaster(), None) .await .expect("convert"); let command = config.command.as_deref().expect("resolved command"); @@ -1976,7 +2067,7 @@ mod tests { }; let mut assembled = resolve_mcp_servers(&overrides); for (name, config) in - load_user_mcp_servers(&repo, None, TEST_USER_ID, "conv-assembly", test_broadcaster()).await + load_user_mcp_servers(&repo, None, TEST_USER_ID, "conv-assembly", test_broadcaster(), None).await { assembled.entry(name).or_insert(config); } diff --git a/crates/aionui-ai-agent/src/mcp_resolve.rs b/crates/aionui-ai-agent/src/mcp_resolve.rs index 9dd73bdb3..ba2d082ca 100644 --- a/crates/aionui-ai-agent/src/mcp_resolve.rs +++ b/crates/aionui-ai-agent/src/mcp_resolve.rs @@ -183,7 +183,7 @@ pub async fn row_to_session_mcp_server( /// Best-effort: an expired token or a repo error just means the server /// won't have a token attached and the tool call fails with a normal /// auth-required error downstream, exactly as if OAuth had never run. -async fn inject_oauth_bearer_header( +pub(crate) async fn inject_oauth_bearer_header( headers: &mut std::collections::HashMap, user_id: &str, server_url: &str,