gitops/optimclaw
1
0
Fork 0
mirror of https://github.com/outbackdingo/optimclaw.git synced 2026-08-25 14:53:34 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: add OAuth support for WASM tools in web gateway (#489)

Browse Source 
* feat: add OAuth support for WASM tools in web gateway

Extract reusable OAuth functions (build_oauth_url, exchange_oauth_code,
store_oauth_tokens, validate_oauth_token) from CLI into shared
oauth_defaults module, then wire them into the web gateway's
ExtensionManager.

Key changes:
- Install auto-activates WASM tools (no separate Activate button)
- Configure button triggers OAuth flow via save_setup_secrets
- Scope merging: installing a second Google tool triggers re-auth with
  merged scopes from all tools sharing the same secret_name
- Cancel-and-retry: aborting stale OAuth listeners prevents port conflicts
- Post-auth validation: wrong account detected via validation_endpoint
- Reconfigure always re-auths (deletes old token before starting fresh)
- UI shows error toast on OAuth failure, refreshes extension list

Flow: Install → Active → Configure (enter client_id/secret) → Save →
OAuth popup → authorize → done. Second Google tool install auto-triggers
scope expansion OAuth.

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

* fix: address PR review comments

- Add custom headers support to ValidationEndpointSchema (fixes
  missing Notion-Version header regression)
- Guard activate handler auth check with status == "awaiting_authorization"
  to prevent unexpected OAuth popups
- Add window dimensions to OAuth popup in activateExtension()
- Simplify UTF-8 truncation boundary check

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

* fix: address Copilot PR review comments (security, UX, bugs)

- Add CSRF state parameter to OAuth flow (random state in auth URL, validated in callback)
- Restore MCP server Activate button in web UI (was hidden for all non-channel extensions)
- Abort JoinHandle in cleanup_expired_auths to prevent port 9876 conflicts
- Fix Google-specific error message for non-Google OAuth providers
- Add has_auth field to ExtensionInfo API response (fixes Configure button visibility)
- Use oauth_defaults::callback_url() instead of hardcoded redirect_uri (both CLI and manager)
- Update auth check comment to match actual behavior (scope expansion + first-time auth)
- Add unit tests for build_oauth_url (basic, PKCE, extra params, state uniqueness)
- Check all required setup secrets (client_id + client_secret) before starting OAuth

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

---------

