From 704d63f16aef3406c25007967eccbc6f94cb3980 Mon Sep 17 00:00:00 2001 From: Henry Park Date: Wed, 4 Mar 2026 15:47:45 -0800 Subject: [PATCH] feat(oauth): route callbacks through web gateway for hosted instances (#555) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat: route OAuth callbacks through web gateway for hosted instances On hosted instances (e.g., NEAR AI), OAuth callbacks can't reach the local TCP listener on port 9876. This adds a gateway-routed OAuth flow that works behind reverse proxies and load balancers. Backend changes: - Add /oauth/callback as a public route on the web gateway - PendingOAuthFlow registry shared between ExtensionManager and handler - Gateway mode auto-detected via IRONCLAW_OAUTH_CALLBACK_URL env var - Platform state format (instance:nonce) for nginx routing - Token exchange proxy support via IRONCLAW_OAUTH_EXCHANGE_URL - Local TCP listener mode preserved as backward-compatible fallback UX improvements: - Hide Configure button for tools with auto-resolved OAuth credentials (builtin defaults or platform-injected env vars) - Skip client_id/client_secret fields in setup schema when auto-resolved - Show Reconfigure only after successful authentication Co-Authored-By: Claude Opus 4.6 (1M context) * fix(oauth): harden gateway callback and refactor AuthResult - Add 60s timeout to exchange_via_proxy HTTP client (matching exchange_oauth_code) - Read GATEWAY_AUTH_TOKEN once at ExtensionManager construction instead of per-flow from env (prevents coupling and clarifies token provenance) - Extract oauth_error_page() helper to deduplicate error landing pages - Remove IRONCLAW_FORCE_GATEWAY_CALLBACK env var (auto-detection suffices) - Refactor AuthResult into typed AuthStatus enum with constructors, eliminating stringly-typed status and Option fields that were always None - Adapt all handlers (chat, extensions, ws) to new AuthResult/AuthStatus API - Use setup_url (not validation_endpoint) for awaiting_token responses [skip-regression-check] Co-Authored-By: Claude Sonnet 4.6 * fix(oauth): address review feedback — empty token guard, test flakiness, doc typos - Fail early in exchange_via_proxy() when gateway_token is empty instead of sending an unauthenticated request to the exchange proxy - Fix test_oauth_callback_strips_instance_prefix to use an expired flow so it never attempts a real HTTP token exchange (prevents CI flakiness) - Fix doc comments: /auth/callback → /oauth/callback in PendingOAuthFlow and ExtensionManager pending_oauth_flows docs [skip-regression-check] Co-Authored-By: Claude Sonnet 4.6 * fix: clarify strip_instance_prefix safety, wrapper credential fix, test assertion - Add comment to strip_instance_prefix noting nonces are base64url (no colons) - Expand wrapper.rs comment explaining the credential_user_id bug fix - Fix test_oauth_callback_strips_instance_prefix assertion: landing_html does not include provider_name on error pages [skip-regression-check] Co-Authored-By: Claude Sonnet 4.6 --------- Co-authored-by: Claude Opus 4.6 (1M context) --- src/agent/thread_ops.rs | 10 +- src/channels/web/handlers/chat.rs | 11 +- src/channels/web/handlers/extensions.rs | 12 +- src/channels/web/server.rs | 552 +++++++++++++++- src/channels/web/static/app.js | 6 +- src/channels/web/ws.rs | 8 +- src/cli/oauth_defaults.rs | 372 +++++++++++ src/extensions/manager.rs | 843 +++++++++++++----------- src/extensions/mod.rs | 397 ++++++++++- src/tools/builtin/extension_tools.rs | 4 +- src/tools/wasm/wrapper.rs | 9 +- 11 files changed, 1759 insertions(+), 465 deletions(-) diff --git a/src/agent/thread_ops.rs b/src/agent/thread_ops.rs index 39cd22ed..b52ad3dd 100644 --- a/src/agent/thread_ops.rs +++ b/src/agent/thread_ops.rs @@ -1295,7 +1295,7 @@ impl Agent { }; match ext_mgr.auth(&pending.extension_name, Some(token)).await { - Ok(result) if result.status == "authenticated" => { + Ok(result) if result.is_authenticated() => { tracing::info!( "Extension '{}' authenticated via auth mode", pending.extension_name @@ -1364,8 +1364,8 @@ impl Agent { } } let msg = result - .instructions - .clone() + .instructions() + .map(String::from) .unwrap_or_else(|| "Invalid token. Please try again.".to_string()); // Re-emit AuthRequired so web UI re-shows the card let _ = self @@ -1375,8 +1375,8 @@ impl Agent { StatusUpdate::AuthRequired { extension_name: pending.extension_name.clone(), instructions: Some(msg.clone()), - auth_url: result.auth_url, - setup_url: result.setup_url, + auth_url: result.auth_url().map(String::from), + setup_url: result.setup_url().map(String::from), }, &message.metadata, ) diff --git a/src/channels/web/handlers/chat.rs b/src/channels/web/handlers/chat.rs index 6cab65e2..934a02dd 100644 --- a/src/channels/web/handlers/chat.rs +++ b/src/channels/web/handlers/chat.rs @@ -142,7 +142,7 @@ pub async fn chat_auth_token_handler( .await .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?; - if result.status == "authenticated" { + if result.is_authenticated() { // Auto-activate so tools are available immediately let msg = match ext_mgr.activate(&req.extension_name).await { Ok(r) => format!( @@ -170,13 +170,14 @@ pub async fn chat_auth_token_handler( // Re-emit auth_required for retry state.sse.broadcast(SseEvent::AuthRequired { extension_name: req.extension_name.clone(), - instructions: result.instructions.clone(), - auth_url: result.auth_url.clone(), - setup_url: result.setup_url.clone(), + instructions: result.instructions().map(String::from), + auth_url: result.auth_url().map(String::from), + setup_url: result.setup_url().map(String::from), }); Ok(Json(ActionResponse::fail( result - .instructions + .instructions() + .map(String::from) .unwrap_or_else(|| "Invalid token".to_string()), ))) } diff --git a/src/channels/web/handlers/extensions.rs b/src/channels/web/handlers/extensions.rs index 8199b63c..0c1f2905 100644 --- a/src/channels/web/handlers/extensions.rs +++ b/src/channels/web/handlers/extensions.rs @@ -141,7 +141,7 @@ pub async fn extensions_activate_handler( // Activation failed due to auth; try authenticating first. match ext_mgr.auth(&name, None).await { - Ok(auth_result) if auth_result.status == "authenticated" => { + Ok(auth_result) if auth_result.is_authenticated() => { // Auth succeeded, retry activation. match ext_mgr.activate(&name).await { Ok(result) => Ok(Json(ActionResponse::ok(result.message))), @@ -152,13 +152,13 @@ pub async fn extensions_activate_handler( // Auth in progress (OAuth URL or awaiting manual token). let mut resp = ActionResponse::fail( auth_result - .instructions - .clone() + .instructions() + .map(String::from) .unwrap_or_else(|| format!("'{}' requires authentication.", name)), ); - resp.auth_url = auth_result.auth_url; - resp.awaiting_token = Some(auth_result.awaiting_token); - resp.instructions = auth_result.instructions; + resp.auth_url = auth_result.auth_url().map(String::from); + resp.awaiting_token = Some(auth_result.is_awaiting_token()); + resp.instructions = auth_result.instructions().map(String::from); Ok(Json(resp)) } Err(auth_err) => Ok(Json(ActionResponse::fail(format!( diff --git a/src/channels/web/server.rs b/src/channels/web/server.rs index bada142a..9fe3ac3d 100644 --- a/src/channels/web/server.rs +++ b/src/channels/web/server.rs @@ -192,7 +192,9 @@ pub async fn start_server( })?; // Public routes (no auth) - let public = Router::new().route("/api/health", get(health_handler)); + let public = Router::new() + .route("/api/health", get(health_handler)) + .route("/oauth/callback", get(oauth_callback_handler)); // Protected routes (require auth) let auth_state = AuthState { token: auth_token }; @@ -424,6 +426,180 @@ async fn health_handler() -> Json { }) } +/// Return an OAuth error landing page response. +fn oauth_error_page(label: &str) -> axum::response::Response { + let html = crate::cli::oauth_defaults::landing_html(label, false); + axum::response::Html(html).into_response() +} + +/// OAuth callback handler for the web gateway. +/// +/// This is a PUBLIC route (no Bearer token required) because OAuth providers +/// redirect the user's browser here. The `state` query parameter correlates +/// the callback with a pending OAuth flow registered by `start_wasm_oauth()`. +/// +/// Used on hosted instances where `IRONCLAW_OAUTH_CALLBACK_URL` points to +/// the gateway (e.g., `https://kind-deer.agent1.near.ai/oauth/callback`). +/// Local/desktop mode continues to use the TCP listener on port 9876. +async fn oauth_callback_handler( + State(state): State>, + Query(params): Query>, +) -> impl IntoResponse { + use crate::cli::oauth_defaults; + + // Check for error from OAuth provider (e.g., user denied consent) + if let Some(error) = params.get("error") { + let description = params + .get("error_description") + .cloned() + .unwrap_or_else(|| error.clone()); + return oauth_error_page(&description); + } + + let state_param = match params.get("state") { + Some(s) if !s.is_empty() => s.clone(), + _ => return oauth_error_page("IronClaw"), + }; + + let code = match params.get("code") { + Some(c) if !c.is_empty() => c.clone(), + _ => return oauth_error_page("IronClaw"), + }; + + // Look up the pending flow by CSRF state (atomic remove prevents replay) + let ext_mgr = match state.extension_manager.as_ref() { + Some(mgr) => mgr, + None => return oauth_error_page("IronClaw"), + }; + + // Strip instance prefix from state for registry lookup. + // Platform nginx sends `state=instance:nonce` but flows are keyed by nonce only. + let lookup_key = oauth_defaults::strip_instance_prefix(&state_param); + + let flow = ext_mgr + .pending_oauth_flows() + .write() + .await + .remove(lookup_key); + + let flow = match flow { + Some(f) => f, + None => { + tracing::warn!( + state = %state_param, + lookup_key = %lookup_key, + "OAuth callback received with unknown or expired state" + ); + return oauth_error_page("IronClaw"); + } + }; + + // Check flow expiry (5 minutes, matching TCP listener timeout) + if flow.created_at.elapsed() > oauth_defaults::OAUTH_FLOW_EXPIRY { + tracing::warn!( + extension = %flow.extension_name, + "OAuth flow expired" + ); + return oauth_error_page(&flow.display_name); + } + + // Exchange the authorization code for tokens. + // Use the platform exchange proxy when configured (keeps client_secret off container), + // otherwise call the provider's token URL directly. + 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 gateway_token = flow.gateway_token.as_deref().unwrap_or_default(); + oauth_defaults::exchange_via_proxy( + proxy_url, + gateway_token, + &code, + &flow.redirect_uri, + flow.code_verifier.as_deref(), + &flow.access_token_field, + ) + .await + .map_err(|e| e.to_string())? + } else { + oauth_defaults::exchange_oauth_code( + &flow.token_url, + &flow.client_id, + flow.client_secret.as_deref(), + &code, + &flow.redirect_uri, + flow.code_verifier.as_deref(), + &flow.access_token_field, + ) + .await + .map_err(|e| e.to_string())? + }; + + // Validate the token before storing (catches wrong account, etc.) + if let Some(ref validation) = flow.validation_endpoint { + oauth_defaults::validate_oauth_token(&token_response.access_token, validation) + .await + .map_err(|e| e.to_string())?; + } + + // Store tokens encrypted in the secrets store + oauth_defaults::store_oauth_tokens( + flow.secrets.as_ref(), + &flow.user_id, + &flow.secret_name, + flow.provider.as_deref(), + &token_response.access_token, + token_response.refresh_token.as_deref(), + token_response.expires_in, + &flow.scopes, + ) + .await + .map_err(|e| e.to_string())?; + + Ok(()) + } + .await; + + let (success, message) = match &result { + Ok(()) => ( + true, + format!("{} authenticated successfully", flow.display_name), + ), + Err(e) => ( + false, + format!("{} authentication failed: {}", flow.display_name, e), + ), + }; + + match &result { + Ok(()) => { + tracing::info!( + extension = %flow.extension_name, + "OAuth completed successfully via gateway callback" + ); + } + Err(e) => { + tracing::warn!( + extension = %flow.extension_name, + error = %e, + "OAuth failed via gateway callback" + ); + } + } + + // 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, + }); + } + + let html = oauth_defaults::landing_html(&flow.display_name, success); + axum::response::Html(html).into_response() +} + // --- Chat handlers --- async fn chat_send_handler( @@ -552,7 +728,7 @@ async fn chat_auth_token_handler( .await .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?; - if result.status == "authenticated" { + if result.is_authenticated() { // Auto-activate so tools are available immediately let msg = match ext_mgr.activate(&req.extension_name).await { Ok(r) => format!( @@ -580,13 +756,14 @@ async fn chat_auth_token_handler( // Re-emit auth_required for retry state.sse.broadcast(SseEvent::AuthRequired { extension_name: req.extension_name.clone(), - instructions: result.instructions.clone(), - auth_url: result.auth_url.clone(), - setup_url: result.setup_url.clone(), + instructions: result.instructions().map(String::from), + auth_url: result.auth_url().map(String::from), + setup_url: result.setup_url().map(String::from), }); Ok(Json(ActionResponse::fail( result - .instructions + .instructions() + .map(String::from) .unwrap_or_else(|| "Invalid token".to_string()), ))) } @@ -1332,12 +1509,9 @@ async fn extensions_install_handler( // configured (e.g., built-in providers). We only surface an auth_url // when the extension reports it is awaiting authorization. match ext_mgr.auth(&req.name, None).await { - Ok(auth_result) - if auth_result.auth_url.is_some() - && auth_result.status == "awaiting_authorization" => - { + Ok(auth_result) if auth_result.auth_url().is_some() => { // Scope expansion or initial OAuth: user needs to authorize - resp.auth_url = auth_result.auth_url; + resp.auth_url = auth_result.auth_url().map(String::from); } _ => {} } @@ -1366,10 +1540,9 @@ async fn extensions_activate_handler( // Initial OAuth setup is triggered via save_setup_secrets. let mut resp = ActionResponse::ok(result.message); if let Ok(auth_result) = ext_mgr.auth(&name, None).await - && auth_result.auth_url.is_some() - && auth_result.status == "awaiting_authorization" + && auth_result.auth_url().is_some() { - resp.auth_url = auth_result.auth_url; + resp.auth_url = auth_result.auth_url().map(String::from); } Ok(Json(resp)) } @@ -1385,7 +1558,7 @@ async fn extensions_activate_handler( // Activation failed due to auth; try authenticating first. match ext_mgr.auth(&name, None).await { - Ok(auth_result) if auth_result.status == "authenticated" => { + Ok(auth_result) if auth_result.is_authenticated() => { // Auth succeeded, retry activation. match ext_mgr.activate(&name).await { Ok(result) => Ok(Json(ActionResponse::ok(result.message))), @@ -1396,13 +1569,13 @@ async fn extensions_activate_handler( // Auth in progress (OAuth URL or awaiting manual token). let mut resp = ActionResponse::fail( auth_result - .instructions - .clone() + .instructions() + .map(String::from) .unwrap_or_else(|| format!("'{}' requires authentication.", name)), ); - resp.auth_url = auth_result.auth_url; - resp.awaiting_token = Some(auth_result.awaiting_token); - resp.instructions = auth_result.instructions; + resp.auth_url = auth_result.auth_url().map(String::from); + resp.awaiting_token = Some(auth_result.is_awaiting_token()); + resp.instructions = auth_result.instructions().map(String::from); Ok(Json(resp)) } Err(auth_err) => Ok(Json(ActionResponse::fail(format!( @@ -2239,4 +2412,343 @@ mod tests { let turns = build_turns_from_db_messages(&[]); assert!(turns.is_empty()); } + + // --- OAuth callback handler tests --- + + /// Build a minimal `GatewayState` for testing the OAuth callback handler. + fn test_gateway_state(ext_mgr: Option>) -> Arc { + Arc::new(GatewayState { + msg_tx: tokio::sync::RwLock::new(None), + sse: SseManager::new(), + workspace: None, + session_manager: None, + log_broadcaster: None, + log_level_handle: None, + extension_manager: ext_mgr, + tool_registry: None, + store: None, + job_manager: None, + prompt_queue: None, + user_id: "test".to_string(), + shutdown_tx: tokio::sync::RwLock::new(None), + ws_tracker: None, + llm_provider: None, + skill_registry: None, + skill_catalog: None, + scheduler: None, + chat_rate_limiter: RateLimiter::new(30, 60), + registry_entries: vec![], + cost_guard: None, + startup_time: std::time::Instant::now(), + }) + } + + /// Build a test router with just the OAuth callback route. + fn test_oauth_router(state: Arc) -> Router { + Router::new() + .route("/oauth/callback", get(oauth_callback_handler)) + .with_state(state) + } + + #[tokio::test] + async fn test_oauth_callback_missing_params() { + use axum::body::Body; + use tower::ServiceExt; + + let state = test_gateway_state(None); + let app = test_oauth_router(state); + + let req = axum::http::Request::builder() + .uri("/oauth/callback") + .body(Body::empty()) + .expect("request"); + + let resp = ServiceExt::>::oneshot(app, req) + .await + .expect("response"); + assert_eq!(resp.status(), StatusCode::OK); + + let body = axum::body::to_bytes(resp.into_body(), 1024 * 64) + .await + .expect("body"); + let html = String::from_utf8_lossy(&body); + assert!(html.contains("Authorization Failed")); + } + + #[tokio::test] + async fn test_oauth_callback_error_from_provider() { + use axum::body::Body; + use tower::ServiceExt; + + let state = test_gateway_state(None); + let app = test_oauth_router(state); + + let req = axum::http::Request::builder() + .uri("/oauth/callback?error=access_denied&error_description=access_denied") + .body(Body::empty()) + .expect("request"); + + let resp = ServiceExt::>::oneshot(app, req) + .await + .expect("response"); + assert_eq!(resp.status(), StatusCode::OK); + + let body = axum::body::to_bytes(resp.into_body(), 1024 * 64) + .await + .expect("body"); + let html = String::from_utf8_lossy(&body); + assert!(html.contains("Authorization Failed")); + } + + #[tokio::test] + async fn test_oauth_callback_unknown_state() { + use axum::body::Body; + use tower::ServiceExt; + + // Build an ExtensionManager so the handler can look up flows + let secrets = Arc::new(crate::secrets::InMemorySecretsStore::new(Arc::new( + crate::secrets::SecretsCrypto::new(secrecy::SecretString::from( + "test-key-at-least-32-chars-long!!".to_string(), + )) + .expect("crypto"), + ))); + let tool_registry = Arc::new(ToolRegistry::new()); + let mcp_sm = Arc::new(crate::tools::mcp::session::McpSessionManager::new()); + + let ext_mgr = Arc::new(ExtensionManager::new( + mcp_sm, + secrets, + tool_registry, + None, + None, + std::path::PathBuf::from("/tmp/wasm_tools"), + std::path::PathBuf::from("/tmp/wasm_channels"), + None, + "test".to_string(), + None, + vec![], + )); + + let state = test_gateway_state(Some(ext_mgr)); + let app = test_oauth_router(state); + + let req = axum::http::Request::builder() + .uri("/oauth/callback?code=test_code&state=unknown_state_value") + .body(Body::empty()) + .expect("request"); + + let resp = ServiceExt::>::oneshot(app, req) + .await + .expect("response"); + assert_eq!(resp.status(), StatusCode::OK); + + let body = axum::body::to_bytes(resp.into_body(), 1024 * 64) + .await + .expect("body"); + let html = String::from_utf8_lossy(&body); + assert!(html.contains("Authorization Failed")); + } + + #[tokio::test] + async fn test_oauth_callback_expired_flow() { + use axum::body::Body; + use tower::ServiceExt; + + let secrets: Arc = + Arc::new(crate::secrets::InMemorySecretsStore::new(Arc::new( + crate::secrets::SecretsCrypto::new(secrecy::SecretString::from( + "test-key-at-least-32-chars-long!!".to_string(), + )) + .expect("crypto"), + ))); + let tool_registry = Arc::new(ToolRegistry::new()); + let mcp_sm = Arc::new(crate::tools::mcp::session::McpSessionManager::new()); + + let ext_mgr = Arc::new(ExtensionManager::new( + mcp_sm, + secrets.clone(), + tool_registry, + None, + None, + std::path::PathBuf::from("/tmp/wasm_tools"), + std::path::PathBuf::from("/tmp/wasm_channels"), + None, + "test".to_string(), + None, + vec![], + )); + + // Insert an expired flow (created 10 minutes ago) + let flow = crate::cli::oauth_defaults::PendingOAuthFlow { + extension_name: "test_tool".to_string(), + display_name: "Test Tool".to_string(), + token_url: "https://example.com/token".to_string(), + client_id: "client123".to_string(), + client_secret: None, + redirect_uri: "https://example.com/oauth/callback".to_string(), + code_verifier: None, + access_token_field: "access_token".to_string(), + secret_name: "test_token".to_string(), + provider: None, + validation_endpoint: None, + scopes: vec![], + user_id: "test".to_string(), + secrets, + sse_sender: None, + gateway_token: None, + created_at: std::time::Instant::now() - std::time::Duration::from_secs(600), + }; + + ext_mgr + .pending_oauth_flows() + .write() + .await + .insert("expired_state".to_string(), flow); + + let state = test_gateway_state(Some(ext_mgr)); + let app = test_oauth_router(state); + + let req = axum::http::Request::builder() + .uri("/oauth/callback?code=test_code&state=expired_state") + .body(Body::empty()) + .expect("request"); + + let resp = ServiceExt::>::oneshot(app, req) + .await + .expect("response"); + assert_eq!(resp.status(), StatusCode::OK); + + let body = axum::body::to_bytes(resp.into_body(), 1024 * 64) + .await + .expect("body"); + let html = String::from_utf8_lossy(&body); + // Expired flow → error landing page + assert!(html.contains("Authorization Failed")); + } + + #[tokio::test] + async fn test_oauth_callback_no_extension_manager() { + use axum::body::Body; + use tower::ServiceExt; + + // No extension manager set → graceful error + let state = test_gateway_state(None); + let app = test_oauth_router(state); + + let req = axum::http::Request::builder() + .uri("/oauth/callback?code=test_code&state=some_state") + .body(Body::empty()) + .expect("request"); + + let resp = ServiceExt::>::oneshot(app, req) + .await + .expect("response"); + assert_eq!(resp.status(), StatusCode::OK); + + let body = axum::body::to_bytes(resp.into_body(), 1024 * 64) + .await + .expect("body"); + let html = String::from_utf8_lossy(&body); + assert!(html.contains("Authorization Failed")); + } + + #[tokio::test] + async fn test_oauth_callback_strips_instance_prefix() { + use axum::body::Body; + use tower::ServiceExt; + + let secrets: Arc = + Arc::new(crate::secrets::InMemorySecretsStore::new(Arc::new( + crate::secrets::SecretsCrypto::new(secrecy::SecretString::from( + "test-key-at-least-32-chars-long!!".to_string(), + )) + .expect("crypto"), + ))); + let tool_registry = Arc::new(ToolRegistry::new()); + let mcp_sm = Arc::new(crate::tools::mcp::session::McpSessionManager::new()); + + let ext_mgr = Arc::new(ExtensionManager::new( + mcp_sm, + secrets.clone(), + tool_registry, + None, + None, + std::path::PathBuf::from("/tmp/wasm_tools"), + std::path::PathBuf::from("/tmp/wasm_channels"), + None, + "test".to_string(), + None, + vec![], + )); + + // Insert a flow keyed by raw nonce "test_nonce" (without instance prefix). + // Use an expired flow so the handler exits before attempting a real HTTP + // token exchange — we only need to verify that the instance prefix was + // stripped and the flow was found by the raw nonce. + let flow = crate::cli::oauth_defaults::PendingOAuthFlow { + extension_name: "test_tool".to_string(), + display_name: "Test Tool".to_string(), + token_url: "https://example.com/token".to_string(), + client_id: "client123".to_string(), + client_secret: None, + redirect_uri: "https://example.com/oauth/callback".to_string(), + code_verifier: None, + access_token_field: "access_token".to_string(), + secret_name: "test_token".to_string(), + provider: None, + validation_endpoint: None, + scopes: vec![], + user_id: "test".to_string(), + secrets, + sse_sender: None, + gateway_token: None, + // Expired — handler will reject after lookup (no network I/O) + created_at: std::time::Instant::now() - std::time::Duration::from_secs(600), + }; + + ext_mgr + .pending_oauth_flows() + .write() + .await + .insert("test_nonce".to_string(), flow); + + let state = test_gateway_state(Some(ext_mgr.clone())); + let app = test_oauth_router(state); + + // Send callback with instance prefix: "myinstance:test_nonce" + // The handler should strip "myinstance:" and find the flow keyed by "test_nonce" + let req = axum::http::Request::builder() + .uri("/oauth/callback?code=fake_code&state=myinstance:test_nonce") + .body(Body::empty()) + .expect("request"); + + let resp = ServiceExt::>::oneshot(app, req) + .await + .expect("response"); + assert_eq!(resp.status(), StatusCode::OK); + + let body = axum::body::to_bytes(resp.into_body(), 1024 * 64) + .await + .expect("body"); + let html = String::from_utf8_lossy(&body); + + // The flow was found (stripped prefix matched) but is expired, so the + // handler returns an error landing page. The flow being consumed from + // the registry (checked below) proves the prefix was stripped correctly. + assert!( + html.contains("Authorization Failed"), + "Expected error page, html was: {}", + &html[..html.len().min(500)] + ); + + // Verify the flow was consumed (removed from registry) + assert!( + ext_mgr + .pending_oauth_flows() + .read() + .await + .get("test_nonce") + .is_none() + ); + } } diff --git a/src/channels/web/static/app.js b/src/channels/web/static/app.js index 138b0dee..1d956cf3 100644 --- a/src/channels/web/static/app.js +++ b/src/channels/web/static/app.js @@ -2013,7 +2013,11 @@ function renderExtensionCard(ext) { actions.appendChild(activateBtn); } - if (ext.needs_setup || ext.has_auth) { + // Show Configure/Reconfigure button when there are secrets to enter. + // Skip when has_auth is true but needs_setup is false and not yet authenticated — + // this means OAuth credentials resolve automatically (builtin/env) and the user + // just needs to complete the OAuth flow, not fill in a config form. + if (ext.needs_setup || (ext.has_auth && ext.authenticated)) { const configBtn = document.createElement('button'); configBtn.className = 'btn-ext configure'; configBtn.textContent = ext.authenticated ? 'Reconfigure' : 'Configure'; diff --git a/src/channels/web/ws.rs b/src/channels/web/ws.rs index 527daf4a..2477217e 100644 --- a/src/channels/web/ws.rs +++ b/src/channels/web/ws.rs @@ -242,7 +242,7 @@ async fn handle_client_message( } => { if let Some(ref ext_mgr) = state.extension_manager { match ext_mgr.auth(&extension_name, Some(&token)).await { - Ok(result) if result.status == "authenticated" => { + Ok(result) if result.is_authenticated() => { let msg = match ext_mgr.activate(&extension_name).await { Ok(r) => format!( "{} authenticated ({} tools loaded)", @@ -268,9 +268,9 @@ async fn handle_client_message( .sse .broadcast(crate::channels::web::types::SseEvent::AuthRequired { extension_name, - instructions: result.instructions, - auth_url: result.auth_url, - setup_url: result.setup_url, + instructions: result.instructions().map(String::from), + auth_url: result.auth_url().map(String::from), + setup_url: result.setup_url().map(String::from), }); } Err(e) => { diff --git a/src/cli/oauth_defaults.rs b/src/cli/oauth_defaults.rs index 8f8cd3a7..75ab7856 100644 --- a/src/cli/oauth_defaults.rs +++ b/src/cli/oauth_defaults.rs @@ -18,6 +18,7 @@ //! env vars, which take priority over built-in defaults. use std::collections::HashMap; +use std::sync::Arc; use std::time::Duration; use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD}; @@ -25,6 +26,7 @@ use rand::RngCore; use sha2::{Digest, Sha256}; use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader}; use tokio::net::TcpListener; +use tokio::sync::RwLock; use crate::secrets::{CreateSecretParams, SecretsStore}; @@ -683,6 +685,219 @@ pub fn landing_html(provider_name: &str, success: bool) -> String { ) } +// ── Gateway callback support ───────────────────────────────────────── + +/// State for an in-progress OAuth flow, keyed by CSRF `state` parameter. +/// +/// Created by `start_wasm_oauth()` and consumed by the web gateway's +/// `/oauth/callback` handler when running in hosted mode. +pub struct PendingOAuthFlow { + /// Extension name (e.g., "google_calendar"). + pub extension_name: String, + /// Human-readable display name (e.g., "Google Calendar"). + pub display_name: String, + /// OAuth token exchange URL. + pub token_url: String, + /// OAuth client ID. + pub client_id: String, + /// OAuth client secret (optional for PKCE-only flows). + pub client_secret: Option, + /// The redirect_uri used in the authorization request. + pub redirect_uri: String, + /// PKCE code verifier (must match the code_challenge sent in the auth URL). + pub code_verifier: Option, + /// Field name in token response containing the access token. + pub access_token_field: String, + /// Secret name for storage (e.g., "google_oauth_token"). + pub secret_name: String, + /// Provider hint (e.g., "google"). + pub provider: Option, + /// Token validation endpoint (optional). + pub validation_endpoint: Option, + /// Scopes that were requested. + pub scopes: Vec, + /// User ID for secret storage. + pub user_id: String, + /// Secrets store reference for token persistence. + pub secrets: Arc, + /// SSE broadcast sender for notifying the web UI. + pub sse_sender: Option>, + /// Gateway auth token for authenticating with the platform token exchange proxy. + pub gateway_token: Option, + /// When this flow was created (for expiry). + pub created_at: std::time::Instant, +} + +impl std::fmt::Debug for PendingOAuthFlow { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("PendingOAuthFlow") + .field("extension_name", &self.extension_name) + .field("display_name", &self.display_name) + .field("secret_name", &self.secret_name) + .field("created_at", &self.created_at) + .finish_non_exhaustive() + } +} + +/// Thread-safe registry of pending OAuth flows, keyed by CSRF `state` parameter. +pub type PendingOAuthRegistry = Arc>>; + +/// Create a new empty pending OAuth flow registry. +pub fn new_pending_oauth_registry() -> PendingOAuthRegistry { + Arc::new(RwLock::new(HashMap::new())) +} + +/// Returns `true` if OAuth callbacks should be routed through the web gateway +/// instead of the local TCP listener. +/// +/// This is the case when `IRONCLAW_OAUTH_CALLBACK_URL` is set to a non-loopback +/// URL, meaning the user's browser will redirect to a hosted gateway rather than +/// localhost. +pub fn use_gateway_callback() -> bool { + std::env::var("IRONCLAW_OAUTH_CALLBACK_URL") + .ok() + .filter(|v| !v.is_empty()) + .map(|raw| { + url::Url::parse(&raw) + .ok() + .and_then(|u| u.host_str().map(String::from)) + .map(|host| !is_loopback_host(&host)) + .unwrap_or(false) + }) + .unwrap_or(false) +} + +/// Maximum age for pending OAuth flows (5 minutes, matching TCP listener timeout). +pub const OAUTH_FLOW_EXPIRY: Duration = Duration::from_secs(300); + +/// Remove expired flows from the registry. +/// +/// Called when inserting new flows to prevent accumulation from abandoned +/// OAuth attempts. +pub async fn sweep_expired_flows(registry: &PendingOAuthRegistry) { + let mut flows = registry.write().await; + flows.retain(|_, flow| flow.created_at.elapsed() < OAUTH_FLOW_EXPIRY); +} + +// ── Platform routing helpers ──────────────────────────────────────── + +/// Prepend instance name to CSRF state for platform routing. +/// +/// The NEAR AI platform nginx proxy at `auth.DOMAIN` parses the instance name +/// from the `state` query parameter (format: `instance:nonce`) to route the +/// OAuth callback to the correct container. +/// +/// Returns the nonce unchanged when `IRONCLAW_INSTANCE_NAME` is not set +/// (local/non-platform mode). +pub fn build_platform_state(nonce: &str) -> String { + let instance = std::env::var("IRONCLAW_INSTANCE_NAME") + .or_else(|_| std::env::var("OPENCLAW_INSTANCE_NAME")) + .ok() + .filter(|v| !v.is_empty()); + match instance { + Some(name) => format!("{}:{}", name, nonce), + None => nonce.to_string(), + } +} + +/// Strip the instance prefix from a state parameter to recover the lookup nonce. +/// +/// `"myinstance:abc123"` → `"abc123"`, `"abc123"` → `"abc123"` (no prefix). +/// +/// Safe because nonces are base64url-encoded (`[A-Za-z0-9_-]`, no colons). +pub fn strip_instance_prefix(state: &str) -> &str { + state + .split_once(':') + .map(|(_, nonce)| nonce) + .unwrap_or(state) +} + +/// Exchange an OAuth authorization code via the platform's token exchange proxy. +/// +/// The proxy holds `client_secret` server-side so the container never sees it. +/// Authenticated via the gateway auth token (Bearer header). +/// +/// The proxy expects form params `{code, redirect_uri, code_verifier}` and +/// returns a standard Google token response `{access_token, refresh_token, expires_in}`. +pub async fn exchange_via_proxy( + proxy_url: &str, + gateway_token: &str, + code: &str, + redirect_uri: &str, + code_verifier: Option<&str>, + access_token_field: &str, +) -> Result { + if gateway_token.is_empty() { + return Err(OAuthCallbackError::Io( + "Gateway auth token is required for proxy token exchange".to_string(), + )); + } + let exchange_url = format!("{}/oauth/exchange", proxy_url.trim_end_matches('/')); + + let client = reqwest::Client::builder() + .timeout(Duration::from_secs(60)) + .build() + .map_err(|e| OAuthCallbackError::Io(format!("Failed to build HTTP client: {}", e)))?; + let mut params = vec![ + ("code", code.to_string()), + ("redirect_uri", redirect_uri.to_string()), + ]; + if let Some(verifier) = code_verifier { + params.push(("code_verifier", verifier.to_string())); + } + + let response = client + .post(&exchange_url) + .bearer_auth(gateway_token) + .form(¶ms) + .send() + .await + .map_err(|e| { + OAuthCallbackError::Io(format!("Token exchange proxy request failed: {}", e)) + })?; + + if !response.status().is_success() { + let status = response.status(); + let body = response.text().await.unwrap_or_default(); + return Err(OAuthCallbackError::Io(format!( + "Token exchange proxy failed: {} - {}", + status, body + ))); + } + + let token_data: serde_json::Value = response + .json() + .await + .map_err(|e| OAuthCallbackError::Io(format!("Failed to parse proxy response: {}", e)))?; + + let access_token = token_data + .get(access_token_field) + .and_then(|v| v.as_str()) + .ok_or_else(|| { + let fields: Vec<&str> = token_data + .as_object() + .map(|o| o.keys().map(|k| k.as_str()).collect()) + .unwrap_or_default(); + OAuthCallbackError::Io(format!( + "No '{}' field in proxy response (fields present: {:?})", + access_token_field, fields + )) + })? + .to_string(); + + let refresh_token = token_data + .get("refresh_token") + .and_then(|v| v.as_str()) + .map(String::from); + let expires_in = token_data.get("expires_in").and_then(|v| v.as_u64()); + + Ok(OAuthTokenResponse { + access_token, + refresh_token, + expires_in, + }) +} + #[cfg(test)] mod tests { use std::sync::Mutex; @@ -939,4 +1154,161 @@ mod tests { // State should be different each time (random) assert_ne!(result1.state, result2.state); } + + #[test] + fn test_use_gateway_callback_false_by_default() { + let _guard = ENV_MUTEX.lock().expect("env mutex poisoned"); + let original = std::env::var("IRONCLAW_OAUTH_CALLBACK_URL").ok(); + // SAFETY: Under ENV_MUTEX, no concurrent env access. + unsafe { + std::env::remove_var("IRONCLAW_OAUTH_CALLBACK_URL"); + } + assert!(!crate::cli::oauth_defaults::use_gateway_callback()); + unsafe { + if let Some(val) = original { + std::env::set_var("IRONCLAW_OAUTH_CALLBACK_URL", val); + } + } + } + + #[test] + fn test_use_gateway_callback_true_for_hosted() { + let _guard = ENV_MUTEX.lock().expect("env mutex poisoned"); + let original = std::env::var("IRONCLAW_OAUTH_CALLBACK_URL").ok(); + // SAFETY: Under ENV_MUTEX, no concurrent env access. + unsafe { + std::env::set_var( + "IRONCLAW_OAUTH_CALLBACK_URL", + "https://kind-deer.agent1.near.ai", + ); + } + assert!(crate::cli::oauth_defaults::use_gateway_callback()); + unsafe { + if let Some(val) = original { + std::env::set_var("IRONCLAW_OAUTH_CALLBACK_URL", val); + } else { + std::env::remove_var("IRONCLAW_OAUTH_CALLBACK_URL"); + } + } + } + + #[test] + fn test_use_gateway_callback_false_for_localhost() { + let _guard = ENV_MUTEX.lock().expect("env mutex poisoned"); + let original = std::env::var("IRONCLAW_OAUTH_CALLBACK_URL").ok(); + // SAFETY: Under ENV_MUTEX, no concurrent env access. + unsafe { + std::env::set_var("IRONCLAW_OAUTH_CALLBACK_URL", "http://127.0.0.1:3001"); + } + assert!(!crate::cli::oauth_defaults::use_gateway_callback()); + unsafe { + if let Some(val) = original { + std::env::set_var("IRONCLAW_OAUTH_CALLBACK_URL", val); + } else { + std::env::remove_var("IRONCLAW_OAUTH_CALLBACK_URL"); + } + } + } + + #[test] + fn test_use_gateway_callback_false_for_empty() { + let _guard = ENV_MUTEX.lock().expect("env mutex poisoned"); + let original = std::env::var("IRONCLAW_OAUTH_CALLBACK_URL").ok(); + // SAFETY: Under ENV_MUTEX, no concurrent env access. + unsafe { + std::env::set_var("IRONCLAW_OAUTH_CALLBACK_URL", ""); + } + assert!(!crate::cli::oauth_defaults::use_gateway_callback()); + unsafe { + if let Some(val) = original { + std::env::set_var("IRONCLAW_OAUTH_CALLBACK_URL", val); + } else { + std::env::remove_var("IRONCLAW_OAUTH_CALLBACK_URL"); + } + } + } + + #[test] + fn test_build_platform_state_with_instance() { + use crate::cli::oauth_defaults::build_platform_state; + + let _guard = ENV_MUTEX.lock().expect("env mutex poisoned"); + let original = std::env::var("IRONCLAW_INSTANCE_NAME").ok(); + // SAFETY: Under ENV_MUTEX, no concurrent env access. + unsafe { + std::env::set_var("IRONCLAW_INSTANCE_NAME", "kind-deer"); + } + assert_eq!(build_platform_state("abc123"), "kind-deer:abc123"); + unsafe { + if let Some(val) = original { + std::env::set_var("IRONCLAW_INSTANCE_NAME", val); + } else { + std::env::remove_var("IRONCLAW_INSTANCE_NAME"); + } + } + } + + #[test] + fn test_build_platform_state_without_instance() { + use crate::cli::oauth_defaults::build_platform_state; + + let _guard = ENV_MUTEX.lock().expect("env mutex poisoned"); + let original = std::env::var("IRONCLAW_INSTANCE_NAME").ok(); + let original_oc = std::env::var("OPENCLAW_INSTANCE_NAME").ok(); + // SAFETY: Under ENV_MUTEX, no concurrent env access. + unsafe { + std::env::remove_var("IRONCLAW_INSTANCE_NAME"); + std::env::remove_var("OPENCLAW_INSTANCE_NAME"); + } + assert_eq!(build_platform_state("abc123"), "abc123"); + unsafe { + if let Some(val) = original { + std::env::set_var("IRONCLAW_INSTANCE_NAME", val); + } + if let Some(val) = original_oc { + std::env::set_var("OPENCLAW_INSTANCE_NAME", val); + } + } + } + + #[test] + fn test_build_platform_state_with_openclaw_instance() { + use crate::cli::oauth_defaults::build_platform_state; + + let _guard = ENV_MUTEX.lock().expect("env mutex poisoned"); + let original_ic = std::env::var("IRONCLAW_INSTANCE_NAME").ok(); + let original_oc = std::env::var("OPENCLAW_INSTANCE_NAME").ok(); + // SAFETY: Under ENV_MUTEX, no concurrent env access. + unsafe { + std::env::remove_var("IRONCLAW_INSTANCE_NAME"); + std::env::set_var("OPENCLAW_INSTANCE_NAME", "quiet-lion"); + } + assert_eq!(build_platform_state("xyz789"), "quiet-lion:xyz789"); + unsafe { + if let Some(val) = original_ic { + std::env::set_var("IRONCLAW_INSTANCE_NAME", val); + } + if let Some(val) = original_oc { + std::env::set_var("OPENCLAW_INSTANCE_NAME", val); + } else { + std::env::remove_var("OPENCLAW_INSTANCE_NAME"); + } + } + } + + #[test] + fn test_strip_instance_prefix_with_colon() { + use crate::cli::oauth_defaults::strip_instance_prefix; + + assert_eq!(strip_instance_prefix("kind-deer:abc123"), "abc123"); + assert_eq!(strip_instance_prefix("my-instance:xyz"), "xyz"); + } + + #[test] + fn test_strip_instance_prefix_without_colon() { + use crate::cli::oauth_defaults::strip_instance_prefix; + + assert_eq!(strip_instance_prefix("abc123"), "abc123"); + assert_eq!(strip_instance_prefix(""), ""); + } } diff --git a/src/extensions/manager.rs b/src/extensions/manager.rs index 7f941ea4..3bae444d 100644 --- a/src/extensions/manager.rs +++ b/src/extensions/manager.rs @@ -18,7 +18,7 @@ use crate::extensions::discovery::OnlineDiscovery; use crate::extensions::registry::ExtensionRegistry; use crate::extensions::{ ActivateResult, AuthResult, ExtensionError, ExtensionKind, ExtensionSource, InstallResult, - InstalledExtension, RegistryEntry, ResultSource, SearchResult, + InstalledExtension, RegistryEntry, ResultSource, SearchResult, ToolAuthState, }; use crate::hooks::HookRegistry; use crate::pairing::PairingStore; @@ -100,6 +100,15 @@ pub struct ExtensionManager { /// SSE broadcast sender (set post-construction via `set_sse_sender()`). sse_sender: RwLock>>, + /// Shared registry of pending OAuth flows for gateway-routed callbacks. + /// + /// Keyed by CSRF `state` parameter. Populated in `start_wasm_oauth()` + /// when running in gateway mode, consumed by the web gateway's + /// `/oauth/callback` handler. + pending_oauth_flows: crate::cli::oauth_defaults::PendingOAuthRegistry, + /// Gateway auth token for authenticating with the platform token exchange proxy. + /// Read once at construction from `GATEWAY_AUTH_TOKEN` env var. + gateway_token: Option, } impl ExtensionManager { @@ -141,6 +150,8 @@ impl ExtensionManager { active_channel_names: RwLock::new(HashSet::new()), activation_errors: RwLock::new(HashMap::new()), sse_sender: RwLock::new(None), + pending_oauth_flows: crate::cli::oauth_defaults::new_pending_oauth_registry(), + gateway_token: std::env::var("GATEWAY_AUTH_TOKEN").ok(), } } @@ -228,6 +239,14 @@ impl ExtensionManager { *self.sse_sender.write().await = Some(sender); } + /// Returns the pending OAuth flow registry for sharing with the web gateway. + /// + /// The gateway's `/oauth/callback` handler uses this to look up pending flows + /// by CSRF `state` parameter and complete the token exchange. + pub fn pending_oauth_flows(&self) -> &crate::cli::oauth_defaults::PendingOAuthRegistry { + &self.pending_oauth_flows + } + /// Broadcast an extension status change to the web UI via SSE. async fn broadcast_extension_status(&self, name: &str, status: &str, message: Option<&str>) { if let Some(ref sender) = *self.sse_sender.read().await { @@ -416,23 +435,18 @@ impl ExtensionManager { .get_with_kind(&name, Some(ExtensionKind::WasmTool)) .await .map(|e| e.display_name); - let (authenticated, needs_setup) = self.check_tool_auth_status(&name).await; - let has_auth = self - .load_tool_capabilities(&name) - .await - .and_then(|c| c.auth) - .is_some(); + let auth_state = self.check_tool_auth_status(&name).await; extensions.push(InstalledExtension { name: name.clone(), kind: ExtensionKind::WasmTool, display_name, description: None, url: None, - authenticated, + authenticated: auth_state == ToolAuthState::Ready, active, tools: if active { vec![name] } else { Vec::new() }, - needs_setup, - has_auth, + needs_setup: auth_state == ToolAuthState::NeedsSetup, + has_auth: auth_state != ToolAuthState::NoAuth, installed: true, activation_error: None, }); @@ -454,8 +468,7 @@ impl ExtensionManager { let errors = self.activation_errors.read().await; for (name, _discovered) in channels { let active = active_names.contains(&name); - let (authenticated, needs_setup) = - self.check_channel_auth_status(&name).await; + let auth_state = self.check_channel_auth_status(&name).await; let activation_error = errors.get(&name).cloned(); let display_name = self .registry @@ -468,10 +481,10 @@ impl ExtensionManager { display_name, description: None, url: None, - authenticated, + authenticated: auth_state == ToolAuthState::Ready, active, tools: Vec::new(), - needs_setup, + needs_setup: auth_state == ToolAuthState::NeedsSetup, has_auth: false, installed: true, activation_error, @@ -1215,46 +1228,19 @@ impl ExtensionManager { .map_err(|e| ExtensionError::AuthFailed(e.to_string()))?; tracing::info!("MCP server '{}' authenticated via manual token", name); - return Ok(AuthResult { - name: name.to_string(), - kind: ExtensionKind::McpServer, - auth_url: None, - callback_type: None, - instructions: None, - setup_url: None, - awaiting_token: false, - status: "authenticated".to_string(), - }); + return Ok(AuthResult::authenticated(name, ExtensionKind::McpServer)); } // Check if already authenticated if is_authenticated(&server, &self.secrets, &self.user_id).await { - return Ok(AuthResult { - name: name.to_string(), - kind: ExtensionKind::McpServer, - auth_url: None, - callback_type: None, - instructions: None, - setup_url: None, - awaiting_token: false, - status: "authenticated".to_string(), - }); + return Ok(AuthResult::authenticated(name, ExtensionKind::McpServer)); } // Run the full 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 { - name: name.to_string(), - kind: ExtensionKind::McpServer, - auth_url: None, - callback_type: None, - instructions: None, - setup_url: None, - awaiting_token: false, - status: "authenticated".to_string(), - }) + Ok(AuthResult::authenticated(name, ExtensionKind::McpServer)) } Err(crate::tools::mcp::auth::AuthError::NotSupported) => { // Server doesn't support OAuth, try building a URL first @@ -1262,39 +1248,31 @@ impl ExtensionManager { Ok(result) => Ok(result), Err(_) => { // No OAuth, no DCR: fall back to manual token entry - Ok(AuthResult { - name: name.to_string(), - kind: ExtensionKind::McpServer, - auth_url: None, - callback_type: None, - instructions: Some(format!( + Ok(AuthResult::awaiting_token( + name, + ExtensionKind::McpServer, + format!( "Server '{}' does not support OAuth. \ Please provide an API token/key for this server.", name - )), - setup_url: None, - awaiting_token: true, - status: "awaiting_token".to_string(), - }) + ), + None, + )) } } } Err(e) => { // OAuth failed for some other reason, fall back to manual token - Ok(AuthResult { - name: name.to_string(), - kind: ExtensionKind::McpServer, - auth_url: None, - callback_type: None, - instructions: Some(format!( + Ok(AuthResult::awaiting_token( + name, + ExtensionKind::McpServer, + format!( "OAuth failed for '{}': {}. \ Please provide an API token/key manually.", name, e - )), - setup_url: None, - awaiting_token: true, - status: "awaiting_token".to_string(), - }) + ), + None, + )) } } } @@ -1356,16 +1334,12 @@ impl ExtensionManager { }, ); - Ok(AuthResult { - name: name.to_string(), - kind: ExtensionKind::McpServer, - auth_url: Some(auth_url), - callback_type: Some("local".to_string()), - instructions: None, - setup_url: None, - awaiting_token: false, - status: "awaiting_authorization".to_string(), - }) + Ok(AuthResult::awaiting_authorization( + name, + ExtensionKind::McpServer, + auth_url, + "local".to_string(), + )) } async fn auth_wasm_tool( @@ -1379,17 +1353,7 @@ impl ExtensionManager { .join(format!("{}.capabilities.json", name)); if !cap_path.exists() { - // No capabilities = no auth needed - return Ok(AuthResult { - name: name.to_string(), - kind: ExtensionKind::WasmTool, - auth_url: None, - callback_type: None, - instructions: None, - setup_url: None, - awaiting_token: false, - status: "no_auth_required".to_string(), - }); + return Ok(AuthResult::no_auth_required(name, ExtensionKind::WasmTool)); } let cap_bytes = tokio::fs::read(&cap_path) @@ -1402,16 +1366,7 @@ impl ExtensionManager { let auth = match cap_file.auth { Some(auth) => auth, None => { - return Ok(AuthResult { - name: name.to_string(), - kind: ExtensionKind::WasmTool, - auth_url: None, - callback_type: None, - instructions: None, - setup_url: None, - awaiting_token: false, - status: "no_auth_required".to_string(), - }); + return Ok(AuthResult::no_auth_required(name, ExtensionKind::WasmTool)); } }; @@ -1427,16 +1382,7 @@ impl ExtensionManager { .await .map_err(|e| ExtensionError::AuthFailed(e.to_string()))?; - return Ok(AuthResult { - name: name.to_string(), - kind: ExtensionKind::WasmTool, - auth_url: None, - callback_type: None, - instructions: None, - setup_url: None, - awaiting_token: false, - status: "authenticated".to_string(), - }); + return Ok(AuthResult::authenticated(name, ExtensionKind::WasmTool)); } // Check if already authenticated (with scope expansion detection) @@ -1466,16 +1412,7 @@ impl ExtensionManager { }; if !needs_reauth { - return Ok(AuthResult { - name: name.to_string(), - kind: ExtensionKind::WasmTool, - auth_url: None, - callback_type: None, - instructions: None, - setup_url: None, - awaiting_token: false, - status: "authenticated".to_string(), - }); + return Ok(AuthResult::authenticated(name, ExtensionKind::WasmTool)); } // Fall through to OAuth branch for scope expansion } @@ -1489,66 +1426,24 @@ impl ExtensionManager { .await .map_err(|e| ExtensionError::AuthFailed(e.to_string()))?; - return Ok(AuthResult { - name: name.to_string(), - kind: ExtensionKind::WasmTool, - auth_url: None, - callback_type: None, - instructions: None, - setup_url: None, - awaiting_token: false, - status: "authenticated".to_string(), - }); + return Ok(AuthResult::authenticated(name, ExtensionKind::WasmTool)); } // OAuth flow: if the tool has OAuth config, start the browser-based flow. // But only if credentials are available — if the tool has setup secrets // for client_id/secret that aren't configured yet, return needs_setup. if let Some(ref oauth) = auth.oauth { - let (setup_client_id_entry, setup_client_secret_entry) = - self.find_setup_credential_names(name).await; - - // Check all required (non-optional) setup credentials before starting - // OAuth, to avoid starting a flow that will fail during token exchange - // due to missing credentials. - let mut needs_setup = false; - if let Some((ref id_name, optional)) = setup_client_id_entry - && !optional - && !self - .secrets - .exists(&self.user_id, id_name) - .await - .unwrap_or(false) - { - needs_setup = true; - } - if !needs_setup - && let Some((ref secret_name, optional)) = setup_client_secret_entry - && !optional - && !self - .secrets - .exists(&self.user_id, secret_name) - .await - .unwrap_or(false) - { - needs_setup = true; - } - - if needs_setup { + if self.needs_setup_credentials(name, &auth, oauth).await { let display = auth.display_name.as_deref().unwrap_or(name); - return Ok(AuthResult { - name: name.to_string(), - kind: ExtensionKind::WasmTool, - auth_url: None, - callback_type: None, - instructions: Some(format!( + return Ok(AuthResult::needs_setup( + name, + ExtensionKind::WasmTool, + format!( "Configure OAuth credentials for {} in the Setup tab.", display - )), - setup_url: auth.setup_url.clone(), - awaiting_token: false, - status: "needs_setup".to_string(), - }); + ), + auth.setup_url.clone(), + )); } return self @@ -1563,54 +1458,51 @@ impl ExtensionManager { .instructions .unwrap_or_else(|| format!("Please provide your {} API token/key.", display)); - Ok(AuthResult { - name: name.to_string(), - kind: ExtensionKind::WasmTool, - auth_url: None, - callback_type: None, - instructions: Some(instructions), - setup_url: auth.setup_url, - awaiting_token: true, - status: "awaiting_token".to_string(), - }) + Ok(AuthResult::awaiting_token( + name, + ExtensionKind::WasmTool, + instructions, + auth.setup_url, + )) } - /// Check whether a WASM channel has all required secrets stored. - /// Returns `(authenticated, needs_setup)`. - async fn check_channel_auth_status(&self, name: &str) -> (bool, bool) { + /// Determine the auth readiness of a WASM channel. + async fn check_channel_auth_status(&self, name: &str) -> ToolAuthState { let cap_path = self .wasm_channels_dir .join(format!("{}.capabilities.json", name)); - if !cap_path.exists() { - return (true, false); - } let Ok(cap_bytes) = tokio::fs::read(&cap_path).await else { - return (true, false); + return ToolAuthState::NoAuth; }; let Ok(cap_file) = crate::channels::wasm::ChannelCapabilitiesFile::from_bytes(&cap_bytes) else { - return (true, false); + return ToolAuthState::NoAuth; }; - let required = &cap_file.setup.required_secrets; + + let required: Vec<_> = cap_file + .setup + .required_secrets + .iter() + .filter(|s| !s.optional) + .collect(); if required.is_empty() { - return (true, false); + return ToolAuthState::NoAuth; } - let mut all_provided = true; - for secret in required { - if secret.optional { - continue; - } - if !self - .secrets - .exists(&self.user_id, &secret.name) - .await - .unwrap_or(false) - { - all_provided = false; - break; - } + + let all_provided = futures::future::join_all( + required + .iter() + .map(|s| self.secrets.exists(&self.user_id, &s.name)), + ) + .await + .into_iter() + .all(|r| r.unwrap_or(false)); + + if all_provided { + ToolAuthState::Ready + } else { + ToolAuthState::NeedsSetup } - (all_provided, true) } /// Load and parse a WASM tool's capabilities file. @@ -1722,6 +1614,50 @@ impl ExtensionManager { (client_id_entry, client_secret_entry) } + /// Check if OAuth client credentials (client_id / client_secret) require + /// user input via the Setup tab. Returns `true` when at least one required + /// credential cannot be resolved through the full chain: + /// secrets store → inline → env var → builtin. + async fn needs_setup_credentials( + &self, + name: &str, + auth: &crate::tools::wasm::AuthCapabilitySchema, + oauth: &crate::tools::wasm::OAuthConfigSchema, + ) -> bool { + let builtin = crate::cli::oauth_defaults::builtin_credentials(&auth.secret_name); + let (id_entry, secret_entry) = self.find_setup_credential_names(name).await; + + for (entry, inline, env, fallback) in [ + ( + &id_entry, + &oauth.client_id, + &oauth.client_id_env, + builtin.as_ref().map(|c| c.client_id), + ), + ( + &secret_entry, + &oauth.client_secret, + &oauth.client_secret_env, + builtin.as_ref().map(|c| c.client_secret), + ), + ] { + let Some((ref setup_name, optional)) = *entry else { + continue; + }; + if optional { + continue; + } + let resolved = self + .resolve_oauth_credential(inline, env, fallback, Some(setup_name)) + .await + .is_some(); + if !resolved { + return true; + } + } + false + } + /// Resolve an OAuth credential value via: secrets store → inline → env var → builtin. /// /// For web gateway users, the secrets store is checked first because client_id/secret @@ -1819,7 +1755,7 @@ impl ExtensionManager { ) .await; - // Cancel any existing pending auth for this tool (frees port 9876) + // Cancel any existing pending auth for this tool (frees port 9876 in TCP mode) { let mut pending = self.pending_auth.write().await; if let Some(old) = pending.remove(name) @@ -1828,11 +1764,11 @@ impl ExtensionManager { handle.abort(); } } - - // Bind callback listener - let listener = oauth_defaults::bind_callback_listener() - .await - .map_err(|e| format!("Failed to start OAuth callback listener: {}", e))?; + // Also clean up any gateway-mode pending flows for this tool + { + let mut flows = self.pending_oauth_flows.write().await; + flows.retain(|_, flow| flow.extension_name != name); + } let redirect_uri = format!("{}/callback", oauth_defaults::callback_url()); @@ -1854,155 +1790,290 @@ impl ExtensionManager { let code_verifier = oauth_result.code_verifier; let expected_state = oauth_result.state; - // Spawn background task: wait for callback → exchange code → validate → store tokens let display_name = auth .display_name .clone() .unwrap_or_else(|| name.to_string()); - let token_url = oauth.token_url.clone(); - let access_token_field = oauth.access_token_field.clone(); - let secret_name = auth.secret_name.clone(); - let provider = auth.provider.clone(); - let validation_endpoint = auth.validation_endpoint.clone(); - let user_id = self.user_id.clone(); - let secrets = Arc::clone(&self.secrets); - let sse_sender = self.sse_sender.read().await.clone(); - let ext_name = name.to_string(); - let task_handle = tokio::spawn(async move { - let result: Result<(), String> = async { - let code = oauth_defaults::wait_for_callback( - listener, - "/callback", - "code", - &display_name, - Some(&expected_state), + if oauth_defaults::use_gateway_callback() { + // 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. + oauth_defaults::sweep_expired_flows(&self.pending_oauth_flows).await; + + // Wrap the CSRF nonce with instance name for platform routing. + // Nginx at auth.DOMAIN parses `instance:nonce` to route the callback + // to the correct container. The flow is keyed by the raw nonce. + let platform_state = oauth_defaults::build_platform_state(&expected_state); + let auth_url = if platform_state != expected_state { + auth_url.replace( + &format!("state={}", urlencoding::encode(&expected_state)), + &format!("state={}", urlencoding::encode(&platform_state)), ) - .await - .map_err(|e| e.to_string())?; - - let token_response = oauth_defaults::exchange_oauth_code( - &token_url, - &client_id, - client_secret.as_deref(), - &code, - &redirect_uri, - code_verifier.as_deref(), - &access_token_field, - ) - .await - .map_err(|e| e.to_string())?; - - // Validate the token before storing (catches wrong account, etc.) - if let Some(ref validation) = validation_endpoint { - oauth_defaults::validate_oauth_token(&token_response.access_token, validation) - .await - .map_err(|e| e.to_string())?; - } - - oauth_defaults::store_oauth_tokens( - secrets.as_ref(), - &user_id, - &secret_name, - provider.as_deref(), - &token_response.access_token, - token_response.refresh_token.as_deref(), - token_response.expires_in, - &merged_scopes, - ) - .await - .map_err(|e| e.to_string())?; - - Ok(()) - } - .await; - - // Broadcast SSE event - let (success, message) = match result { - Ok(()) => (true, format!("{} authenticated successfully", display_name)), - Err(ref e) => ( - false, - format!("{} authentication failed: {}", display_name, e), - ), + } else { + auth_url }; - match &result { - Ok(()) => { - tracing::info!( - tool = %ext_name, - "OAuth completed successfully" - ); - } - Err(e) => { - tracing::warn!( - tool = %ext_name, - error = %e, - "WASM tool OAuth failed" - ); - } - } - - if let Some(ref sender) = sse_sender { - let _ = sender.send(crate::channels::web::types::SseEvent::AuthCompleted { - extension_name: ext_name, - success, - message, - }); - } - }); - - // Store pending auth with task handle - self.pending_auth.write().await.insert( - name.to_string(), - PendingAuth { - _name: name.to_string(), - _kind: ExtensionKind::WasmTool, + let flow = oauth_defaults::PendingOAuthFlow { + extension_name: name.to_string(), + display_name: display_name.clone(), + token_url: oauth.token_url.clone(), + client_id: client_id.clone(), + client_secret: client_secret.clone(), + redirect_uri: redirect_uri.clone(), + code_verifier, + access_token_field: oauth.access_token_field.clone(), + secret_name: auth.secret_name.clone(), + provider: auth.provider.clone(), + validation_endpoint: auth.validation_endpoint.clone(), + scopes: merged_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(), created_at: std::time::Instant::now(), - task_handle: Some(task_handle), - }, - ); + }; - Ok(AuthResult { - name: name.to_string(), - kind: ExtensionKind::WasmTool, - auth_url: Some(auth_url), - callback_type: Some("local".to_string()), - instructions: None, - setup_url: None, - awaiting_token: false, - status: "awaiting_authorization".to_string(), - }) + // Key by raw nonce (without instance prefix) — the callback handler + // strips the prefix before lookup. + self.pending_oauth_flows + .write() + .await + .insert(expected_state, flow); + + // Register pending auth without a task handle (gateway handles completion) + self.pending_auth.write().await.insert( + name.to_string(), + PendingAuth { + _name: name.to_string(), + _kind: ExtensionKind::WasmTool, + created_at: std::time::Instant::now(), + task_handle: None, + }, + ); + + Ok(AuthResult::awaiting_authorization( + name, + ExtensionKind::WasmTool, + auth_url, + "gateway".to_string(), + )) + } else { + // TCP listener mode: bind port 9876 and spawn a background task + // to wait for the callback. This is the original flow for local/desktop use. + let listener = oauth_defaults::bind_callback_listener() + .await + .map_err(|e| format!("Failed to start OAuth callback listener: {}", e))?; + + let token_url = oauth.token_url.clone(); + let access_token_field = oauth.access_token_field.clone(); + let secret_name = auth.secret_name.clone(); + let provider = auth.provider.clone(); + let validation_endpoint = auth.validation_endpoint.clone(); + let user_id = self.user_id.clone(); + let secrets = Arc::clone(&self.secrets); + let sse_sender = self.sse_sender.read().await.clone(); + let ext_name = name.to_string(); + + let task_handle = tokio::spawn(async move { + let result: Result<(), String> = async { + let code = oauth_defaults::wait_for_callback( + listener, + "/callback", + "code", + &display_name, + Some(&expected_state), + ) + .await + .map_err(|e| e.to_string())?; + + let token_response = oauth_defaults::exchange_oauth_code( + &token_url, + &client_id, + client_secret.as_deref(), + &code, + &redirect_uri, + code_verifier.as_deref(), + &access_token_field, + ) + .await + .map_err(|e| e.to_string())?; + + // Validate the token before storing (catches wrong account, etc.) + if let Some(ref validation) = validation_endpoint { + oauth_defaults::validate_oauth_token( + &token_response.access_token, + validation, + ) + .await + .map_err(|e| e.to_string())?; + } + + oauth_defaults::store_oauth_tokens( + secrets.as_ref(), + &user_id, + &secret_name, + provider.as_deref(), + &token_response.access_token, + token_response.refresh_token.as_deref(), + token_response.expires_in, + &merged_scopes, + ) + .await + .map_err(|e| e.to_string())?; + + Ok(()) + } + .await; + + // Broadcast SSE event + let (success, message) = match result { + Ok(()) => (true, format!("{} authenticated successfully", display_name)), + Err(ref e) => ( + false, + format!("{} authentication failed: {}", display_name, e), + ), + }; + + match &result { + Ok(()) => { + tracing::info!( + tool = %ext_name, + "OAuth completed successfully" + ); + } + Err(e) => { + tracing::warn!( + tool = %ext_name, + error = %e, + "WASM tool OAuth failed" + ); + } + } + + if let Some(ref sender) = sse_sender { + let _ = sender.send(crate::channels::web::types::SseEvent::AuthCompleted { + extension_name: ext_name, + success, + message, + }); + } + }); + + // Store pending auth with task handle + self.pending_auth.write().await.insert( + name.to_string(), + PendingAuth { + _name: name.to_string(), + _kind: ExtensionKind::WasmTool, + created_at: std::time::Instant::now(), + task_handle: Some(task_handle), + }, + ); + + Ok(AuthResult::awaiting_authorization( + name, + ExtensionKind::WasmTool, + auth_url, + "local".to_string(), + )) + } } - /// Check whether a WASM tool's required setup secrets are provided. + /// Returns `true` if a setup secret is an OAuth credential (client_id or client_secret) + /// that can be resolved without user input — via inline capabilities, env var, or + /// builtin defaults. /// - /// Returns `(authenticated, needs_setup)` — same semantics as `check_channel_auth_status`. - async fn check_tool_auth_status(&self, name: &str) -> (bool, bool) { - let Some(cap_file) = self.load_tool_capabilities(name).await else { - return (true, false); - }; - let Some(setup) = &cap_file.setup else { - return (true, false); - }; - if setup.required_secrets.is_empty() { - return (true, false); + /// Used by `check_tool_auth_status()` and `get_setup_schema()` to hide setup fields + /// that the user doesn't need to fill (e.g., Google tools with builtin credentials). + fn is_auto_resolved_oauth_field( + secret_name: &str, + cap_file: &crate::tools::wasm::CapabilitiesFile, + ) -> bool { + let lower = secret_name.to_lowercase(); + let is_client_id = lower.ends_with("client_id") || lower == "client_id"; + let is_client_secret = lower.ends_with("client_secret") || lower == "client_secret"; + if !is_client_id && !is_client_secret { + return false; } - let mut all_provided = true; - for secret in &setup.required_secrets { - if secret.optional { - continue; - } - if !self + let Some(ref auth) = cap_file.auth else { + return false; + }; + let Some(ref oauth) = auth.oauth else { + return false; + }; + let builtin = crate::cli::oauth_defaults::builtin_credentials(&auth.secret_name); + + if is_client_id { + oauth.client_id.is_some() + || oauth + .client_id_env + .as_ref() + .is_some_and(|e| std::env::var(e).is_ok()) + || builtin.is_some() + } else { + oauth.client_secret.is_some() + || oauth + .client_secret_env + .as_ref() + .is_some_and(|e| std::env::var(e).is_ok()) + || builtin.is_some() + } + } + + /// Determine the auth readiness of a WASM tool. + async fn check_tool_auth_status(&self, name: &str) -> ToolAuthState { + let Some(cap_file) = self.load_tool_capabilities(name).await else { + return ToolAuthState::NoAuth; + }; + + // If the tool declares an auth section, the access token is the + // authoritative signal — setup secrets (client_id/secret) are + // intermediate and may be auto-resolved via builtins. + if let Some(ref auth) = cap_file.auth { + let has_token = self .secrets - .exists(&self.user_id, &secret.name) + .exists(&self.user_id, &auth.secret_name) .await .unwrap_or(false) - { - all_provided = false; - break; - } + || auth + .env_var + .as_ref() + .is_some_and(|v| std::env::var(v).is_ok()); + return if has_token { + ToolAuthState::Ready + } else if auth.oauth.is_some() { + ToolAuthState::NeedsAuth + } else { + ToolAuthState::NeedsSetup + }; + } + + // No auth section — fall back to checking setup.required_secrets. + let Some(setup) = &cap_file.setup else { + return ToolAuthState::NoAuth; + }; + if setup.required_secrets.is_empty() { + return ToolAuthState::NoAuth; + } + + let all_provided = futures::future::join_all( + setup + .required_secrets + .iter() + .filter(|s| !s.optional) + .filter(|s| !Self::is_auto_resolved_oauth_field(&s.name, &cap_file)) + .map(|s| self.secrets.exists(&self.user_id, &s.name)), + ) + .await + .into_iter() + .all(|r| r.unwrap_or(false)); + + if all_provided { + ToolAuthState::Ready + } else { + ToolAuthState::NeedsSetup } - (all_provided, true) } async fn auth_wasm_channel( @@ -2015,16 +2086,10 @@ impl ExtensionManager { .join(format!("{}.capabilities.json", name)); if !cap_path.exists() { - return Ok(AuthResult { - name: name.to_string(), - kind: ExtensionKind::WasmChannel, - auth_url: None, - callback_type: None, - instructions: None, - setup_url: None, - awaiting_token: false, - status: "no_auth_required".to_string(), - }); + return Ok(AuthResult::no_auth_required( + name, + ExtensionKind::WasmChannel, + )); } let cap_bytes = tokio::fs::read(&cap_path) @@ -2037,16 +2102,10 @@ impl ExtensionManager { // Get required secrets from the setup section let required_secrets = &cap_file.setup.required_secrets; if required_secrets.is_empty() { - return Ok(AuthResult { - name: name.to_string(), - kind: ExtensionKind::WasmChannel, - auth_url: None, - callback_type: None, - instructions: None, - setup_url: None, - awaiting_token: false, - status: "no_auth_required".to_string(), - }); + return Ok(AuthResult::no_auth_required( + name, + ExtensionKind::WasmChannel, + )); } // Find the first non-optional secret that isn't yet stored @@ -2066,16 +2125,7 @@ impl ExtensionManager { } if missing.is_empty() { - return Ok(AuthResult { - name: name.to_string(), - kind: ExtensionKind::WasmChannel, - auth_url: None, - callback_type: None, - instructions: None, - setup_url: None, - awaiting_token: false, - status: "authenticated".to_string(), - }); + return Ok(AuthResult::authenticated(name, ExtensionKind::WasmChannel)); } // If a token was provided, store it for the first missing secret @@ -2090,44 +2140,27 @@ impl ExtensionManager { // Check if there are more missing secrets if missing.len() <= 1 { - return Ok(AuthResult { - name: name.to_string(), - kind: ExtensionKind::WasmChannel, - auth_url: None, - callback_type: None, - instructions: None, - setup_url: None, - awaiting_token: false, - status: "authenticated".to_string(), - }); + return Ok(AuthResult::authenticated(name, ExtensionKind::WasmChannel)); } // More secrets needed; prompt for the next one let next = &missing[1]; - return Ok(AuthResult { - name: name.to_string(), - kind: ExtensionKind::WasmChannel, - auth_url: None, - callback_type: None, - instructions: Some(next.prompt.clone()), - setup_url: cap_file.setup.setup_url.clone(), - awaiting_token: true, - status: "awaiting_token".to_string(), - }); + return Ok(AuthResult::awaiting_token( + name, + ExtensionKind::WasmChannel, + next.prompt.clone(), + cap_file.setup.setup_url.clone(), + )); } // Prompt for the first missing secret let secret = &missing[0]; - Ok(AuthResult { - name: name.to_string(), - kind: ExtensionKind::WasmChannel, - auth_url: None, - callback_type: None, - instructions: Some(secret.prompt.clone()), - setup_url: cap_file.setup.setup_url.clone(), - awaiting_token: true, - status: "awaiting_token".to_string(), - }) + Ok(AuthResult::awaiting_token( + name, + ExtensionKind::WasmChannel, + secret.prompt.clone(), + cap_file.setup.setup_url.clone(), + )) } async fn activate_mcp(&self, name: &str) -> Result { @@ -2328,8 +2361,8 @@ impl ExtensionManager { }; // Check auth status first - let (authenticated, _needs_setup) = self.check_channel_auth_status(name).await; - if !authenticated { + let auth_state = self.check_channel_auth_status(name).await; + if auth_state != ToolAuthState::Ready && auth_state != ToolAuthState::NoAuth { return Err(ExtensionError::ActivationFailed(format!( "Channel '{}' requires configuration. Use the setup form to provide credentials.", name @@ -2759,6 +2792,10 @@ impl ExtensionManager { let mut fields = Vec::new(); if let Some(setup) = &cap_file.setup { for secret in &setup.required_secrets { + // Skip OAuth client_id/secret fields that resolve automatically + if Self::is_auto_resolved_oauth_field(&secret.name, &cap_file) { + continue; + } let provided = self .secrets .exists(&self.user_id, &secret.name) @@ -2959,7 +2996,7 @@ impl ExtensionManager { // This is safe to call here — cancel-and-retry prevents port conflicts. let mut auth_url = None; if let Ok(auth_result) = self.auth(name, None).await { - auth_url = auth_result.auth_url; + auth_url = auth_result.auth_url().map(String::from); } let message = if auth_url.is_some() { format!( diff --git a/src/extensions/mod.rs b/src/extensions/mod.rs index 353b6ff9..51a173f7 100644 --- a/src/extensions/mod.rs +++ b/src/extensions/mod.rs @@ -24,6 +24,7 @@ pub use discovery::OnlineDiscovery; pub use manager::ExtensionManager; pub use registry::ExtensionRegistry; +use serde::ser::SerializeMap; use serde::{Deserialize, Serialize}; /// The kind of extension, determining how it's installed, authenticated, and activated. @@ -145,28 +146,267 @@ pub struct InstallResult { pub message: String, } +/// Auth readiness state for the extensions list UI. +/// +/// Used by `check_tool_auth_status` and `check_channel_auth_status` to +/// communicate a tool's credential state to the list handler without +/// ambiguous `(bool, bool)` tuples. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ToolAuthState { + /// Token/credentials are present — ready to use. + Ready, + /// Auth section exists but the access token is missing (OAuth not completed). + NeedsAuth, + /// Setup credentials (client_id/secret) must be configured before OAuth can start. + NeedsSetup, + /// No auth configuration at all (no capabilities or auth section). + NoAuth, +} + +/// The typed auth status, carrying only the data relevant to each state. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum AuthStatus { + /// Authentication is complete; no further action needed. + Authenticated, + /// No authentication is required for this extension. + NoAuthRequired, + /// OAuth flow started — user must open `auth_url` in their browser. + AwaitingAuthorization { + auth_url: String, + callback_type: String, + }, + /// Waiting for user to provide a token/key manually. + AwaitingToken { + instructions: String, + setup_url: Option, + }, + /// OAuth client credentials need to be configured before auth can proceed. + NeedsSetup { + instructions: String, + setup_url: Option, + }, +} + +impl AuthStatus { + /// The wire-format status string (backward-compatible with JS consumers). + pub fn as_str(&self) -> &'static str { + match self { + AuthStatus::Authenticated => "authenticated", + AuthStatus::NoAuthRequired => "no_auth_required", + AuthStatus::AwaitingAuthorization { .. } => "awaiting_authorization", + AuthStatus::AwaitingToken { .. } => "awaiting_token", + AuthStatus::NeedsSetup { .. } => "needs_setup", + } + } +} + /// Result of authenticating an extension. -#[derive(Debug, Clone, Serialize, Deserialize)] +#[derive(Debug, Clone)] pub struct AuthResult { pub name: String, pub kind: ExtensionKind, - /// OAuth URL to open (for OAuth flows). - #[serde(skip_serializing_if = "Option::is_none")] - pub auth_url: Option, - /// Whether using local or remote callback. - #[serde(skip_serializing_if = "Option::is_none")] - pub callback_type: Option, - /// Instructions for manual token entry (for WASM tools). - #[serde(skip_serializing_if = "Option::is_none")] - pub instructions: Option, - /// URL for manual token setup. - #[serde(skip_serializing_if = "Option::is_none")] - pub setup_url: Option, - /// Whether the tool is waiting for a token from the user. - #[serde(default)] - pub awaiting_token: bool, - /// Current auth status. - pub status: String, + pub status: AuthStatus, +} + +impl AuthResult { + // ── Constructors ────────────────────────────────────────────────── + + pub fn authenticated(name: impl Into, kind: ExtensionKind) -> Self { + Self { + name: name.into(), + kind, + status: AuthStatus::Authenticated, + } + } + + pub fn no_auth_required(name: impl Into, kind: ExtensionKind) -> Self { + Self { + name: name.into(), + kind, + status: AuthStatus::NoAuthRequired, + } + } + + pub fn awaiting_authorization( + name: impl Into, + kind: ExtensionKind, + auth_url: String, + callback_type: String, + ) -> Self { + Self { + name: name.into(), + kind, + status: AuthStatus::AwaitingAuthorization { + auth_url, + callback_type, + }, + } + } + + pub fn awaiting_token( + name: impl Into, + kind: ExtensionKind, + instructions: String, + setup_url: Option, + ) -> Self { + Self { + name: name.into(), + kind, + status: AuthStatus::AwaitingToken { + instructions, + setup_url, + }, + } + } + + pub fn needs_setup( + name: impl Into, + kind: ExtensionKind, + instructions: String, + setup_url: Option, + ) -> Self { + Self { + name: name.into(), + kind, + status: AuthStatus::NeedsSetup { + instructions, + setup_url, + }, + } + } + + // ── Accessors ───────────────────────────────────────────────────── + + pub fn is_authenticated(&self) -> bool { + matches!(self.status, AuthStatus::Authenticated) + } + + pub fn auth_url(&self) -> Option<&str> { + match &self.status { + AuthStatus::AwaitingAuthorization { auth_url, .. } => Some(auth_url), + _ => None, + } + } + + pub fn callback_type(&self) -> Option<&str> { + match &self.status { + AuthStatus::AwaitingAuthorization { callback_type, .. } => Some(callback_type), + _ => None, + } + } + + pub fn instructions(&self) -> Option<&str> { + match &self.status { + AuthStatus::AwaitingToken { instructions, .. } + | AuthStatus::NeedsSetup { instructions, .. } => Some(instructions), + _ => None, + } + } + + pub fn setup_url(&self) -> Option<&str> { + match &self.status { + AuthStatus::AwaitingToken { setup_url, .. } + | AuthStatus::NeedsSetup { setup_url, .. } => setup_url.as_deref(), + _ => None, + } + } + + pub fn is_awaiting_token(&self) -> bool { + matches!(self.status, AuthStatus::AwaitingToken { .. }) + } + + pub fn status_str(&self) -> &'static str { + self.status.as_str() + } +} + +/// Serialize `AuthResult` to the same flat JSON shape the JS frontend expects. +impl Serialize for AuthResult { + fn serialize(&self, serializer: S) -> Result { + // Count fields: name + kind + status + optional fields + let optional_count = self.auth_url().is_some() as usize + + self.callback_type().is_some() as usize + + self.instructions().is_some() as usize + + self.setup_url().is_some() as usize; + let mut map = serializer.serialize_map(Some(4 + optional_count))?; + + map.serialize_entry("name", &self.name)?; + map.serialize_entry("kind", &self.kind)?; + if let Some(url) = self.auth_url() { + map.serialize_entry("auth_url", url)?; + } + if let Some(cb) = self.callback_type() { + map.serialize_entry("callback_type", cb)?; + } + if let Some(inst) = self.instructions() { + map.serialize_entry("instructions", inst)?; + } + if let Some(url) = self.setup_url() { + map.serialize_entry("setup_url", url)?; + } + map.serialize_entry("awaiting_token", &self.is_awaiting_token())?; + map.serialize_entry("status", self.status_str())?; + map.end() + } +} + +/// Deserialize from the flat JSON shape back into the typed enum. +impl<'de> Deserialize<'de> for AuthResult { + fn deserialize>(deserializer: D) -> Result { + /// Flat helper matching the old JSON shape. + #[derive(Deserialize)] + #[allow(dead_code)] + struct Raw { + name: String, + kind: ExtensionKind, + #[serde(default)] + auth_url: Option, + #[serde(default)] + callback_type: Option, + #[serde(default)] + instructions: Option, + #[serde(default)] + setup_url: Option, + #[serde(default)] + awaiting_token: bool, + status: String, + } + + let raw = Raw::deserialize(deserializer)?; + let status = match raw.status.as_str() { + "authenticated" => AuthStatus::Authenticated, + "no_auth_required" => AuthStatus::NoAuthRequired, + "awaiting_authorization" => AuthStatus::AwaitingAuthorization { + auth_url: raw.auth_url.unwrap_or_default(), + callback_type: raw.callback_type.unwrap_or_default(), + }, + "awaiting_token" => AuthStatus::AwaitingToken { + instructions: raw.instructions.unwrap_or_default(), + setup_url: raw.setup_url, + }, + "needs_setup" => AuthStatus::NeedsSetup { + instructions: raw.instructions.unwrap_or_default(), + setup_url: raw.setup_url, + }, + other => { + return Err(serde::de::Error::unknown_variant( + other, + &[ + "authenticated", + "no_auth_required", + "awaiting_authorization", + "awaiting_token", + "needs_setup", + ], + )); + } + }; + Ok(AuthResult { + name: raw.name, + kind: raw.kind, + status, + }) + } } /// Result of activating an extension. @@ -257,3 +497,124 @@ pub enum ExtensionError { #[error("{0}")] Other(String), } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn auth_result_authenticated_round_trip() { + let result = AuthResult::authenticated("gmail", ExtensionKind::WasmTool); + let json = serde_json::to_value(&result).unwrap(); + + assert_eq!(json["status"], "authenticated"); + assert_eq!(json["name"], "gmail"); + assert_eq!(json["kind"], "wasm_tool"); + assert_eq!(json["awaiting_token"], false); + assert!(json.get("auth_url").is_none()); + assert!(json.get("instructions").is_none()); + + let back: AuthResult = serde_json::from_value(json).unwrap(); + assert!(back.is_authenticated()); + assert!(back.auth_url().is_none()); + } + + #[test] + fn auth_result_awaiting_authorization_round_trip() { + let result = AuthResult::awaiting_authorization( + "google-drive", + ExtensionKind::WasmTool, + "https://accounts.google.com/o/oauth2/v2/auth?state=abc".to_string(), + "local".to_string(), + ); + let json = serde_json::to_value(&result).unwrap(); + + assert_eq!(json["status"], "awaiting_authorization"); + assert_eq!( + json["auth_url"], + "https://accounts.google.com/o/oauth2/v2/auth?state=abc" + ); + assert_eq!(json["callback_type"], "local"); + assert_eq!(json["awaiting_token"], false); + + let back: AuthResult = serde_json::from_value(json).unwrap(); + assert_eq!( + back.auth_url(), + Some("https://accounts.google.com/o/oauth2/v2/auth?state=abc") + ); + assert_eq!(back.callback_type(), Some("local")); + assert!(!back.is_authenticated()); + } + + #[test] + fn auth_result_awaiting_token_round_trip() { + let result = AuthResult::awaiting_token( + "telegram", + ExtensionKind::WasmChannel, + "Enter your bot token".to_string(), + None, + ); + let json = serde_json::to_value(&result).unwrap(); + + assert_eq!(json["status"], "awaiting_token"); + assert_eq!(json["instructions"], "Enter your bot token"); + assert_eq!(json["awaiting_token"], true); + assert!(json.get("auth_url").is_none()); + + let back: AuthResult = serde_json::from_value(json).unwrap(); + assert!(back.is_awaiting_token()); + assert_eq!(back.instructions(), Some("Enter your bot token")); + } + + #[test] + fn auth_result_needs_setup_round_trip() { + let result = AuthResult::needs_setup( + "custom-tool", + ExtensionKind::WasmTool, + "Configure OAuth credentials in the Setup tab.".to_string(), + Some("https://console.cloud.google.com".to_string()), + ); + let json = serde_json::to_value(&result).unwrap(); + + assert_eq!(json["status"], "needs_setup"); + assert_eq!(json["setup_url"], "https://console.cloud.google.com"); + assert_eq!(json["awaiting_token"], false); + + let back: AuthResult = serde_json::from_value(json).unwrap(); + assert!(!back.is_authenticated()); + assert!(!back.is_awaiting_token()); + assert_eq!(back.setup_url(), Some("https://console.cloud.google.com")); + } + + #[test] + fn auth_result_no_auth_required_round_trip() { + let result = AuthResult::no_auth_required("echo", ExtensionKind::WasmTool); + let json = serde_json::to_value(&result).unwrap(); + + assert_eq!(json["status"], "no_auth_required"); + assert_eq!(json["awaiting_token"], false); + + let back: AuthResult = serde_json::from_value(json).unwrap(); + assert!(!back.is_authenticated()); + assert_eq!(back.status, AuthStatus::NoAuthRequired); + } + + #[test] + fn auth_status_type_safety() { + // AwaitingAuthorization always has auth_url + let result = AuthResult::awaiting_authorization( + "test", + ExtensionKind::WasmTool, + "https://example.com".to_string(), + "local".to_string(), + ); + assert!(result.auth_url().is_some()); + assert!(!result.is_awaiting_token()); + + // Authenticated never has auth_url + let result = AuthResult::authenticated("test", ExtensionKind::WasmTool); + assert!(result.auth_url().is_none()); + assert!(result.instructions().is_none()); + assert!(result.setup_url().is_none()); + } +} diff --git a/src/tools/builtin/extension_tools.rs b/src/tools/builtin/extension_tools.rs index 0fcfbda3..f82049e9 100644 --- a/src/tools/builtin/extension_tools.rs +++ b/src/tools/builtin/extension_tools.rs @@ -218,7 +218,7 @@ impl Tool for ToolAuthTool { .map_err(|e| ToolError::ExecutionFailed(e.to_string()))?; // Auto-activate after successful auth so tools are available immediately - if result.status == "authenticated" { + if result.is_authenticated() { match self.manager.activate(name).await { Ok(activate_result) => { let output = serde_json::json!({ @@ -324,7 +324,7 @@ impl Tool for ToolActivateTool { // Activation failed due to missing auth; initiate auth flow // so the agent loop can show the auth card. match self.manager.auth(name, None).await { - Ok(auth_result) if auth_result.status == "authenticated" => { + Ok(auth_result) if auth_result.is_authenticated() => { // Auth succeeded (e.g. env var was set); retry activation. let result = self .manager diff --git a/src/tools/wasm/wrapper.rs b/src/tools/wasm/wrapper.rs index 1f545f77..4328eb9e 100644 --- a/src/tools/wasm/wrapper.rs +++ b/src/tools/wasm/wrapper.rs @@ -656,10 +656,17 @@ impl Tool for WasmToolWrapper { // Pre-resolve host credentials from secrets store (async, before blocking task). // This decrypts the secrets once so the sync http_request() host function // can inject them without needing async access. + // + // BUG FIX: ExtensionManager stores OAuth tokens under user_id "default" + // (hardcoded at construction in app.rs), but this was previously looking + // them up under ctx.user_id — which could be a Telegram user ID, web + // gateway user, etc. — causing credential resolution to silently fail. + // Must match the storage key until per-user credential isolation is added. + let credential_user_id = "default"; let host_credentials = resolve_host_credentials( &self.capabilities, self.secrets_store.as_deref(), - &ctx.user_id, + credential_user_id, self.oauth_refresh.as_ref(), ) .await;