diff --git a/.githooks/pre-push b/.githooks/pre-push new file mode 100755 index 00000000..cd6b5cd4 --- /dev/null +++ b/.githooks/pre-push @@ -0,0 +1,23 @@ +#!/usr/bin/env bash +set -euo pipefail + +# Pre-push hook: run clippy and tests before pushing. +# Install: git config core.hooksPath .githooks + +echo "pre-push: running clippy..." +if ! cargo clippy --all --benches --tests --examples --all-features -- -D warnings; then + echo "" + echo "Push blocked: clippy warnings found." + echo "To bypass: git push --no-verify" + exit 1 +fi + +echo "pre-push: running tests..." +if ! cargo test; then + echo "" + echo "Push blocked: tests failed." + echo "To bypass: git push --no-verify" + exit 1 +fi + +echo "pre-push: all checks passed." diff --git a/src/agent/agent_loop.rs b/src/agent/agent_loop.rs index 97c9b2b9..8fda4143 100644 --- a/src/agent/agent_loop.rs +++ b/src/agent/agent_loop.rs @@ -18,7 +18,7 @@ use crate::agent::self_repair::{DefaultSelfRepair, RepairResult, SelfRepair}; use crate::agent::session_manager::SessionManager; use crate::agent::submission::{Submission, SubmissionParser, SubmissionResult}; use crate::agent::{HeartbeatConfig as AgentHeartbeatConfig, Router, Scheduler}; -use crate::channels::{ChannelManager, IncomingMessage, OutgoingResponse, StatusUpdate}; +use crate::channels::{ChannelManager, IncomingMessage, OutgoingResponse}; use crate::config::{AgentConfig, HeartbeatConfig, RoutineConfig, SkillsConfig}; use crate::context::ContextManager; use crate::db::Database; @@ -936,29 +936,10 @@ impl Agent { SubmissionResult::Ok { message } => Ok(message), SubmissionResult::Error { message } => Ok(Some(format!("Error: {}", message))), SubmissionResult::Interrupted => Ok(Some("Interrupted.".into())), - SubmissionResult::NeedApproval { - request_id, - tool_name, - description, - parameters, - } => { - // Each channel renders the approval prompt via send_status. - // Web gateway shows an inline card, REPL prints a formatted prompt, etc. - let _ = self - .channels - .send_status( - &message.channel, - StatusUpdate::ApprovalNeeded { - request_id: request_id.to_string(), - tool_name, - description, - parameters, - }, - &message.metadata, - ) - .await; - - // Empty string signals the caller to skip respond() (no duplicate text) + SubmissionResult::NeedApproval { .. } => { + // ApprovalNeeded status was already sent by thread_ops.rs before + // returning this result. Empty string signals the caller to skip + // respond() (no duplicate text). Ok(Some(String::new())) } } diff --git a/src/agent/thread_ops.rs b/src/agent/thread_ops.rs index 566f5140..f3673781 100644 --- a/src/agent/thread_ops.rs +++ b/src/agent/thread_ops.rs @@ -486,7 +486,12 @@ impl Agent { .channels .send_status( &message.channel, - StatusUpdate::Status("Awaiting approval".into()), + StatusUpdate::ApprovalNeeded { + request_id: request_id.to_string(), + tool_name: tool_name.clone(), + description: description.clone(), + parameters: parameters.clone(), + }, &message.metadata, ) .await; @@ -1297,7 +1302,12 @@ impl Agent { .channels .send_status( &message.channel, - StatusUpdate::Status("Awaiting approval".into()), + StatusUpdate::ApprovalNeeded { + request_id: request_id.to_string(), + tool_name: tool_name.clone(), + description: description.clone(), + parameters: parameters.clone(), + }, &message.metadata, ) .await; @@ -1368,7 +1378,12 @@ impl Agent { .channels .send_status( &message.channel, - StatusUpdate::Status("Awaiting approval".into()), + StatusUpdate::ApprovalNeeded { + request_id: request_id.to_string(), + tool_name: tool_name.clone(), + description: description.clone(), + parameters: parameters.clone(), + }, &message.metadata, ) .await; diff --git a/src/channels/web/server.rs b/src/channels/web/server.rs index f08c95c2..904971fc 100644 --- a/src/channels/web/server.rs +++ b/src/channels/web/server.rs @@ -581,7 +581,12 @@ async fn oauth_callback_handler( let exchange_proxy_url = std::env::var("IRONCLAW_OAUTH_EXCHANGE_URL").ok(); let result: Result<(), String> = async { - let token_response = if let Some(ref proxy_url) = exchange_proxy_url { + let token_response = if let (Some(proxy_url), None) = (&exchange_proxy_url, &flow.resource) + { + // Use the platform exchange proxy when configured and no resource + // parameter is needed. The proxy holds client_secret server-side so + // the container never sees it. MCP flows (resource.is_some()) bypass + // the proxy because it doesn't forward the RFC 8707 resource param. let gateway_token = flow.gateway_token.as_deref().unwrap_or_default(); oauth_defaults::exchange_via_proxy( proxy_url, @@ -594,7 +599,10 @@ async fn oauth_callback_handler( .await .map_err(|e| e.to_string())? } else { - oauth_defaults::exchange_oauth_code( + // Direct token exchange: uses exchange_oauth_code_with_resource so MCP + // flows can include the RFC 8707 `resource` parameter to scope the + // issued token to the specific MCP server. + oauth_defaults::exchange_oauth_code_with_resource( &flow.token_url, &flow.client_id, flow.client_secret.as_deref(), @@ -602,6 +610,7 @@ async fn oauth_callback_handler( &flow.redirect_uri, flow.code_verifier.as_deref(), &flow.access_token_field, + flow.resource.as_deref(), ) .await .map_err(|e| e.to_string())? @@ -628,6 +637,19 @@ async fn oauth_callback_handler( .await .map_err(|e| e.to_string())?; + // For MCP OAuth flows (identified by resource field), persist the + // client_id so token refresh works without re-authentication. + // The CLI flow stores this in authorize_mcp_server(); the gateway + // callback must do the same. + if let Some(ref client_id_secret) = flow.client_id_secret_name { + let params = crate::secrets::CreateSecretParams::new(client_id_secret, &flow.client_id) + .with_provider(flow.provider.as_ref().cloned().unwrap_or_default()); + flow.secrets + .create(&flow.user_id, params) + .await + .map_err(|e| e.to_string())?; + } + Ok(()) } .await; @@ -659,12 +681,35 @@ async fn oauth_callback_handler( } } + // After successful OAuth, auto-activate the extension so it moves + // from "Installed (Authenticate)" → "Active" without a second click. + // OAuth success is independent of activation — tokens are already stored. + // Report auth as successful and attempt activation as a bonus step. + let final_message = if success { + match ext_mgr.activate(&flow.extension_name).await { + Ok(result) => result.message, + Err(e) => { + tracing::warn!( + extension = %flow.extension_name, + error = %e, + "Auto-activation after OAuth failed" + ); + format!( + "{} authenticated successfully. Activation failed: {}. Try activating manually.", + flow.display_name, e + ) + } + } + } else { + message + }; + // Broadcast SSE event to notify the web UI if let Some(ref sender) = flow.sse_sender { let _ = sender.send(SseEvent::AuthCompleted { extension_name: flow.extension_name, success, - message, + message: final_message.clone(), }); } @@ -2966,6 +3011,8 @@ mod tests { secrets, sse_sender: None, gateway_token: None, + resource: None, + client_id_secret_name: None, created_at: std::time::Instant::now() .checked_sub(std::time::Duration::from_secs(600)) .expect("System uptime is too low to run expired flow test"), @@ -3075,6 +3122,8 @@ mod tests { secrets, sse_sender: None, gateway_token: None, + resource: None, + client_id_secret_name: None, // Expired — handler will reject after lookup (no network I/O) created_at: std::time::Instant::now() .checked_sub(std::time::Duration::from_secs(600)) diff --git a/src/cli/oauth_defaults.rs b/src/cli/oauth_defaults.rs index 2da14f0a..a625f718 100644 --- a/src/cli/oauth_defaults.rs +++ b/src/cli/oauth_defaults.rs @@ -172,6 +172,35 @@ pub async fn exchange_oauth_code( redirect_uri: &str, code_verifier: Option<&str>, access_token_field: &str, +) -> Result { + // Delegates to exchange_oauth_code_with_resource with resource=None. + // Non-MCP OAuth flows don't need the RFC 8707 resource parameter. + exchange_oauth_code_with_resource( + token_url, + client_id, + client_secret, + code, + redirect_uri, + code_verifier, + access_token_field, + None, + ) + .await +} + +/// Exchange an OAuth authorization code for tokens, with optional RFC 8707 `resource` parameter. +/// +/// The `resource` parameter scopes the issued token to a specific server (used by MCP OAuth). +#[allow(clippy::too_many_arguments)] +pub async fn exchange_oauth_code_with_resource( + token_url: &str, + client_id: &str, + client_secret: Option<&str>, + code: &str, + redirect_uri: &str, + code_verifier: Option<&str>, + access_token_field: &str, + resource: Option<&str>, ) -> Result { let client = reqwest::Client::new(); let mut token_params = vec![ @@ -184,6 +213,12 @@ pub async fn exchange_oauth_code( token_params.push(("code_verifier", verifier.to_string())); } + // RFC 8707: include the `resource` parameter so the authorization server + // scopes the issued token to the specific MCP server (protected resource). + if let Some(resource) = resource { + token_params.push(("resource", resource.to_string())); + } + let mut request = client.post(token_url); if let Some(secret) = client_secret { @@ -388,6 +423,12 @@ pub struct PendingOAuthFlow { pub sse_sender: Option>, /// Gateway auth token for authenticating with the platform token exchange proxy. pub gateway_token: Option, + /// RFC 8707 resource parameter (MCP OAuth only). + /// Sent during token exchange to scope the token to a specific MCP server. + pub resource: Option, + /// Secret name for persisting the client ID (MCP OAuth only). + /// Needed so token refresh can find the client_id after the session ends. + pub client_id_secret_name: Option, /// When this flow was created (for expiry). pub created_at: std::time::Instant, } @@ -975,4 +1016,42 @@ mod tests { assert_eq!(strip_instance_prefix("abc123"), "abc123"); assert_eq!(strip_instance_prefix(""), ""); } + + /// Verify that `build_oauth_url` includes the RFC 8707 `resource` parameter + /// when passed through `extra_params`, which is how MCP OAuth gateway mode + /// scopes tokens to a specific MCP server. + #[test] + fn test_build_oauth_url_includes_resource_via_extra_params() { + use std::collections::HashMap; + + use crate::cli::oauth_defaults::build_oauth_url; + + let mut extra = HashMap::new(); + extra.insert( + "resource".to_string(), + "https://mcp.example.com".to_string(), + ); + + let result = build_oauth_url( + "https://auth.example.com/authorize", + "client-123", + "https://gateway.example.com/oauth/callback", + &["read".to_string()], + true, + &extra, + ); + + // The resource parameter should be URL-encoded in the auth URL + assert!( + result + .url + .contains("resource=https%3A%2F%2Fmcp.example.com"), + "Expected resource param in URL: {}", + result.url + ); + // State and PKCE should be present + assert!(result.url.contains("state=")); + assert!(result.url.contains("code_challenge=")); + assert!(result.code_verifier.is_some()); + } } diff --git a/src/extensions/manager.rs b/src/extensions/manager.rs index 3f02dd5e..2a6cc6d1 100644 --- a/src/extensions/manager.rs +++ b/src/extensions/manager.rs @@ -27,7 +27,7 @@ use crate::secrets::{CreateSecretParams, SecretsStore}; use crate::tools::ToolRegistry; use crate::tools::mcp::McpClient; use crate::tools::mcp::auth::{ - PkceChallenge, authorize_mcp_server, build_authorization_url, discover_full_oauth_metadata, + authorize_mcp_server, canonical_resource_uri, discover_full_oauth_metadata, find_available_port, is_authenticated, register_client, }; use crate::tools::mcp::config::McpServerConfig; @@ -108,6 +108,13 @@ pub struct ExtensionManager { /// Relay config captured at startup. Used by `auth_channel_relay` and /// `activate_channel_relay` instead of re-reading env vars. relay_config: Option, + /// When `true`, OAuth flows always return an auth URL to the caller + /// instead of opening a browser on the server via `open::that()`. + /// Set by the web gateway at startup via `enable_gateway_mode()`. + gateway_mode: std::sync::atomic::AtomicBool, + /// The gateway's own base URL for building OAuth redirect URIs. + /// Set by the web gateway at startup via `enable_gateway_mode()`. + gateway_base_url: RwLock>, } /// Sanitize a URL for logging by removing query parameters and credentials. @@ -181,9 +188,75 @@ impl ExtensionManager { pending_oauth_flows: crate::cli::oauth_defaults::new_pending_oauth_registry(), gateway_token: std::env::var("GATEWAY_AUTH_TOKEN").ok(), relay_config: crate::config::RelayConfig::from_env(), + gateway_mode: std::sync::atomic::AtomicBool::new(false), + gateway_base_url: RwLock::new(None), } } + /// Enable gateway mode so OAuth flows return auth URLs to the frontend + /// instead of calling `open::that()` on the server. + /// + /// `base_url` is the gateway's own public URL (e.g. `https://my-gateway.example.com`), + /// used to build OAuth redirect URIs when `IRONCLAW_OAUTH_CALLBACK_URL` is not set. + pub async fn enable_gateway_mode(&self, base_url: String) { + self.gateway_mode + .store(true, std::sync::atomic::Ordering::Release); + *self.gateway_base_url.write().await = Some(base_url); + } + + /// Returns `true` if OAuth should use gateway mode (return auth URL to + /// frontend) rather than CLI mode (open browser on server via `open::that`). + /// + /// Gateway mode is active when any of: + /// - `enable_gateway_mode()` was called (web gateway is running), OR + /// - `IRONCLAW_OAUTH_CALLBACK_URL` is set to a non-loopback URL, OR + /// - `self.tunnel_url` is set to a non-loopback URL + pub fn should_use_gateway_mode(&self) -> bool { + if self.gateway_mode.load(std::sync::atomic::Ordering::Acquire) { + return true; + } + if crate::cli::oauth_defaults::use_gateway_callback() { + return true; + } + self.tunnel_url + .as_ref() + .filter(|u| !u.is_empty()) + .and_then(|raw| url::Url::parse(raw).ok()) + .and_then(|u| u.host_str().map(String::from)) + .map(|host| !crate::cli::oauth_defaults::is_loopback_host(&host)) + .unwrap_or(false) + } + + /// Returns the OAuth redirect URI for gateway mode, or `None` for local mode. + /// + /// Priority: + /// 1. `IRONCLAW_OAUTH_CALLBACK_URL` env var (via `callback_url()`) + /// 2. `gateway_base_url` (set by `enable_gateway_mode()`) + /// 3. `tunnel_url` (from config) + /// 4. `None` (local/CLI mode) + async fn gateway_callback_redirect_uri(&self) -> Option { + use crate::cli::oauth_defaults; + if oauth_defaults::use_gateway_callback() { + return Some(format!("{}/oauth/callback", oauth_defaults::callback_url())); + } + // Use gateway_base_url from enable_gateway_mode() + if let Some(ref base) = *self.gateway_base_url.read().await { + let base = base.trim_end_matches('/'); + return Some(format!("{}/oauth/callback", base)); + } + // Fall back to tunnel_url + self.tunnel_url + .as_ref() + .filter(|u| !u.is_empty()) + .and_then(|raw| url::Url::parse(raw).ok()) + .and_then(|u| u.host_str().map(String::from)) + .filter(|host| !oauth_defaults::is_loopback_host(host)) + .map(|_| { + let base = self.tunnel_url.as_ref().unwrap().trim_end_matches('/'); + format!("{}/oauth/callback", base) + }) + } + /// Get the relay config stored at startup. fn relay_config(&self) -> Result<&crate::config::RelayConfig, ExtensionError> { self.relay_config.as_ref().ok_or_else(|| { @@ -193,6 +266,12 @@ impl ExtensionManager { }) } + /// Inject a registry entry for testing. The entry is added to the discovery + /// cache so it appears in search results alongside built-in entries. + pub async fn inject_registry_entry(&self, entry: crate::extensions::RegistryEntry) { + self.registry.cache_discovered(vec![entry]).await; + } + /// Configure the channel runtime infrastructure for hot-activating WASM channels. /// /// Call after construction (and after wrapping in `Arc`) once the channel @@ -1684,29 +1763,46 @@ impl ExtensionManager { return Ok(AuthResult::authenticated(name, ExtensionKind::McpServer)); } - // Run the full OAuth flow (opens browser, waits for callback) + // In gateway mode, build an auth URL and return it for the frontend to + // open in the same browser. The gateway's /oauth/callback handler will + // complete the token exchange. + if self.should_use_gateway_mode() { + return match self.auth_mcp_build_url(name, &server).await { + Ok(result) => Ok(result), + Err(ExtensionError::AuthNotSupported(_)) => Ok(AuthResult::awaiting_token( + name, + ExtensionKind::McpServer, + format!( + "Server '{}' does not support OAuth. \ + Please provide an API token/key for this server.", + name + ), + None, + )), + Err(e) => Err(e), + }; + } + + // CLI/local mode: run the full blocking OAuth flow (opens browser, waits for callback) match authorize_mcp_server(&server, &self.secrets, &self.user_id).await { Ok(_token) => { tracing::info!("MCP server '{}' authenticated via OAuth", name); Ok(AuthResult::authenticated(name, ExtensionKind::McpServer)) } Err(crate::tools::mcp::auth::AuthError::NotSupported) => { - // Server doesn't support OAuth, try building a URL first + // Server doesn't support OAuth, try building a URL match self.auth_mcp_build_url(name, &server).await { Ok(result) => Ok(result), - Err(_) => { - // No OAuth, no DCR: fall back to manual token entry - Ok(AuthResult::awaiting_token( - name, - ExtensionKind::McpServer, - format!( - "Server '{}' does not support OAuth. \ - Please provide an API token/key for this server.", - name - ), - None, - )) - } + Err(_) => Ok(AuthResult::awaiting_token( + name, + ExtensionKind::McpServer, + format!( + "Server '{}' does not support OAuth. \ + Please provide an API token/key for this server.", + name + ), + None, + )), } } Err(e) => { @@ -1725,8 +1821,12 @@ impl ExtensionManager { } } - /// Build an auth URL for cases where non-interactive auth is needed - /// (e.g., running via Telegram where we can't open a browser). + /// Build an auth URL for MCP OAuth. + /// + /// In gateway mode, stores a `PendingOAuthFlow` so the web gateway's + /// `/oauth/callback` handler can complete the token exchange — the auth + /// URL is sent to the frontend which opens it in the same browser. + /// In local/CLI mode, builds the URL for the user to open manually. async fn auth_mcp_build_url( &self, name: &str, @@ -1735,60 +1835,153 @@ impl ExtensionManager { // Try to discover OAuth metadata and build a URL the user can open manually let metadata = discover_full_oauth_metadata(&server.url) .await - .map_err(|e| ExtensionError::AuthFailed(e.to_string()))?; + .map_err(|e| match e { + crate::tools::mcp::auth::AuthError::NotSupported => { + ExtensionError::AuthNotSupported(e.to_string()) + } + _ => ExtensionError::AuthFailed(e.to_string()), + })?; + + use crate::cli::oauth_defaults; + + let is_gateway = self.should_use_gateway_mode(); + + // Build redirect URI: gateway uses the public callback URL, + // local mode binds a random port. + let redirect_uri = if let Some(uri) = self.gateway_callback_redirect_uri().await { + uri + } else { + let port = find_available_port() + .await + .map_err(|e| ExtensionError::AuthFailed(e.to_string()))?; + format!("http://localhost:{}/callback", port.1) + }; // Try DCR if no client_id configured - let (client_id, redirect_uri) = if let Some(ref oauth) = server.oauth { - let port = find_available_port() - .await - .map_err(|e| ExtensionError::AuthFailed(e.to_string()))?; - let redirect = format!("http://localhost:{}/callback", port.1); - (oauth.client_id.clone(), redirect) + let (client_id, client_secret) = if let Some(ref oauth) = server.oauth { + (oauth.client_id.clone(), None) } else if let Some(ref reg_endpoint) = metadata.registration_endpoint { - let port = find_available_port() - .await - .map_err(|e| ExtensionError::AuthFailed(e.to_string()))?; - let redirect = format!("http://localhost:{}/callback", port.1); - - let registration = register_client(reg_endpoint, &redirect) + let registration = register_client(reg_endpoint, &redirect_uri) .await .map_err(|e| ExtensionError::AuthFailed(e.to_string()))?; - (registration.client_id, redirect) + (registration.client_id, None) } else { - return Err(ExtensionError::AuthFailed( + return Err(ExtensionError::AuthNotSupported( "Server doesn't support OAuth or Dynamic Client Registration".to_string(), )); }; - let pkce = PkceChallenge::generate(); - let auth_url = build_authorization_url( + // RFC 8707: resource parameter to scope the token to this MCP server + let resource = canonical_resource_uri(&server.url); + + // Build authorization URL with CSRF state using the shared oauth_defaults + // builder, which generates PKCE + state for us. + let mut extra_params = server + .oauth + .as_ref() + .map(|o| o.extra_params.clone()) + .unwrap_or_default(); + extra_params.insert("resource".to_string(), resource.clone()); + + let scopes = server + .oauth + .as_ref() + .map(|o| o.scopes.clone()) + .unwrap_or_else(|| metadata.scopes_supported.clone()); + + let oauth_result = oauth_defaults::build_oauth_url( &metadata.authorization_endpoint, &client_id, &redirect_uri, - &metadata.scopes_supported, - Some(&pkce), - &std::collections::HashMap::new(), - None, + &scopes, + true, // Always use PKCE for MCP + &extra_params, ); + let expected_state = oauth_result.state; + let code_verifier = oauth_result.code_verifier; - // Store pending auth for later callback handling - self.pending_auth.write().await.insert( - name.to_string(), - PendingAuth { - _name: name.to_string(), - _kind: ExtensionKind::McpServer, + if is_gateway { + // Gateway mode: store pending flow for the /oauth/callback handler. + oauth_defaults::sweep_expired_flows(&self.pending_oauth_flows).await; + + // Platform routing: prepend instance name to state + let platform_state = oauth_defaults::build_platform_state(&expected_state); + let auth_url = if platform_state != expected_state { + oauth_result.url.replace( + &format!("state={}", urlencoding::encode(&expected_state)), + &format!("state={}", urlencoding::encode(&platform_state)), + ) + } else { + oauth_result.url + }; + + let flow = oauth_defaults::PendingOAuthFlow { + extension_name: name.to_string(), + display_name: server.name.clone(), + token_url: metadata.token_endpoint, + client_id, + client_secret, + redirect_uri, + code_verifier, + access_token_field: "access_token".to_string(), + secret_name: server.token_secret_name(), + provider: Some(format!("mcp:{}", name)), + validation_endpoint: None, + scopes, + user_id: self.user_id.clone(), + secrets: Arc::clone(&self.secrets), + sse_sender: self.sse_sender.read().await.clone(), + gateway_token: self.gateway_token.clone(), + resource: Some(resource), + client_id_secret_name: if server.oauth.is_none() { + Some(server.client_id_secret_name()) + } else { + None + }, created_at: std::time::Instant::now(), - task_handle: None, - }, - ); + }; - Ok(AuthResult::awaiting_authorization( - name, - ExtensionKind::McpServer, - auth_url, - "local".to_string(), - )) + self.pending_oauth_flows + .write() + .await + .insert(expected_state, flow); + + self.pending_auth.write().await.insert( + name.to_string(), + PendingAuth { + _name: name.to_string(), + _kind: ExtensionKind::McpServer, + created_at: std::time::Instant::now(), + task_handle: None, + }, + ); + + Ok(AuthResult::awaiting_authorization( + name, + ExtensionKind::McpServer, + auth_url, + "gateway".to_string(), + )) + } else { + // Local mode: return URL for manual opening + self.pending_auth.write().await.insert( + name.to_string(), + PendingAuth { + _name: name.to_string(), + _kind: ExtensionKind::McpServer, + created_at: std::time::Instant::now(), + task_handle: None, + }, + ); + + Ok(AuthResult::awaiting_authorization( + name, + ExtensionKind::McpServer, + oauth_result.url, + "local".to_string(), + )) + } } async fn auth_wasm_tool(&self, name: &str) -> Result { @@ -2203,7 +2396,10 @@ impl ExtensionManager { flows.retain(|_, flow| flow.extension_name != name); } - let redirect_uri = format!("{}/callback", oauth_defaults::callback_url()); + let redirect_uri = self + .gateway_callback_redirect_uri() + .await + .unwrap_or_else(|| format!("{}/callback", oauth_defaults::callback_url())); // Merge scopes from all tools sharing this provider let merged_scopes = self @@ -2228,7 +2424,7 @@ impl ExtensionManager { .clone() .unwrap_or_else(|| name.to_string()); - if oauth_defaults::use_gateway_callback() { + if self.should_use_gateway_mode() { // Gateway mode: store pending flow state for the web gateway's // `/oauth/callback` handler to complete the exchange. No TCP listener // needed — the OAuth provider redirects to the gateway URL. @@ -2264,6 +2460,8 @@ impl ExtensionManager { secrets: Arc::clone(&self.secrets), sse_sender: self.sse_sender.read().await.clone(), gateway_token: self.gateway_token.clone(), + resource: None, + client_id_secret_name: None, created_at: std::time::Instant::now(), }; @@ -2605,11 +2803,17 @@ impl ExtensionManager { .await .map_err(|e| ExtensionError::ActivationFailed(e.to_string()))?; - // Try to list and create tools - let mcp_tools = client - .list_tools() - .await - .map_err(|e| ExtensionError::ActivationFailed(e.to_string()))?; + // Try to list and create tools. + // A 401/auth error means the server requires OAuth — surface as + // AuthRequired so the activate handler triggers the OAuth flow. + let mcp_tools = client.list_tools().await.map_err(|e| { + let msg = e.to_string(); + if msg.contains("requires authentication") || msg.contains("401") { + ExtensionError::AuthRequired + } else { + ExtensionError::ActivationFailed(msg) + } + })?; let tool_impls = client .create_tools() @@ -4766,6 +4970,190 @@ mod tests { assert!(result.contains("/v1/users/123/profile")); } + // ---- gateway mode detection tests ---- + // Regression tests for a bug where MCP OAuth called `open::that()` on the + // server machine instead of returning an auth URL to the gateway frontend. + // The root cause was that `should_use_gateway_mode()` only checked the + // `IRONCLAW_OAUTH_CALLBACK_URL` env var, ignoring `self.tunnel_url`. + + /// Serializes env-mutating tests to prevent parallel races. + static GATEWAY_ENV_MUTEX: std::sync::Mutex<()> = std::sync::Mutex::new(()); + + /// Build a minimal ExtensionManager with a custom tunnel_url. + fn make_manager_with_tunnel(tunnel_url: Option) -> ExtensionManager { + use crate::secrets::{InMemorySecretsStore, SecretsCrypto}; + use crate::tools::mcp::process::McpProcessManager; + use crate::tools::mcp::session::McpSessionManager; + + let key = secrecy::SecretString::from(crate::secrets::keychain::generate_master_key_hex()); + let crypto = Arc::new(SecretsCrypto::new(key).expect("crypto")); + let secrets: Arc = + Arc::new(InMemorySecretsStore::new(crypto)); + let tools = Arc::new(crate::tools::ToolRegistry::new()); + let mcp = Arc::new(McpSessionManager::new()); + let dir = std::env::temp_dir().join("ironclaw-test-gateway-mode"); + + ExtensionManager::new( + mcp, + Arc::new(McpProcessManager::new()), + secrets, + tools, + None, + None, + dir.clone(), + dir, + tunnel_url, + "test".to_string(), + None, + vec![], + ) + } + + #[test] + fn should_use_gateway_mode_true_for_tunnel_url() { + let _guard = GATEWAY_ENV_MUTEX.lock().expect("env mutex poisoned"); + let original = std::env::var("IRONCLAW_OAUTH_CALLBACK_URL").ok(); + // SAFETY: Under GATEWAY_ENV_MUTEX, no concurrent env access. + unsafe { + std::env::remove_var("IRONCLAW_OAUTH_CALLBACK_URL"); + } + + let mgr = make_manager_with_tunnel(Some("https://my-gateway.example.com".into())); + assert!( + mgr.should_use_gateway_mode(), + "should detect gateway mode from tunnel_url" + ); + + unsafe { + if let Some(val) = original { + std::env::set_var("IRONCLAW_OAUTH_CALLBACK_URL", val); + } + } + } + + #[test] + fn should_use_gateway_mode_false_without_tunnel() { + let _guard = GATEWAY_ENV_MUTEX.lock().expect("env mutex poisoned"); + let original = std::env::var("IRONCLAW_OAUTH_CALLBACK_URL").ok(); + unsafe { + std::env::remove_var("IRONCLAW_OAUTH_CALLBACK_URL"); + } + + let mgr = make_manager_with_tunnel(None); + assert!( + !mgr.should_use_gateway_mode(), + "should not detect gateway mode without tunnel_url or env var" + ); + + unsafe { + if let Some(val) = original { + std::env::set_var("IRONCLAW_OAUTH_CALLBACK_URL", val); + } + } + } + + #[test] + fn should_use_gateway_mode_false_for_loopback_tunnel() { + let _guard = GATEWAY_ENV_MUTEX.lock().expect("env mutex poisoned"); + let original = std::env::var("IRONCLAW_OAUTH_CALLBACK_URL").ok(); + unsafe { + std::env::remove_var("IRONCLAW_OAUTH_CALLBACK_URL"); + } + + let mgr = make_manager_with_tunnel(Some("http://127.0.0.1:3001".into())); + assert!( + !mgr.should_use_gateway_mode(), + "should not detect gateway mode for loopback tunnel_url" + ); + + unsafe { + if let Some(val) = original { + std::env::set_var("IRONCLAW_OAUTH_CALLBACK_URL", val); + } + } + } + + /// Helper to run an async test body while holding the env mutex. + /// Clears `IRONCLAW_OAUTH_CALLBACK_URL` for the duration, restoring on drop. + struct EnvGuard { + original: Option, + _mutex: std::sync::MutexGuard<'static, ()>, + } + + impl EnvGuard { + fn new() -> Self { + let guard = GATEWAY_ENV_MUTEX.lock().expect("env mutex poisoned"); + let original = std::env::var("IRONCLAW_OAUTH_CALLBACK_URL").ok(); + // SAFETY: Under GATEWAY_ENV_MUTEX, no concurrent env access. + unsafe { + std::env::remove_var("IRONCLAW_OAUTH_CALLBACK_URL"); + } + Self { + original, + _mutex: guard, + } + } + } + + impl Drop for EnvGuard { + fn drop(&mut self) { + // SAFETY: Under GATEWAY_ENV_MUTEX (still held by _mutex), no concurrent env access. + unsafe { + if let Some(ref val) = self.original { + std::env::set_var("IRONCLAW_OAUTH_CALLBACK_URL", val); + } else { + std::env::remove_var("IRONCLAW_OAUTH_CALLBACK_URL"); + } + } + } + } + + #[tokio::test] + async fn gateway_callback_redirect_uri_from_tunnel_url() { + let _env = EnvGuard::new(); + + let mgr = make_manager_with_tunnel(Some("https://my-gateway.example.com".into())); + assert_eq!( + mgr.gateway_callback_redirect_uri().await, + Some("https://my-gateway.example.com/oauth/callback".to_string()), + ); + } + + #[tokio::test] + async fn gateway_callback_redirect_uri_none_without_tunnel() { + let _env = EnvGuard::new(); + + let mgr = make_manager_with_tunnel(None); + assert_eq!(mgr.gateway_callback_redirect_uri().await, None); + } + + #[tokio::test] + async fn gateway_callback_redirect_uri_trims_trailing_slash() { + let _env = EnvGuard::new(); + + let mgr = make_manager_with_tunnel(Some("https://my-gateway.example.com/".into())); + assert_eq!( + mgr.gateway_callback_redirect_uri().await, + Some("https://my-gateway.example.com/oauth/callback".to_string()), + ); + } + + #[tokio::test] + async fn gateway_mode_enabled_explicitly() { + let _env = EnvGuard::new(); + + let mgr = make_manager_with_tunnel(None); + assert!(!mgr.should_use_gateway_mode()); + + mgr.enable_gateway_mode("https://my-gateway.example.com".into()) + .await; + assert!(mgr.should_use_gateway_mode()); + assert_eq!( + mgr.gateway_callback_redirect_uri().await, + Some("https://my-gateway.example.com/oauth/callback".to_string()), + ); + } + // ── Regression tests for PR #677 (unify-extension-lifecycle) ───────── #[tokio::test] diff --git a/src/extensions/mod.rs b/src/extensions/mod.rs index ce91a998..428d9b42 100644 --- a/src/extensions/mod.rs +++ b/src/extensions/mod.rs @@ -517,6 +517,9 @@ pub enum ExtensionError { #[error("Authentication failed: {0}")] AuthFailed(String), + #[error("Server does not support OAuth: {0}")] + AuthNotSupported(String), + #[error("Activation failed: {0}")] ActivationFailed(String), diff --git a/src/main.rs b/src/main.rs index 6444242c..89fa2068 100644 --- a/src/main.rs +++ b/src/main.rs @@ -472,6 +472,14 @@ async fn async_main() -> anyhow::Result<()> { gw = gw.with_log_level_handle(Arc::clone(&log_level_handle)); gw = gw.with_tool_registry(Arc::clone(&components.tools)); if let Some(ref ext_mgr) = components.extension_manager { + // Enable gateway mode so MCP OAuth returns auth URLs to the frontend + // instead of calling open::that() on the server. + let gw_base = config + .tunnel + .public_url + .clone() + .unwrap_or_else(|| format!("http://{}:{}", gw_config.host, gw_config.port)); + ext_mgr.enable_gateway_mode(gw_base).await; gw = gw.with_extension_manager(Arc::clone(ext_mgr)); } if !components.catalog_entries.is_empty() { diff --git a/src/tools/builtin/extension_tools.rs b/src/tools/builtin/extension_tools.rs index fd8a9d3b..cb0f71dd 100644 --- a/src/tools/builtin/extension_tools.rs +++ b/src/tools/builtin/extension_tools.rs @@ -256,7 +256,13 @@ impl Tool for ToolAuthTool { } fn requires_approval(&self, _params: &serde_json::Value) -> ApprovalRequirement { - ApprovalRequirement::UnlessAutoApproved + // In gateway mode, tool_auth only returns an auth URL for the frontend + // to open — no browser is launched server-side, so no approval needed. + if self.manager.should_use_gateway_mode() { + ApprovalRequirement::Never + } else { + ApprovalRequirement::UnlessAutoApproved + } } } @@ -733,6 +739,22 @@ mod tests { } } + #[tokio::test] + async fn tool_auth_no_approval_in_gateway_mode() { + let manager = test_manager_stub(); + manager + .enable_gateway_mode("http://localhost:3000".to_string()) + .await; + let tool = ToolAuthTool { + manager: manager.clone(), + }; + assert_eq!( + tool.requires_approval(&serde_json::json!({})), + ApprovalRequirement::Never, + "tool_auth should not require approval in gateway mode" + ); + } + #[test] fn test_tool_upgrade_schema() { use crate::tools::tool::ApprovalRequirement; diff --git a/tests/e2e_advanced_traces.rs b/tests/e2e_advanced_traces.rs index 5a307cd8..7b114d28 100644 --- a/tests/e2e_advanced_traces.rs +++ b/tests/e2e_advanced_traces.rs @@ -403,4 +403,136 @@ mod advanced { rig.verify_trace_expects(&trace, &responses); rig.shutdown(); } + + // ----------------------------------------------------------------------- + // 8. MCP extension lifecycle (search → install → activate → use) + // + // Exercises the MCP extension flow with a mock MCP server: + // Turn 1: tool_search → tool_install → text + // (inject token + activate between turns) + // Turn 2: mock-notion_notion-search → mock-notion_notion-fetch → text + // ----------------------------------------------------------------------- + + #[tokio::test] + async fn mcp_extension_lifecycle() { + use crate::support::mock_mcp_server::{MockToolResponse, start_mock_mcp_server}; + use ironclaw::extensions::{AuthHint, ExtensionKind, ExtensionSource, RegistryEntry}; + + // 1. Start mock MCP server with pre-configured tool responses. + let mock_server = start_mock_mcp_server(vec![ + MockToolResponse { + name: "notion-search".into(), + content: serde_json::json!({ + "results": [ + {"id": "page-001", "title": "Project Alpha", "type": "page"}, + {"id": "page-002", "title": "Sprint Planning", "type": "page"} + ] + }), + }, + MockToolResponse { + name: "notion-fetch".into(), + content: serde_json::json!({ + "id": "page-001", + "title": "Project Alpha", + "content": "Status: In Progress\n- Sprint planning on March 15\n- API redesign review pending" + }), + }, + ]) + .await; + + // 2. Load trace fixture. + let trace = + LlmTrace::from_file(format!("{FIXTURES}/mcp_extension_lifecycle.json")).unwrap(); + + // 3. Build rig with auto-approve (so tool_install doesn't block). + let rig = TestRigBuilder::new() + .with_trace(trace.clone()) + .with_auto_approve_tools(true) + .with_max_tool_iterations(15) + .build() + .await; + + // 4. Inject mock-notion registry entry pointing to the mock server. + let ext_mgr = rig + .extension_manager() + .expect("test rig must expose extension manager"); + ext_mgr + .inject_registry_entry(RegistryEntry { + name: "mock-notion".to_string(), + display_name: "Mock Notion".to_string(), + kind: ExtensionKind::McpServer, + description: "Test MCP server for E2E lifecycle test".to_string(), + keywords: vec!["mock-notion".into(), "notion".into()], + source: ExtensionSource::McpUrl { + url: mock_server.mcp_url(), + }, + fallback_source: None, + auth_hint: AuthHint::Dcr, + version: None, + }) + .await; + + // 5. Turn 1: "setup mock-notion" → search → install → text. + rig.send_message("setup mock-notion").await; + let r1 = rig.wait_for_responses(1, TIMEOUT).await; + assert!(!r1.is_empty(), "Turn 1: no response"); + + // 6. Simulate OAuth completion: inject token + activate. + // This mirrors what the gateway's oauth_callback_handler does after + // the user completes the OAuth flow in their browser. + let secret_name = "mcp_mock-notion_access_token"; + ext_mgr + .secrets() + .create( + "default", + ironclaw::secrets::CreateSecretParams::new(secret_name, "mock-access-token") + .with_provider("mcp:mock-notion".to_string()), + ) + .await + .expect("failed to inject test token"); + + let activate_result = ext_mgr.activate("mock-notion").await; + assert!( + activate_result.is_ok(), + "activation failed: {:?}", + activate_result.err() + ); + + // 7. Turn 2: "check what's in my notion" → notion-search → notion-fetch → text. + // Wait for r1.len() + 1 to ensure we observe at least one new turn-2 response. + let turn1_count = r1.len(); + rig.send_message("it's done, check what's in my notion") + .await; + let r2 = rig.wait_for_responses(turn1_count + 1, TIMEOUT).await; + assert!( + r2.len() > turn1_count, + "Turn 2: expected new responses beyond turn 1's {turn1_count}, got {}", + r2.len() + ); + + // 8. Verify tool calls across both turns. + let started = rig.tool_calls_started(); + assert!( + started.iter().any(|s| s == "tool_search"), + "tool_search not called: {started:?}" + ); + assert!( + started.iter().any(|s| s == "tool_install"), + "tool_install not called: {started:?}" + ); + + // Verify MCP tools were called in turn 2. + assert!( + started.iter().any(|s| s.starts_with("mock-notion_")), + "No mock-notion MCP tools called: {started:?}" + ); + + // Verify all tools that completed did so successfully. + let completed = rig.tool_calls_completed(); + let failed: Vec<_> = completed.iter().filter(|(_, success)| !success).collect(); + assert!(failed.is_empty(), "Tools failed: {failed:?}"); + + mock_server.shutdown().await; + rig.shutdown(); + } } diff --git a/tests/fixtures/llm_traces/advanced/mcp_extension_lifecycle.json b/tests/fixtures/llm_traces/advanced/mcp_extension_lifecycle.json new file mode 100644 index 00000000..59655a65 --- /dev/null +++ b/tests/fixtures/llm_traces/advanced/mcp_extension_lifecycle.json @@ -0,0 +1,98 @@ +{ + "model_name": "advanced-mcp-extension-lifecycle", + "expects": { + "tools_used": ["tool_search", "tool_install"], + "tools_order": ["tool_search", "tool_install"], + "all_tools_succeeded": true, + "min_responses": 2 + }, + "turns": [ + { + "user_input": "setup mock-notion", + "steps": [ + { + "request_hint": { "last_user_message_contains": "setup mock-notion" }, + "response": { + "type": "tool_calls", + "tool_calls": [ + { + "id": "call_search_1", + "name": "tool_search", + "arguments": { "query": "mock-notion" } + } + ], + "input_tokens": 500, + "output_tokens": 30 + } + }, + { + "request_hint": { "last_user_message_contains": "setup mock-notion", "min_message_count": 4 }, + "response": { + "type": "tool_calls", + "tool_calls": [ + { + "id": "call_install_1", + "name": "tool_install", + "arguments": { "name": "mock-notion" } + } + ], + "input_tokens": 600, + "output_tokens": 30 + } + }, + { + "request_hint": { "last_user_message_contains": "setup mock-notion", "min_message_count": 6 }, + "response": { + "type": "text", + "content": "I've installed Mock Notion. Please authenticate to connect your account — once done, tell me and I'll load the MCP tools.", + "input_tokens": 700, + "output_tokens": 35 + } + } + ] + }, + { + "user_input": "it's done, check what's in my notion", + "steps": [ + { + "request_hint": { "last_user_message_contains": "notion" }, + "response": { + "type": "tool_calls", + "tool_calls": [ + { + "id": "call_ns_1", + "name": "mock-notion_notion-search", + "arguments": { "query": "recent notes" } + } + ], + "input_tokens": 900, + "output_tokens": 30 + } + }, + { + "request_hint": { "min_message_count": 4 }, + "response": { + "type": "tool_calls", + "tool_calls": [ + { + "id": "call_nf_1", + "name": "mock-notion_notion-fetch", + "arguments": { "query": "page-001" } + } + ], + "input_tokens": 1000, + "output_tokens": 30 + } + }, + { + "response": { + "type": "text", + "content": "Here's what I found in your Notion:\n\n**Project Alpha** — Status: In Progress\n- Sprint planning on March 15\n- API redesign review pending\n\nLet me know if you want more details on any item.", + "input_tokens": 1100, + "output_tokens": 50 + } + } + ] + } + ] +} diff --git a/tests/support/mock_mcp_server.rs b/tests/support/mock_mcp_server.rs new file mode 100644 index 00000000..7919045c --- /dev/null +++ b/tests/support/mock_mcp_server.rs @@ -0,0 +1,340 @@ +//! Mock MCP server for E2E testing of the extension lifecycle. +//! +//! Provides a minimal HTTP server with: +//! - OAuth 2.1 discovery (`.well-known/oauth-protected-resource`, `.well-known/oauth-authorization-server`) +//! - Dynamic Client Registration (`/register`) +//! - Token exchange (`/token`) +//! - MCP JSON-RPC endpoint (`/mcp`) with `initialize`, `tools/list`, `tools/call` +//! +//! Tool call responses are pre-configured via `MockToolResponse`. + +#![allow(dead_code)] + +use std::collections::HashMap; +use std::net::SocketAddr; +use std::sync::Arc; + +use axum::extract::State; +use axum::http::{HeaderMap, StatusCode}; +use axum::response::IntoResponse; +use axum::routing::{get, post}; +use axum::{Json, Router}; +use serde::{Deserialize, Serialize}; +use tokio::sync::oneshot; + +/// A pre-configured response for a specific MCP tool call. +#[derive(Clone, Debug)] +pub struct MockToolResponse { + /// Tool name (e.g., "notion-search"). + pub name: String, + /// JSON response content for `tools/call`. + pub content: serde_json::Value, +} + +/// A running mock MCP server. +pub struct MockMcpServer { + /// Base URL including port (e.g., "http://127.0.0.1:12345"). + pub base_url: String, + /// Shutdown signal sender. + shutdown_tx: Option>, + /// Server task handle. + handle: Option>, +} + +impl MockMcpServer { + /// The MCP endpoint URL for use in registry entries. + pub fn mcp_url(&self) -> String { + format!("{}/mcp", self.base_url) + } + + /// Shut down the server. + pub async fn shutdown(mut self) { + if let Some(tx) = self.shutdown_tx.take() { + let _ = tx.send(()); + } + if let Some(h) = self.handle.take() { + let _ = h.await; + } + } +} + +impl Drop for MockMcpServer { + fn drop(&mut self) { + if let Some(tx) = self.shutdown_tx.take() { + let _ = tx.send(()); + } + if let Some(h) = self.handle.take() { + h.abort(); + } + } +} + +/// Shared state for the mock server handlers. +struct MockState { + /// Base URL (filled after bind). + base_url: String, + /// Tool definitions served by tools/list. + tools: Vec, + /// Pre-configured tool call responses keyed by tool name. + /// Multiple calls to the same tool return responses in order. + tool_responses: HashMap>, + /// Counter for tool_responses consumption (per tool name). + tool_response_idx: std::sync::Mutex>, +} + +#[derive(Clone, Serialize)] +struct McpToolDef { + name: String, + description: String, + #[serde(rename = "inputSchema")] + input_schema: serde_json::Value, +} + +/// Start a mock MCP server on a random port. +/// +/// `tool_responses` configures what `tools/call` returns for each tool name. +/// Multiple responses for the same tool are returned in order. +pub async fn start_mock_mcp_server(tool_responses: Vec) -> MockMcpServer { + // Build tool definitions and response map. + let mut tools = Vec::new(); + let mut response_map: HashMap> = HashMap::new(); + let mut seen_tools = std::collections::HashSet::new(); + + for tr in &tool_responses { + if seen_tools.insert(tr.name.clone()) { + tools.push(McpToolDef { + name: tr.name.clone(), + description: format!("Mock tool: {}", tr.name), + input_schema: serde_json::json!({"type": "object", "properties": {}}), + }); + } + response_map + .entry(tr.name.clone()) + .or_default() + .push(tr.content.clone()); + } + + // Bind to a random port. + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("failed to bind mock MCP server"); + let addr: SocketAddr = listener.local_addr().expect("no local addr"); + let base_url = format!("http://127.0.0.1:{}", addr.port()); + + let state = Arc::new(MockState { + base_url: base_url.clone(), + tools, + tool_responses: response_map, + tool_response_idx: std::sync::Mutex::new(HashMap::new()), + }); + + let app = Router::new() + .route( + "/.well-known/oauth-protected-resource/mcp", + get(handle_protected_resource), + ) + .route( + "/.well-known/oauth-authorization-server", + get(handle_auth_server_metadata), + ) + .route("/register", post(handle_register)) + .route("/authorize", get(handle_authorize)) + .route("/token", post(handle_token)) + .route("/mcp", post(handle_mcp)) + .with_state(state); + + let (shutdown_tx, shutdown_rx) = oneshot::channel::<()>(); + let handle = tokio::spawn(async move { + axum::serve(listener, app) + .with_graceful_shutdown(async { + let _ = shutdown_rx.await; + }) + .await + .expect("mock MCP server failed"); + }); + + // Wait briefly for the server to start accepting. + tokio::time::sleep(std::time::Duration::from_millis(50)).await; + + MockMcpServer { + base_url, + shutdown_tx: Some(shutdown_tx), + handle: Some(handle), + } +} + +// ── OAuth discovery endpoints ─────────────────────────────────────────── + +async fn handle_protected_resource(State(state): State>) -> impl IntoResponse { + Json(serde_json::json!({ + "resource": format!("{}/mcp", state.base_url), + "authorization_servers": [state.base_url], + "scopes_supported": ["read", "write"] + })) +} + +async fn handle_auth_server_metadata(State(state): State>) -> impl IntoResponse { + Json(serde_json::json!({ + "issuer": state.base_url, + "authorization_endpoint": format!("{}/authorize", state.base_url), + "token_endpoint": format!("{}/token", state.base_url), + "registration_endpoint": format!("{}/register", state.base_url), + "response_types_supported": ["code"], + "grant_types_supported": ["authorization_code"], + "code_challenge_methods_supported": ["S256"], + "scopes_supported": ["read", "write"] + })) +} + +// ── OAuth DCR ─────────────────────────────────────────────────────────── + +async fn handle_register() -> impl IntoResponse { + Json(serde_json::json!({ + "client_id": "mock-client-id", + "client_name": "ironclaw-test", + "redirect_uris": [], + "grant_types": ["authorization_code"], + "response_types": ["code"], + "token_endpoint_auth_method": "none" + })) +} + +// ── OAuth authorize (auto-approve) ────────────────────────────────────── + +/// In a real flow, this would show a consent screen. For testing, we just +/// need the endpoint to exist. The test will bypass OAuth by injecting +/// tokens directly. +async fn handle_authorize() -> impl IntoResponse { + // Return a simple HTML page; in practice the test injects tokens directly. + axum::response::Html( + "Mock OAuth: authorize endpoint. Tests bypass this.", + ) +} + +// ── OAuth token exchange ──────────────────────────────────────────────── + +async fn handle_token() -> impl IntoResponse { + Json(serde_json::json!({ + "access_token": "mock-access-token", + "token_type": "Bearer", + "expires_in": 3600, + "refresh_token": "mock-refresh-token" + })) +} + +// ── MCP JSON-RPC endpoint ─────────────────────────────────────────────── + +#[derive(Deserialize)] +struct JsonRpcRequest { + jsonrpc: String, + id: Option, + method: String, + #[serde(default)] + params: Option, +} + +async fn handle_mcp( + State(state): State>, + headers: HeaderMap, + Json(req): Json, +) -> impl IntoResponse { + // Check for auth header. + let auth = headers + .get("authorization") + .and_then(|v| v.to_str().ok()) + .unwrap_or(""); + + if !auth.starts_with("Bearer ") || &auth[7..] != "mock-access-token" { + // Return 401 with WWW-Authenticate header per MCP OAuth spec. + let www_auth = format!( + "Bearer resource_metadata=\"{}/.well-known/oauth-protected-resource/mcp\"", + state.base_url + ); + return ( + StatusCode::UNAUTHORIZED, + [("www-authenticate", www_auth.as_str())], + Json(serde_json::json!({ + "jsonrpc": "2.0", + "id": req.id, + "error": {"code": -32000, "message": "Unauthorized"} + })), + ) + .into_response(); + } + + // Handle notifications (no id) silently. + if req.id.is_none() { + return StatusCode::OK.into_response(); + } + + let response = match req.method.as_str() { + "initialize" => serde_json::json!({ + "jsonrpc": "2.0", + "id": req.id, + "result": { + "protocolVersion": "2024-11-05", + "serverInfo": { + "name": "mock-mcp-server", + "version": "1.0.0" + }, + "capabilities": { + "tools": {} + } + } + }), + "tools/list" => { + let tools: Vec = state + .tools + .iter() + .map(|t| serde_json::to_value(t).unwrap()) + .collect(); + serde_json::json!({ + "jsonrpc": "2.0", + "id": req.id, + "result": { + "tools": tools + } + }) + } + "tools/call" => { + let tool_name = req + .params + .as_ref() + .and_then(|p| p.get("name")) + .and_then(|n| n.as_str()) + .unwrap_or("unknown"); + + let content = { + let mut idx_map = state.tool_response_idx.lock().unwrap(); + let idx = idx_map.entry(tool_name.to_string()).or_insert(0); + let responses = state.tool_responses.get(tool_name); + let result = responses + .and_then(|r| r.get(*idx)) + .cloned() + .unwrap_or_else(|| serde_json::json!({"error": "no mock response configured"})); + *idx += 1; + result + }; + + serde_json::json!({ + "jsonrpc": "2.0", + "id": req.id, + "result": { + "content": [ + { + "type": "text", + "text": serde_json::to_string(&content).unwrap_or_default() + } + ] + } + }) + } + _ => serde_json::json!({ + "jsonrpc": "2.0", + "id": req.id, + "error": {"code": -32601, "message": format!("Method not found: {}", req.method)} + }), + }; + + Json(response).into_response() +} diff --git a/tests/support/mod.rs b/tests/support/mod.rs index 91321a30..3048002f 100644 --- a/tests/support/mod.rs +++ b/tests/support/mod.rs @@ -4,6 +4,7 @@ pub mod cleanup; pub mod gateway_workflow_harness; pub mod instrumented_llm; pub mod metrics; +pub mod mock_mcp_server; pub mod mock_openai_server; pub mod test_channel; pub mod test_rig; diff --git a/tests/support/test_rig.rs b/tests/support/test_rig.rs index 87782f00..07106e42 100644 --- a/tests/support/test_rig.rs +++ b/tests/support/test_rig.rs @@ -50,6 +50,9 @@ pub struct TestRig { /// The underlying TraceLlm for inspecting captured requests. #[cfg(feature = "libsql")] trace_llm: Option>, + /// Extension manager for direct extension operations in tests. + #[cfg(feature = "libsql")] + extension_manager: Option>, /// Temp directory guard -- keeps the libSQL database file alive. #[cfg(feature = "libsql")] _temp_dir: tempfile::TempDir, @@ -76,6 +79,11 @@ impl TestRig { .unwrap_or_default() } + /// Return the extension manager for direct extension operations in tests. + pub fn extension_manager(&self) -> Option<&Arc> { + self.extension_manager.as_ref() + } + /// Wait until at least `n` responses have been captured, or `timeout` elapses. pub async fn wait_for_responses(&self, n: usize, timeout: Duration) -> Vec { self.channel.wait_for_responses(n, timeout).await @@ -600,6 +608,7 @@ impl TestRigBuilder { // Save references for test accessors. let db_ref = components.db.clone().expect("test rig requires a database"); let workspace_ref = components.workspace.clone(); + let ext_mgr_ref = components.extension_manager.clone(); // 7. Construct AgentDeps from AppComponents (mirrors main.rs). let deps = AgentDeps { @@ -695,6 +704,7 @@ impl TestRigBuilder { db: db_ref, workspace: workspace_ref, trace_llm: trace_llm_ref, + extension_manager: ext_mgr_ref, _temp_dir: temp_dir, } }