Co-authored-by: Claude Opus 4.6 (1M context) <[email protected]>
This commit is contained in:
Henry Park
2026-03-04 01:08:40 +08:00
committed by GitHub
co-authored by Claude Opus 4.6
parent f4855962fc
commit 18b59ae9a7
11 changed files with 1109 additions and 232 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/channels/web/handlers/extensions.rs
+6 -1
View File
@@ -59,6 +59,7 @@ pub async fn extensions_list_handler(
active: ext.active,
tools: ext.tools,
needs_setup: ext.needs_setup,
has_auth: ext.has_auth,
activation_status,
activation_error: ext.activation_error,
}
@@ -123,7 +124,11 @@ pub async fn extensions_activate_handler(
))?;
match ext_mgr.activate(&name).await {
Ok(result) => Ok(Json(ActionResponse::ok(result.message))),
Ok(result) => {
// Activation just loads the WASM module. Auth (OAuth/manual) is
// triggered separately via save_setup_secrets or the auth endpoint.
Ok(Json(ActionResponse::ok(result.message)))
}
Err(activate_err) => {
let err_str = activate_err.to_string();
let needs_auth = err_str.contains("authentication")
src/channels/web/server.rs
+47 -2
View File
@@ -1244,6 +1244,7 @@ async fn extensions_list_handler(
active: ext.active,
tools: ext.tools,
needs_setup: ext.needs_setup,
has_auth: ext.has_auth,
activation_status,
activation_error: ext.activation_error,
}
@@ -1313,7 +1314,37 @@ async fn extensions_install_handler(
.install(&req.name, req.url.as_deref(), kind_hint)
.await
{
Ok(result) => Ok(Json(ActionResponse::ok(result.message))),
Ok(result) => {
let mut resp = ActionResponse::ok(result.message);
// Auto-activate WASM tools after install (install = active).
if result.kind == crate::extensions::ExtensionKind::WasmTool {
if let Err(e) = ext_mgr.activate(&req.name).await {
tracing::debug!(
extension = %req.name,
error = %e,
"Auto-activation after install failed"
);
}
// Check auth after activation. This may initiate OAuth both for scope
// expansion and for first-time auth when credentials are already
// 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" =>
{
// Scope expansion or initial OAuth: user needs to authorize
resp.auth_url = auth_result.auth_url;
}
_ => {}
}
}
Ok(Json(resp))
}
Err(e) => Ok(Json(ActionResponse::fail(e.to_string()))),
}
}
@@ -1328,7 +1359,20 @@ async fn extensions_activate_handler(
))?;
match ext_mgr.activate(&name).await {
Ok(result) => Ok(Json(ActionResponse::ok(result.message))),
Ok(result) => {
// Activation loaded the WASM module. Check if the tool needs
// OAuth scope expansion (e.g., adding google-docs when gmail
// already has a token but missing the documents scope).
// 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"
{
resp.auth_url = auth_result.auth_url;
}
Ok(Json(resp))
}
Err(activate_err) => {
let err_str = activate_err.to_string();
let needs_auth = err_str.contains("authentication")
@@ -1550,6 +1594,7 @@ async fn extensions_setup_submit_handler(
Ok(result) => {
let mut resp = ActionResponse::ok(result.message);
resp.activated = Some(result.activated);
resp.auth_url = result.auth_url;
Ok(Json(resp))
}
Err(e) => Ok(Json(ActionResponse::fail(e.to_string()))),
src/channels/web/static/app.js
+32 -11
View File
@@ -228,7 +228,13 @@ function connectSSE() {
eventSource.addEventListener('auth_completed', (e) => {
const data = JSON.parse(e.data);
removeAuthCard(data.extension_name);
showToast(data.message, 'success');
if (data.success) {
showToast(data.message, 'success');
} else {
showToast(data.message, 'error');
}
// Refresh extensions list so status indicators update
if (currentTab === 'extensions') loadExtensions();
enableChatInput();
});
@@ -1760,6 +1766,11 @@ function renderAvailableExtensionCard(entry) {
}).then(function(res) {
if (res.success) {
showToast('Installed ' + entry.display_name, 'success');
// OAuth popup if auth started during install (builtin creds)
if (res.auth_url) {
showToast('Opening authentication for ' + entry.display_name, 'info');
window.open(res.auth_url, '_blank', 'width=600,height=700');
}
loadExtensions();
// Auto-open configure for WASM channels
if (entry.kind === 'wasm_channel') {
@@ -1961,24 +1972,25 @@ function renderExtensionCard(ext) {
actions.appendChild(setupBtn);
}
} else {
// Non-WASM-channel extensions: original behavior
if (!ext.active) {
// WASM tools / MCP servers
const activeLabel = document.createElement('span');
activeLabel.className = 'ext-active-label';
activeLabel.textContent = ext.active ? 'Active' : 'Installed';
actions.appendChild(activeLabel);
// MCP servers may be installed but inactive — show Activate button
if (ext.kind === 'mcp_server' && !ext.active) {
const activateBtn = document.createElement('button');
activateBtn.className = 'btn-ext activate';
activateBtn.textContent = 'Activate';
activateBtn.addEventListener('click', () => activateExtension(ext.name));
actions.appendChild(activateBtn);
} else {
const activeLabel = document.createElement('span');
activeLabel.className = 'ext-active-label';
activeLabel.textContent = 'Active';
actions.appendChild(activeLabel);
}
if (ext.needs_setup) {
if (ext.needs_setup || ext.has_auth) {
const configBtn = document.createElement('button');
configBtn.className = 'btn-ext configure';
configBtn.textContent = ext.authenticated ? 'Reconfigure' : 'Setup';
configBtn.textContent = ext.authenticated ? 'Reconfigure' : 'Configure';
configBtn.addEventListener('click', () => showConfigureModal(ext.name));
actions.appendChild(configBtn);
}
@@ -2008,6 +2020,11 @@ function activateExtension(name) {
apiFetch('/api/extensions/' + encodeURIComponent(name) + '/activate', { method: 'POST' })
.then((res) => {
if (res.success) {
// Even on success, the tool may need OAuth (e.g., WASM loaded but no token yet)
if (res.auth_url) {
showToast('Opening authentication for ' + name, 'info');
window.open(res.auth_url, '_blank', 'width=600,height=700');
}
loadExtensions();
return;
}
@@ -2158,7 +2175,11 @@ function submitConfigureModal(name, fields) {
.then((res) => {
closeConfigureModal();
if (res.success) {
if (res.activated) {
if (res.auth_url) {
// OAuth flow started — open consent popup
showToast('Opening OAuth authorization for ' + name, 'info');
window.open(res.auth_url, '_blank', 'width=600,height=700');
} else if (res.activated) {
showToast('Configured and activated ' + name, 'success');
} else {
showToast(res.message || 'Configuration saved but activation failed', 'warning');
src/channels/web/types.rs
+3
View File
@@ -388,6 +388,9 @@ pub struct ExtensionInfo {
/// Whether this extension has configurable secrets (setup schema).
#[serde(default)]
pub needs_setup: bool,
/// Whether this extension has an auth configuration (OAuth or manual token).
#[serde(default)]
pub has_auth: bool,
/// WASM channel activation status: "installed", "configured", "active", "failed".
#[serde(skip_serializing_if = "Option::is_none")]
pub activation_status: Option<String>,
src/cli/oauth_defaults.rs
+440 -13
View File
@@ -17,11 +17,17 @@
//! - **Runtime**: Users can set GOOGLE_OAUTH_CLIENT_ID / GOOGLE_OAUTH_CLIENT_SECRET
//! env vars, which take priority over built-in defaults.
use std::collections::HashMap;
use std::time::Duration;
use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD};
use rand::RngCore;
use sha2::{Digest, Sha256};
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
use tokio::net::TcpListener;
use crate::secrets::{CreateSecretParams, SecretsStore};
// ── Built-in credentials ────────────────────────────────────────────────
pub struct OAuthCredentials {
@@ -121,6 +127,9 @@ pub enum OAuthCallbackError {
#[error("Timed out waiting for authorization")]
Timeout,
#[error("CSRF state mismatch: expected {expected}, got {actual}")]
StateMismatch { expected: String, actual: String },
#[error("IO error: {0}")]
Io(String),
}
@@ -177,16 +186,22 @@ pub async fn bind_callback_listener() -> Result<TcpListener, OAuthCallbackError>
/// extracts the value of `param_name` (e.g., "code" or "token"), and shows a branded
/// landing page using `display_name` (e.g., "Google", "Notion", "NEAR AI").
///
/// When `expected_state` is `Some`, the callback's `state` query parameter is validated
/// against it to prevent CSRF attacks. If the state doesn't match, the callback is
/// rejected with an error page.
///
/// Times out after 5 minutes.
pub async fn wait_for_callback(
listener: TcpListener,
path_prefix: &str,
param_name: &str,
display_name: &str,
expected_state: Option<&str>,
) -> Result<String, OAuthCallbackError> {
let path_prefix = path_prefix.to_string();
let param_name = param_name.to_string();
let display_name = display_name.to_string();
let expected_state = expected_state.map(String::from);
tokio::time::timeout(Duration::from_secs(300), async move {
loop {
@@ -221,17 +236,29 @@ pub async fn wait_for_callback(
return Err(OAuthCallbackError::Denied);
}
// Look for the target parameter
for param in query.split('&') {
let parts: Vec<&str> = param.splitn(2, '=').collect();
if parts.len() == 2 && parts[0] == param_name {
let value = urlencoding::decode(parts[1])
.unwrap_or_else(|_| parts[1].into())
.into_owned();
// Parse all query params into a map for validation
let params: HashMap<&str, String> = query
.split('&')
.filter_map(|p| {
let mut parts = p.splitn(2, '=');
let key = parts.next()?;
let val = parts.next().unwrap_or("");
Some((
key,
urlencoding::decode(val)
.unwrap_or_else(|_| val.into())
.into_owned(),
))
})
.collect();
let html = landing_html(&display_name, true);
// Validate CSRF state parameter
if let Some(ref expected) = expected_state {
let actual = params.get("state").cloned().unwrap_or_default();
if actual != *expected {
let html = landing_html(&display_name, false);
let response = format!(
"HTTP/1.1 200 OK\r\n\
"HTTP/1.1 403 Forbidden\r\n\
Content-Type: text/html; charset=utf-8\r\n\
Connection: close\r\n\
\r\n\
@@ -239,11 +266,29 @@ pub async fn wait_for_callback(
html
);
let _ = socket.write_all(response.as_bytes()).await;
let _ = socket.shutdown().await;
return Ok(value);
return Err(OAuthCallbackError::StateMismatch {
expected: expected.clone(),
actual,
});
}
}
// Look for the target parameter
if let Some(value) = params.get(param_name.as_str()) {
let html = landing_html(&display_name, true);
let response = format!(
"HTTP/1.1 200 OK\r\n\
Content-Type: text/html; charset=utf-8\r\n\
Connection: close\r\n\
\r\n\
{}",
html
);
let _ = socket.write_all(response.as_bytes()).await;
let _ = socket.shutdown().await;
return Ok(value.clone());
}
}
// Not the callback we're looking for
@@ -271,7 +316,288 @@ fn html_escape(s: &str) -> String {
out
}
/// HTML landing page shown in the browser after an OAuth redirect.
// ── Shared OAuth flow steps ─────────────────────────────────────────
/// Response from the OAuth token exchange.
pub struct OAuthTokenResponse {
pub access_token: String,
pub refresh_token: Option<String>,
pub expires_in: Option<u64>,
}
/// Result of building an OAuth 2.0 authorization URL.
pub struct OAuthUrlResult {
/// The full authorization URL to redirect the user to.
pub url: String,
/// PKCE code verifier (must be sent with the token exchange request).
pub code_verifier: Option<String>,
/// Random state parameter for CSRF protection (must be validated in callback).
pub state: String,
}
/// Build an OAuth 2.0 authorization URL with optional PKCE and CSRF state.
///
/// Returns an `OAuthUrlResult` containing the authorization URL, optional PKCE
/// code verifier, and a random `state` parameter for CSRF protection. The caller
/// must validate the `state` value in the callback before exchanging the code.
pub fn build_oauth_url(
authorization_url: &str,
client_id: &str,
redirect_uri: &str,
scopes: &[String],
use_pkce: bool,
extra_params: &HashMap<String, String>,
) -> OAuthUrlResult {
// Generate PKCE verifier and challenge
let (code_verifier, code_challenge) = if use_pkce {
let mut verifier_bytes = [0u8; 32];
rand::thread_rng().fill_bytes(&mut verifier_bytes);
let verifier = URL_SAFE_NO_PAD.encode(verifier_bytes);
let mut hasher = Sha256::new();
hasher.update(verifier.as_bytes());
let challenge = URL_SAFE_NO_PAD.encode(hasher.finalize());
(Some(verifier), Some(challenge))
} else {
(None, None)
};
// Generate random state for CSRF protection
let mut state_bytes = [0u8; 32];
rand::thread_rng().fill_bytes(&mut state_bytes);
let state = URL_SAFE_NO_PAD.encode(state_bytes);
// Build authorization URL
let mut auth_url = format!(
"{}?client_id={}&response_type=code&redirect_uri={}&state={}",
authorization_url,
urlencoding::encode(client_id),
urlencoding::encode(redirect_uri),
urlencoding::encode(&state),
);
if !scopes.is_empty() {
auth_url.push_str(&format!(
"&scope={}",
urlencoding::encode(&scopes.join(" "))
));
}
if let Some(ref challenge) = code_challenge {
auth_url.push_str(&format!(
"&code_challenge={}&code_challenge_method=S256",
challenge
));
}
for (key, value) in extra_params {
auth_url.push_str(&format!(
"&{}={}",
urlencoding::encode(key),
urlencoding::encode(value)
));
}
OAuthUrlResult {
url: auth_url,
code_verifier,
state,
}
}
/// Exchange an OAuth authorization code for tokens.
///
/// POSTs to `token_url` with the authorization code and optional PKCE verifier.
/// If `client_secret` is provided, uses HTTP Basic auth; otherwise includes
/// `client_id` in the form body (for public clients).
pub async fn exchange_oauth_code(
token_url: &str,
client_id: &str,
client_secret: Option<&str>,
code: &str,
redirect_uri: &str,
code_verifier: Option<&str>,
access_token_field: &str,
) -> Result<OAuthTokenResponse, OAuthCallbackError> {
let client = reqwest::Client::new();
let mut token_params = vec![
("grant_type", "authorization_code".to_string()),
("code", code.to_string()),
("redirect_uri", redirect_uri.to_string()),
];
if let Some(verifier) = code_verifier {
token_params.push(("code_verifier", verifier.to_string()));
}
let mut request = client.post(token_url);
if let Some(secret) = client_secret {
request = request.basic_auth(client_id, Some(secret));
} else {
token_params.push(("client_id", client_id.to_string()));
}
let token_response = request
.form(&token_params)
.send()
.await
.map_err(|e| OAuthCallbackError::Io(format!("Token exchange request failed: {}", e)))?;
if !token_response.status().is_success() {
let status = token_response.status();
let body = token_response.text().await.unwrap_or_default();
return Err(OAuthCallbackError::Io(format!(
"Token exchange failed: {} - {}",
status, body
)));
}
let token_data: serde_json::Value = token_response
.json()
.await
.map_err(|e| OAuthCallbackError::Io(format!("Failed to parse token response: {}", e)))?;
let access_token = token_data
.get(access_token_field)
.and_then(|v| v.as_str())
.ok_or_else(|| {
// Log only the field names present, not values (which may contain tokens)
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 token 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,
})
}
/// Store OAuth tokens (access + refresh) in the secrets store.
///
/// Also stores the granted scopes as `{secret_name}_scopes` so that scope
/// expansion can be detected on subsequent activations.
#[allow(clippy::too_many_arguments)]
pub async fn store_oauth_tokens(
store: &(dyn SecretsStore + Send + Sync),
user_id: &str,
secret_name: &str,
provider: Option<&str>,
access_token: &str,
refresh_token: Option<&str>,
expires_in: Option<u64>,
scopes: &[String],
) -> Result<(), OAuthCallbackError> {
let mut params = CreateSecretParams::new(secret_name, access_token);
if let Some(prov) = provider {
params = params.with_provider(prov);
}
if let Some(secs) = expires_in {
let expires_at = chrono::Utc::now() + chrono::Duration::seconds(secs as i64);
params = params.with_expiry(expires_at);
}
store
.create(user_id, params)
.await
.map_err(|e| OAuthCallbackError::Io(format!("Failed to save token: {}", e)))?;
// Store refresh token separately (no expiry, it's long-lived)
if let Some(rt) = refresh_token {
let refresh_name = format!("{}_refresh_token", secret_name);
let mut refresh_params = CreateSecretParams::new(&refresh_name, rt);
if let Some(prov) = provider {
refresh_params = refresh_params.with_provider(prov);
}
store
.create(user_id, refresh_params)
.await
.map_err(|e| OAuthCallbackError::Io(format!("Failed to save refresh token: {}", e)))?;
}
// Store granted scopes for scope expansion detection
if !scopes.is_empty() {
let scopes_name = format!("{}_scopes", secret_name);
let scopes_value = scopes.join(" ");
let scopes_params = CreateSecretParams::new(&scopes_name, &scopes_value);
// Best-effort: scope tracking failure shouldn't block auth
let _ = store.create(user_id, scopes_params).await;
}
Ok(())
}
/// Validate an OAuth token against a tool's validation endpoint.
///
/// Sends a request to the configured endpoint with the token as a Bearer header.
/// Returns `Ok(())` if the response status matches the expected success status,
/// or an error with details if validation fails (wrong account, expired token, etc.).
pub async fn validate_oauth_token(
token: &str,
validation: &crate::tools::wasm::ValidationEndpointSchema,
) -> Result<(), OAuthCallbackError> {
let client = reqwest::Client::builder()
.timeout(Duration::from_secs(10))
.build()
.map_err(|e| OAuthCallbackError::Io(format!("Failed to build HTTP client: {}", e)))?;
let request = match validation.method.to_uppercase().as_str() {
"POST" => client.post(&validation.url),
_ => client.get(&validation.url),
};
let mut request = request.header("Authorization", format!("Bearer {}", token));
// Add custom headers from the validation schema (e.g., Notion-Version)
for (key, value) in &validation.headers {
request = request.header(key, value);
}
let response = request
.send()
.await
.map_err(|e| OAuthCallbackError::Io(format!("Validation request failed: {}", e)))?;
if response.status().as_u16() == validation.success_status {
Ok(())
} else {
let status = response.status();
let body = response.text().await.unwrap_or_default();
let truncated: String = if body.len() > 200 {
let mut end = 200;
while end > 0 && !body.is_char_boundary(end) {
end -= 1;
}
format!("{}...", &body[..end])
} else {
body
};
Err(OAuthCallbackError::Io(format!(
"Token validation failed: HTTP {} (expected {}): {}",
status, validation.success_status, truncated
)))
}
}
// ── Landing pages ───────────────────────────────────────────────────
pub fn landing_html(provider_name: &str, success: bool) -> String {
let safe_name = html_escape(provider_name);
let (icon, heading, subtitle, accent) = if success {
@@ -512,4 +838,105 @@ mod tests {
assert!(html.contains("#ef4444")); // red accent
assert!(!html.contains("Connected"));
}
#[test]
fn test_build_oauth_url_basic() {
use std::collections::HashMap;
use crate::cli::oauth_defaults::build_oauth_url;
let result = build_oauth_url(
"https://accounts.google.com/o/oauth2/auth",
"my-client-id",
"http://localhost:9876/callback",
&["openid".to_string(), "email".to_string()],
false,
&HashMap::new(),
);
assert!(
result
.url
.starts_with("https://accounts.google.com/o/oauth2/auth?")
);
assert!(result.url.contains("client_id=my-client-id"));
assert!(result.url.contains("response_type=code"));
assert!(result.url.contains("redirect_uri="));
assert!(result.url.contains("scope=openid%20email"));
assert!(result.url.contains("state="));
assert!(result.code_verifier.is_none());
assert!(!result.state.is_empty());
}
#[test]
fn test_build_oauth_url_with_pkce() {
use std::collections::HashMap;
use crate::cli::oauth_defaults::build_oauth_url;
let result = build_oauth_url(
"https://auth.example.com/authorize",
"client-123",
"http://localhost:9876/callback",
&[],
true,
&HashMap::new(),
);
assert!(result.url.contains("code_challenge="));
assert!(result.url.contains("code_challenge_method=S256"));
assert!(result.code_verifier.is_some());
let verifier = result.code_verifier.unwrap();
assert!(!verifier.is_empty());
}
#[test]
fn test_build_oauth_url_with_extra_params() {
use std::collections::HashMap;
use crate::cli::oauth_defaults::build_oauth_url;
let mut extra = HashMap::new();
extra.insert("access_type".to_string(), "offline".to_string());
extra.insert("prompt".to_string(), "consent".to_string());
let result = build_oauth_url(
"https://auth.example.com/authorize",
"client-123",
"http://localhost:9876/callback",
&["read".to_string()],
false,
&extra,
);
assert!(result.url.contains("access_type=offline"));
assert!(result.url.contains("prompt=consent"));
}
#[test]
fn test_build_oauth_url_state_is_unique() {
use std::collections::HashMap;
use crate::cli::oauth_defaults::build_oauth_url;
let result1 = build_oauth_url(
"https://auth.example.com/authorize",
"client",
"http://localhost:9876/callback",
&[],
false,
&HashMap::new(),
);
let result2 = build_oauth_url(
"https://auth.example.com/authorize",
"client",
"http://localhost:9876/callback",
&[],
false,
&HashMap::new(),
);
// State should be different each time (random)
assert_ne!(result1.state, result2.state);
}
}
src/cli/tool.rs
+58 -184
View File
@@ -782,11 +782,7 @@ async fn auth_tool_oauth(
auth: &crate::tools::wasm::AuthCapabilitySchema,
oauth: &crate::tools::wasm::OAuthConfigSchema,
) -> anyhow::Result<()> {
use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD};
use rand::RngCore;
use sha2::{Digest, Sha256};
use crate::cli::oauth_defaults::{self, OAUTH_CALLBACK_PORT};
use crate::cli::oauth_defaults;
let display_name = auth.display_name.as_deref().unwrap_or(&auth.secret_name);
@@ -827,142 +823,69 @@ async fn auth_tool_oauth(
println!();
let listener = oauth_defaults::bind_callback_listener().await?;
let redirect_uri = format!("http://localhost:{}/callback", OAUTH_CALLBACK_PORT);
let redirect_uri = format!("{}/callback", oauth_defaults::callback_url());
// Generate PKCE verifier and challenge
let (code_verifier, code_challenge) = if oauth.use_pkce {
let mut verifier_bytes = [0u8; 32];
rand::thread_rng().fill_bytes(&mut verifier_bytes);
let verifier = URL_SAFE_NO_PAD.encode(verifier_bytes);
let mut hasher = Sha256::new();
hasher.update(verifier.as_bytes());
let challenge = URL_SAFE_NO_PAD.encode(hasher.finalize());
(Some(verifier), Some(challenge))
} else {
(None, None)
};
// Build authorization URL
let mut auth_url = format!(
"{}?client_id={}&response_type=code&redirect_uri={}",
oauth.authorization_url,
urlencoding::encode(&client_id),
urlencoding::encode(&redirect_uri)
// Build authorization URL with PKCE and CSRF state
let oauth_result = oauth_defaults::build_oauth_url(
&oauth.authorization_url,
&client_id,
&redirect_uri,
&oauth.scopes,
oauth.use_pkce,
&oauth.extra_params,
);
if !oauth.scopes.is_empty() {
auth_url.push_str(&format!(
"&scope={}",
urlencoding::encode(&oauth.scopes.join(" "))
));
}
if let Some(ref challenge) = code_challenge {
auth_url.push_str(&format!(
"&code_challenge={}&code_challenge_method=S256",
challenge
));
}
// Add extra params
for (key, value) in &oauth.extra_params {
auth_url.push_str(&format!(
"&{}={}",
urlencoding::encode(key),
urlencoding::encode(value)
));
}
let code_verifier = oauth_result.code_verifier;
println!(" Opening browser for {} login...", display_name);
println!();
if let Err(e) = open::that(&auth_url) {
if let Err(e) = open::that(&oauth_result.url) {
println!(" Could not open browser: {}", e);
println!(" Please open this URL manually:");
println!(" {}", auth_url);
println!(" {}", oauth_result.url);
}
println!(" Waiting for authorization...");
let code =
oauth_defaults::wait_for_callback(listener, "/callback", "code", display_name).await?;
let code = oauth_defaults::wait_for_callback(
listener,
"/callback",
"code",
display_name,
Some(&oauth_result.state),
)
.await?;
println!();
println!(" Exchanging code for token...");
// Exchange code for token
let client = reqwest::Client::new();
let mut token_params = vec![
("grant_type", "authorization_code".to_string()),
("code", code),
("redirect_uri", redirect_uri),
];
if let Some(ref verifier) = code_verifier {
token_params.push(("code_verifier", verifier.to_string()));
}
// Build token request
let mut request = client.post(&oauth.token_url);
// Use Basic auth if client_secret is provided, otherwise include client_id in body
if let Some(ref secret) = client_secret {
request = request.basic_auth(&client_id, Some(secret));
} else {
token_params.push(("client_id", client_id));
}
let token_response = request.form(&token_params).send().await?;
if !token_response.status().is_success() {
let status = token_response.status();
let body = token_response.text().await.unwrap_or_default();
return Err(anyhow::anyhow!(
"Token exchange failed: {} - {}",
status,
body
));
}
let token_data: serde_json::Value = token_response.json().await?;
let access_token = token_data
.get(&oauth.access_token_field)
.and_then(|v| v.as_str())
.ok_or_else(|| {
anyhow::anyhow!(
"No {} in token response: {:?}",
oauth.access_token_field,
token_data
)
})?;
let refresh_token = token_data.get("refresh_token").and_then(|v| v.as_str());
let expires_in = token_data.get("expires_in").and_then(|v| v.as_u64());
// Save the token (with refresh token and expiry if provided)
save_token(
store,
user_id,
auth,
access_token,
refresh_token,
expires_in,
let token_response = oauth_defaults::exchange_oauth_code(
&oauth.token_url,
&client_id,
client_secret.as_deref(),
&code,
&redirect_uri,
code_verifier.as_deref(),
&oauth.access_token_field,
)
.await?;
// Extract any additional info for display
let workspace_name = token_data
.get("workspace_name")
.and_then(|v| v.as_str())
.or_else(|| token_data.get("team_name").and_then(|v| v.as_str()));
// Save tokens (access + refresh + scopes)
oauth_defaults::store_oauth_tokens(
store,
user_id,
&auth.secret_name,
auth.provider.as_deref(),
&token_response.access_token,
token_response.refresh_token.as_deref(),
token_response.expires_in,
&oauth.scopes,
)
.await?;
println!();
println!(" ✓ {} connected!", display_name);
if let Some(workspace) = workspace_name {
println!(" Workspace: {}", workspace);
}
println!();
println!(" The tool can now access the API.");
println!();
@@ -1107,46 +1030,15 @@ async fn validate_token(
validation: &crate::tools::wasm::ValidationEndpointSchema,
_secret_name: &str,
) -> anyhow::Result<()> {
let client = reqwest::Client::builder()
.timeout(std::time::Duration::from_secs(10))
.build()?;
// Build request based on method
let request = match validation.method.to_uppercase().as_str() {
"GET" => client.get(&validation.url),
"POST" => client.post(&validation.url),
_ => client.get(&validation.url),
};
// Add authorization header (assume Bearer for now, could be extended)
let response = request
.header("Authorization", format!("Bearer {}", token))
.header("Notion-Version", "2022-06-28") // Notion-specific, but harmless for others
.send()
.await?;
if response.status().as_u16() == validation.success_status {
Ok(())
} else {
let status = response.status();
let body = response.text().await.unwrap_or_default();
Err(anyhow::anyhow!(
"HTTP {} (expected {}): {}",
status,
validation.success_status,
if body.len() > 100 {
format!("{}...", &body[..100])
} else {
body
}
))
}
crate::cli::oauth_defaults::validate_oauth_token(token, validation)
.await
.map_err(|e| anyhow::anyhow!("{}", e))
}
/// Save token to secrets store.
///
/// Optionally stores a refresh token (as `{secret_name}_refresh_token`) and
/// sets `expires_at` on the access token so the runtime can auto-refresh.
/// Delegates to the shared `store_oauth_tokens` for OAuth tokens, or stores
/// directly for manual/env-var tokens (no scopes or refresh token).
async fn save_token(
store: &(dyn SecretsStore + Send + Sync),
user_id: &str,
@@ -1155,36 +1047,18 @@ async fn save_token(
refresh_token: Option<&str>,
expires_in: Option<u64>,
) -> anyhow::Result<()> {
let mut params = CreateSecretParams::new(&auth.secret_name, token);
if let Some(ref provider) = auth.provider {
params = params.with_provider(provider);
}
if let Some(secs) = expires_in {
let expires_at = chrono::Utc::now() + chrono::Duration::seconds(secs as i64);
params = params.with_expiry(expires_at);
}
store
.create(user_id, params)
.await
.map_err(|e| anyhow::anyhow!("Failed to save token: {}", e))?;
// Store refresh token separately (no expiry, it's long-lived)
if let Some(rt) = refresh_token {
let refresh_name = format!("{}_refresh_token", auth.secret_name);
let mut refresh_params = CreateSecretParams::new(&refresh_name, rt);
if let Some(ref provider) = auth.provider {
refresh_params = refresh_params.with_provider(provider);
}
store
.create(user_id, refresh_params)
.await
.map_err(|e| anyhow::anyhow!("Failed to save refresh token: {}", e))?;
}
Ok(())
crate::cli::oauth_defaults::store_oauth_tokens(
store,
user_id,
&auth.secret_name,
auth.provider.as_deref(),
token,
refresh_token,
expires_in,
&[], // No scopes for manual/env-var tokens
)
.await
.map_err(|e| anyhow::anyhow!("{}", e))
}
/// Print success message.
src/extensions/manager.rs
+510 -19
View File
@@ -38,6 +38,9 @@ struct PendingAuth {
_name: String,
_kind: ExtensionKind,
created_at: std::time::Instant,
/// Background task listening for the OAuth callback.
/// Aborted when a new auth flow starts for the same extension.
task_handle: Option<tokio::task::JoinHandle<()>>,
}
/// Runtime infrastructure needed for hot-activating WASM channels.
@@ -58,6 +61,8 @@ pub struct SetupResult {
pub message: String,
/// Whether the channel was successfully activated after saving secrets.
pub activated: bool,
/// OAuth authorization URL for the UI to open (if OAuth flow was started).
pub auth_url: Option<String>,
}
/// Central manager for extension lifecycle operations.
@@ -385,6 +390,7 @@ impl ExtensionManager {
active,
tools,
needs_setup: false,
has_auth: false,
installed: true,
activation_error: None,
});
@@ -411,6 +417,11 @@ impl ExtensionManager {
.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();
extensions.push(InstalledExtension {
name: name.clone(),
kind: ExtensionKind::WasmTool,
@@ -421,6 +432,7 @@ impl ExtensionManager {
active,
tools: if active { vec![name] } else { Vec::new() },
needs_setup,
has_auth,
installed: true,
activation_error: None,
});
@@ -460,6 +472,7 @@ impl ExtensionManager {
active,
tools: Vec::new(),
needs_setup,
has_auth: false,
installed: true,
activation_error,
});
@@ -497,6 +510,7 @@ impl ExtensionManager {
active: false,
tools: Vec::new(),
needs_setup: false,
has_auth: false,
installed: false,
activation_error: None,
});
@@ -1338,6 +1352,7 @@ impl ExtensionManager {
_name: name.to_string(),
_kind: ExtensionKind::McpServer,
created_at: std::time::Instant::now(),
task_handle: None,
},
);
@@ -1424,23 +1439,45 @@ impl ExtensionManager {
});
}
// Check if already authenticated
if self
// Check if already authenticated (with scope expansion detection)
let token_exists = self
.secrets
.exists(&self.user_id, &auth.secret_name)
.await
.unwrap_or(false)
{
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(),
});
.unwrap_or(false);
if token_exists {
// If this tool has OAuth config, check whether new scopes are needed
let needs_reauth = if let Some(ref oauth) = auth.oauth {
let merged = self
.collect_shared_scopes(&auth.secret_name, &oauth.scopes)
.await;
let needs = self.needs_scope_expansion(&auth.secret_name, &merged).await;
tracing::debug!(
tool = name,
secret_name = %auth.secret_name,
merged_scopes = ?merged,
needs_reauth = needs,
"Scope expansion check"
);
needs
} else {
false
};
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(),
});
}
// Fall through to OAuth branch for scope expansion
}
// If a token was provided, store it
@@ -1464,6 +1501,62 @@ impl ExtensionManager {
});
}
// 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 {
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!(
"Configure OAuth credentials for {} in the Setup tab.",
display
)),
setup_url: auth.setup_url.clone(),
awaiting_token: false,
status: "needs_setup".to_string(),
});
}
return self
.start_wasm_oauth(name, &auth, oauth)
.await
.map_err(|e| ExtensionError::AuthFailed(e.to_string()));
}
// Return instructions for manual token entry
let display = auth.display_name.unwrap_or_else(|| name.to_string());
let instructions = auth
@@ -1534,6 +1627,353 @@ impl ExtensionManager {
crate::tools::wasm::CapabilitiesFile::from_bytes(&cap_bytes).ok()
}
/// Collect merged OAuth scopes from all installed tools sharing the same secret_name.
///
/// When multiple tools share an OAuth provider (e.g., google-calendar and google-drive
/// both use `google_oauth_token`), we request all their scopes in a single OAuth flow
/// so one login covers everything.
async fn collect_shared_scopes(
&self,
secret_name: &str,
base_scopes: &[String],
) -> Vec<String> {
let mut all_scopes: std::collections::BTreeSet<String> =
base_scopes.iter().cloned().collect();
if let Ok(tools) = discover_tools(&self.wasm_tools_dir).await {
for tool_name in tools.keys() {
if let Some(cap) = self.load_tool_capabilities(tool_name).await
&& let Some(auth) = &cap.auth
&& auth.secret_name == secret_name
&& let Some(oauth) = &auth.oauth
{
all_scopes.extend(oauth.scopes.iter().cloned());
}
}
}
all_scopes.into_iter().collect()
}
/// Check whether the stored scopes are insufficient for the merged scopes.
async fn needs_scope_expansion(&self, secret_name: &str, merged_scopes: &[String]) -> bool {
if merged_scopes.is_empty() {
return false;
}
let scopes_key = format!("{}_scopes", secret_name);
let stored_scopes: std::collections::HashSet<String> =
match self.secrets.get_decrypted(&self.user_id, &scopes_key).await {
Ok(secret) => {
let scopes: std::collections::HashSet<String> = secret
.expose()
.split_whitespace()
.map(String::from)
.collect();
tracing::debug!(
secret_name,
stored_scopes = ?scopes,
"Loaded stored scopes for expansion check"
);
scopes
}
Err(_) => {
// No stored scopes record — this is a legacy token created before
// scope tracking. Force re-auth to ensure all required scopes are granted.
tracing::debug!(
secret_name,
"No stored scopes record, forcing re-auth for legacy token"
);
return true;
}
};
// Check if any merged scope is missing from stored scopes
merged_scopes
.iter()
.any(|scope| !stored_scopes.contains(scope))
}
/// Find the setup secret names for OAuth client_id and client_secret.
///
/// Scans `setup.required_secrets` for names containing "client_id" and "client_secret".
/// Returns `(Option<(name, optional)>, Option<(name, optional)>)`.
async fn find_setup_credential_names(
&self,
tool_name: &str,
) -> (Option<(String, bool)>, Option<(String, bool)>) {
let Some(cap) = self.load_tool_capabilities(tool_name).await else {
return (None, None);
};
let Some(setup) = &cap.setup else {
return (None, None);
};
let mut client_id_entry = None;
let mut client_secret_entry = None;
for secret in &setup.required_secrets {
let lower = secret.name.to_lowercase();
if lower.ends_with("client_id") || lower == "client_id" {
client_id_entry = Some((secret.name.clone(), secret.optional));
} else if lower.ends_with("client_secret") || lower == "client_secret" {
client_secret_entry = Some((secret.name.clone(), secret.optional));
}
}
(client_id_entry, client_secret_entry)
}
/// 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
/// may have been entered via the Setup tab (stored as setup secrets).
async fn resolve_oauth_credential(
&self,
inline_value: &Option<String>,
env_var_name: &Option<String>,
builtin_value: Option<&str>,
setup_secret_name: Option<&str>,
) -> Option<String> {
// 1. Check secrets store (entered via Setup tab)
if let Some(secret_name) = setup_secret_name
&& let Ok(secret) = self.secrets.get_decrypted(&self.user_id, secret_name).await
{
let val = secret.expose();
if !val.is_empty() {
return Some(val.to_string());
}
}
// 2. Inline value from capabilities.json
if let Some(val) = inline_value {
return Some(val.clone());
}
// 3. Runtime environment variable
if let Some(env) = env_var_name
&& let Ok(val) = std::env::var(env)
{
return Some(val);
}
// 4. Built-in defaults
builtin_value.map(String::from)
}
/// Start the OAuth browser flow for a WASM tool.
///
/// Binds a callback listener, builds the authorization URL, spawns a background
/// task to wait for the callback and exchange the code, then returns the auth URL
/// immediately so the web UI can open it.
async fn start_wasm_oauth(
&self,
name: &str,
auth: &crate::tools::wasm::AuthCapabilitySchema,
oauth: &crate::tools::wasm::OAuthConfigSchema,
) -> Result<AuthResult, String> {
use crate::cli::oauth_defaults;
let builtin = oauth_defaults::builtin_credentials(&auth.secret_name);
// Find setup secret names for client_id and client_secret from capabilities.
// These are the actual names used in the Setup tab (e.g., "google_oauth_client_id"),
// which may differ from "{secret_name}_client_id".
let (setup_client_id_entry, setup_client_secret_entry) =
self.find_setup_credential_names(name).await;
let setup_client_id_name = setup_client_id_entry.map(|(n, _)| n);
let setup_client_secret_name = setup_client_secret_entry.map(|(n, _)| n);
// Resolve client_id: setup secrets → inline → env var → builtin
let client_id = self
.resolve_oauth_credential(
&oauth.client_id,
&oauth.client_id_env,
builtin.as_ref().map(|c| c.client_id),
setup_client_id_name.as_deref(),
)
.await
.ok_or_else(|| {
let env_name = oauth
.client_id_env
.as_deref()
.unwrap_or("the client_id env var");
let mut msg = format!(
"OAuth client_id not configured for '{}'. \
Enter it in the Setup tab or set {} env var",
name, env_name
);
// Only mention the Google-specific build flag for Google providers
if auth.secret_name.to_lowercase().contains("google") {
msg.push_str(", or build with IRONCLAW_GOOGLE_CLIENT_ID");
}
msg.push('.');
msg
})?;
// Resolve client_secret (optional for PKCE-only flows)
let client_secret = self
.resolve_oauth_credential(
&oauth.client_secret,
&oauth.client_secret_env,
builtin.as_ref().map(|c| c.client_secret),
setup_client_secret_name.as_deref(),
)
.await;
// Cancel any existing pending auth for this tool (frees port 9876)
{
let mut pending = self.pending_auth.write().await;
if let Some(old) = pending.remove(name)
&& let Some(handle) = old.task_handle
{
handle.abort();
}
}
// Bind callback listener
let listener = oauth_defaults::bind_callback_listener()
.await
.map_err(|e| format!("Failed to start OAuth callback listener: {}", e))?;
let redirect_uri = format!("{}/callback", oauth_defaults::callback_url());
// Merge scopes from all tools sharing this provider
let merged_scopes = self
.collect_shared_scopes(&auth.secret_name, &oauth.scopes)
.await;
// Build authorization URL with CSRF state
let oauth_result = oauth_defaults::build_oauth_url(
&oauth.authorization_url,
&client_id,
&redirect_uri,
&merged_scopes,
oauth.use_pkce,
&oauth.extra_params,
);
let auth_url = oauth_result.url.clone();
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),
)
.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 {
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(),
})
}
/// Check whether a WASM tool's required setup secrets are provided.
///
/// Returns `(authenticated, needs_setup)` — same semantics as `check_channel_auth_status`.
@@ -2261,7 +2701,16 @@ impl ExtensionManager {
async fn cleanup_expired_auths(&self) {
let mut pending = self.pending_auth.write().await;
pending.retain(|_, auth| auth.created_at.elapsed() < std::time::Duration::from_secs(300));
pending.retain(|_, auth| {
let expired = auth.created_at.elapsed() >= std::time::Duration::from_secs(300);
if expired {
// Abort the background listener task to free port 9876
if let Some(ref handle) = auth.task_handle {
handle.abort();
}
}
!expired
});
}
/// Get the setup schema for an extension (secret fields and their status).
@@ -2478,16 +2927,55 @@ impl ExtensionManager {
}
}
// For tools, save and attempt auto-activation
// For tools, save and attempt auto-activation, then check auth.
if kind == ExtensionKind::WasmTool {
match self.activate_wasm_tool(name).await {
Ok(result) => {
return Ok(SetupResult {
message: format!(
// Delete existing OAuth token so auth() starts a fresh flow.
// Done AFTER activation succeeds to avoid losing tokens on failure.
// This covers Reconfigure: user wants to re-auth (switch account, update creds).
if let Some(cap) = self.load_tool_capabilities(name).await
&& let Some(ref auth_cfg) = cap.auth
&& auth_cfg.oauth.is_some()
{
let _ = self
.secrets
.delete(&self.user_id, &auth_cfg.secret_name)
.await;
let _ = self
.secrets
.delete(&self.user_id, &format!("{}_scopes", auth_cfg.secret_name))
.await;
let _ = self
.secrets
.delete(
&self.user_id,
&format!("{}_refresh_token", auth_cfg.secret_name),
)
.await;
}
// Check if auth is needed (OAuth or manual token).
// 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;
}
let message = if auth_url.is_some() {
format!(
"Configuration saved and tool '{}' activated. Complete OAuth in your browser.",
name
)
} else {
format!(
"Configuration saved and tool '{}' activated. {}",
name, result.message
),
)
};
return Ok(SetupResult {
message,
activated: true,
auth_url,
});
}
Err(e) => {
@@ -2499,6 +2987,7 @@ impl ExtensionManager {
return Ok(SetupResult {
message: format!("Configuration saved for '{}'.", name),
activated: false,
auth_url: None,
});
}
}
@@ -2515,6 +3004,7 @@ impl ExtensionManager {
name, result.message
),
activated: true,
auth_url: None,
})
}
Err(e) => {
@@ -2536,6 +3026,7 @@ impl ExtensionManager {
name, e
),
activated: false,
auth_url: None,
})
}
}
src/extensions/mod.rs
+3
View File
@@ -204,6 +204,9 @@ pub struct InstalledExtension {
/// Whether this extension has a setup schema (required_secrets) that can be configured.
#[serde(default)]
pub needs_setup: bool,
/// Whether this extension has an auth configuration (OAuth or manual token).
#[serde(default)]
pub has_auth: bool,
/// Whether this extension is installed locally (false = available in registry but not installed).
#[serde(default = "default_true")]
pub installed: bool,
src/llm/session.rs
+1 -1
View File
@@ -347,7 +347,7 @@ impl SessionManager {
// The NEAR AI API redirects to: {frontend_callback}/auth/callback?token=X&...
let session_token =
oauth_defaults::wait_for_callback(listener, "/auth/callback", "token", "NEAR AI")
oauth_defaults::wait_for_callback(listener, "/auth/callback", "token", "NEAR AI", None)
.await
.map_err(|e| LlmError::SessionRenewalFailed {
provider: "nearai".to_string(),
src/tools/mcp/auth.rs
+4 -1
View File
@@ -531,7 +531,7 @@ pub async fn wait_for_authorization_callback(
listener: TcpListener,
server_name: &str,
) -> Result<String, AuthError> {
oauth_defaults::wait_for_callback(listener, "/callback", "code", server_name)
oauth_defaults::wait_for_callback(listener, "/callback", "code", server_name, None)
.await
.map_err(|e| match e {
oauth_defaults::OAuthCallbackError::Denied => AuthError::AuthorizationDenied,
@@ -539,6 +539,9 @@ pub async fn wait_for_authorization_callback(
oauth_defaults::OAuthCallbackError::PortInUse(_, msg) => {
AuthError::Http(format!("Port error: {}", msg))
}
oauth_defaults::OAuthCallbackError::StateMismatch { .. } => {
AuthError::Http("CSRF state mismatch in OAuth callback".to_string())
}
oauth_defaults::OAuthCallbackError::Io(msg) => AuthError::Http(msg),
})
}
src/tools/wasm/capabilities_schema.rs
+5
View File
@@ -512,6 +512,11 @@ pub struct ValidationEndpointSchema {
/// Expected HTTP status code for success (defaults to 200).
#[serde(default = "default_success_status")]
pub success_status: u16,
/// Additional headers to send with the validation request.
/// Used for service-specific requirements (e.g., Notion-Version for Notion API).
#[serde(default)]
pub headers: HashMap<String, String>,
}
fn default_method() -> String {