Add hosted MCP server support with OAuth 2.1 and token refresh

Enables connecting to official MCP servers (like Notion) instead of
building custom WASM tools. Uses OAuth 2.1 with PKCE and supports
Dynamic Client Registration for zero-config authentication.

Key features:
- OAuth 2.1 flow with PKCE for secure browser-based auth
- Dynamic Client Registration (DCR) for servers without pre-configured clients
- Automatic token refresh on 401 responses
- Session management with Mcp-Session-Id headers
- SSE streaming response handling

New CLI commands:
- `mcp add <name> <url>` - Add an MCP server
- `mcp remove <name>` - Remove an MCP server
- `mcp list` - List configured servers
- `mcp auth <name>` - Authenticate with a server
- `mcp test <name>` - Test connection

Also removes the Notion WASM tool example since it's superseded by the
Notion MCP server which provides 13 official tools.

Co-Authored-By: Claude Opus 4.5 <[email protected]>
This commit is contained in:
Illia Polosukhin
2026-02-05 17:43:26 -08:00
co-authored by Claude Opus 4.5
parent 5992e27507
commit 974bc8d407
23 changed files with 3674 additions and 1838 deletions
+484
View File
@@ -0,0 +1,484 @@
//! MCP server management CLI commands.
//!
//! Commands for adding, removing, authenticating, and testing MCP servers.
use std::io::Write;
use std::sync::Arc;
use clap::Subcommand;
use crate::config::Config;
use crate::history::Store;
use crate::secrets::{PostgresSecretsStore, SecretsCrypto, SecretsStore};
use crate::tools::mcp::{
McpClient, McpServerConfig, McpSessionManager, OAuthConfig,
auth::{authorize_mcp_server, is_authenticated},
config::{
add_mcp_server, get_mcp_server, load_mcp_servers, remove_mcp_server, save_mcp_servers,
},
};
#[derive(Subcommand, Debug, Clone)]
pub enum McpCommand {
/// Add an MCP server
Add {
/// Server name (e.g., "notion", "github")
name: String,
/// Server URL (e.g., "https://mcp.notion.com")
url: String,
/// OAuth client ID (if authentication is required)
#[arg(long)]
client_id: Option<String>,
/// OAuth authorization URL (optional, can be discovered)
#[arg(long)]
auth_url: Option<String>,
/// OAuth token URL (optional, can be discovered)
#[arg(long)]
token_url: Option<String>,
/// Scopes to request (comma-separated)
#[arg(long)]
scopes: Option<String>,
/// Server description
#[arg(long)]
description: Option<String>,
},
/// Remove an MCP server
Remove {
/// Server name to remove
name: String,
},
/// List configured MCP servers
List {
/// Show detailed information
#[arg(short, long)]
verbose: bool,
},
/// Authenticate with an MCP server (OAuth flow)
Auth {
/// Server name to authenticate
name: String,
/// User ID for storing the token (default: "default")
#[arg(short, long, default_value = "default")]
user: String,
},
/// Test connection to an MCP server
Test {
/// Server name to test
name: String,
/// User ID for authentication (default: "default")
#[arg(short, long, default_value = "default")]
user: String,
},
/// Enable or disable an MCP server
Toggle {
/// Server name
name: String,
/// Enable the server
#[arg(long, conflicts_with = "disable")]
enable: bool,
/// Disable the server
#[arg(long, conflicts_with = "enable")]
disable: bool,
},
}
/// Run an MCP command.
pub async fn run_mcp_command(cmd: McpCommand) -> anyhow::Result<()> {
match cmd {
McpCommand::Add {
name,
url,
client_id,
auth_url,
token_url,
scopes,
description,
} => {
add_server(
name,
url,
client_id,
auth_url,
token_url,
scopes,
description,
)
.await
}
McpCommand::Remove { name } => remove_server(name).await,
McpCommand::List { verbose } => list_servers(verbose).await,
McpCommand::Auth { name, user } => auth_server(name, user).await,
McpCommand::Test { name, user } => test_server(name, user).await,
McpCommand::Toggle {
name,
enable,
disable,
} => toggle_server(name, enable, disable).await,
}
}
/// Add a new MCP server.
async fn add_server(
name: String,
url: String,
client_id: Option<String>,
auth_url: Option<String>,
token_url: Option<String>,
scopes: Option<String>,
description: Option<String>,
) -> anyhow::Result<()> {
let mut config = McpServerConfig::new(&name, &url);
if let Some(desc) = description {
config = config.with_description(desc);
}
// Track if auth is required
let requires_auth = client_id.is_some();
// Set up OAuth if client_id is provided
if let Some(client_id) = client_id {
let mut oauth = OAuthConfig::new(client_id);
if let (Some(auth), Some(token)) = (auth_url, token_url) {
oauth = oauth.with_endpoints(auth, token);
}
if let Some(scopes_str) = scopes {
let scope_list: Vec<String> = scopes_str
.split(',')
.map(|s| s.trim().to_string())
.collect();
oauth = oauth.with_scopes(scope_list);
}
config = config.with_oauth(oauth);
}
// Validate
config.validate()?;
// Save
add_mcp_server(config).await?;
println!();
println!(" ✓ Added MCP server '{}'", name);
println!(" URL: {}", url);
if requires_auth {
println!();
println!(" Run 'ironclaw mcp auth {}' to authenticate.", name);
}
println!();
Ok(())
}
/// Remove an MCP server.
async fn remove_server(name: String) -> anyhow::Result<()> {
remove_mcp_server(&name).await?;
println!();
println!(" ✓ Removed MCP server '{}'", name);
println!();
Ok(())
}
/// List configured MCP servers.
async fn list_servers(verbose: bool) -> anyhow::Result<()> {
let servers = load_mcp_servers().await?;
if servers.servers.is_empty() {
println!();
println!(" No MCP servers configured.");
println!();
println!(" Add a server with:");
println!(" ironclaw mcp add <name> <url> [--client-id <id>]");
println!();
return Ok(());
}
println!();
println!(" Configured MCP servers:");
println!();
for server in &servers.servers {
let status = if server.enabled { "" } else { "" };
let auth_status = if server.requires_auth() {
" (auth required)"
} else {
""
};
if verbose {
println!(" {} {}{}", status, server.name, auth_status);
println!(" URL: {}", server.url);
if let Some(ref desc) = server.description {
println!(" Description: {}", desc);
}
if let Some(ref oauth) = server.oauth {
println!(" OAuth Client ID: {}", oauth.client_id);
if !oauth.scopes.is_empty() {
println!(" Scopes: {}", oauth.scopes.join(", "));
}
}
println!();
} else {
println!(
" {} {} - {}{}",
status, server.name, server.url, auth_status
);
}
}
if !verbose {
println!();
println!(" Use --verbose for more details.");
}
println!();
Ok(())
}
/// Authenticate with an MCP server.
async fn auth_server(name: String, user_id: String) -> anyhow::Result<()> {
// Get server config
let server = get_mcp_server(&name).await?;
// Initialize secrets store
let secrets = get_secrets_store().await?;
// Check if already authenticated
if is_authenticated(&server, &secrets, &user_id).await {
println!();
println!(" Server '{}' is already authenticated.", name);
println!();
print!(" Re-authenticate? [y/N]: ");
std::io::stdout().flush()?;
let mut input = String::new();
std::io::stdin().read_line(&mut input)?;
if !input.trim().eq_ignore_ascii_case("y") {
return Ok(());
}
println!();
}
println!();
println!("╔════════════════════════════════════════════════════════════════╗");
println!(
"{:^62}",
format!("{} Authentication", name.to_uppercase())
);
println!("╚════════════════════════════════════════════════════════════════╝");
println!();
// Perform OAuth flow (supports both pre-configured OAuth and DCR)
match authorize_mcp_server(&server, &secrets, &user_id).await {
Ok(_token) => {
println!();
println!(" ✓ Successfully authenticated with '{}'!", name);
println!();
println!(" You can now use tools from this server.");
println!();
}
Err(crate::tools::mcp::auth::AuthError::NotSupported) => {
println!();
println!(" ✗ Server does not support OAuth authentication.");
println!();
println!(" The server may require a different authentication method,");
println!(" or you may need to configure OAuth manually:");
println!();
println!(" ironclaw mcp remove {}", name);
println!(
" ironclaw mcp add {} {} --client-id YOUR_CLIENT_ID",
name, server.url
);
println!();
}
Err(e) => {
println!();
println!(" ✗ Authentication failed: {}", e);
println!();
return Err(e.into());
}
}
Ok(())
}
/// Test connection to an MCP server.
async fn test_server(name: String, user_id: String) -> anyhow::Result<()> {
// Get server config
let server = get_mcp_server(&name).await?;
println!();
println!(" Testing connection to '{}'...", name);
// Create client
let session_manager = Arc::new(McpSessionManager::new());
// Always check for stored tokens (from either pre-configured OAuth or DCR)
let secrets = get_secrets_store().await?;
let has_tokens = is_authenticated(&server, &secrets, &user_id).await;
let client = if has_tokens {
// We have stored tokens, use authenticated client
McpClient::new_authenticated(server.clone(), session_manager, secrets, user_id)
} else if server.requires_auth() {
// OAuth configured but no tokens - need to authenticate
println!();
println!(
" ✗ Not authenticated. Run 'ironclaw mcp auth {}' first.",
name
);
println!();
return Ok(());
} else {
// No OAuth and no tokens - try unauthenticated
McpClient::new_with_name(&server.name, &server.url)
};
// Test connection
match client.test_connection().await {
Ok(()) => {
println!(" ✓ Connection successful!");
println!();
// List tools
match client.list_tools().await {
Ok(tools) => {
println!(" Available tools ({}):", tools.len());
for tool in tools {
let approval = if tool.requires_approval() {
" [approval required]"
} else {
""
};
println!("{}{}", tool.name, approval);
if !tool.description.is_empty() {
// Truncate long descriptions
let desc = if tool.description.len() > 60 {
format!("{}...", &tool.description[..57])
} else {
tool.description.clone()
};
println!(" {}", desc);
}
}
}
Err(e) => {
println!(" ✗ Failed to list tools: {}", e);
}
}
}
Err(e) => {
let err_str = e.to_string();
// Check if server requires auth but we don't have valid tokens
if err_str.contains("401") || err_str.contains("requires authentication") {
if has_tokens {
// We had tokens but they failed - need to re-authenticate
println!(
" ✗ Authentication failed (token may be expired). Try re-authenticating:"
);
println!(" ironclaw mcp auth {}", name);
} else {
// No tokens - server requires auth
println!(" ✗ Server requires authentication.");
println!();
println!(" Run 'ironclaw mcp auth {}' to authenticate.", name);
}
} else {
println!(" ✗ Connection failed: {}", e);
}
}
}
println!();
Ok(())
}
/// Toggle server enabled/disabled state.
async fn toggle_server(name: String, enable: bool, disable: bool) -> anyhow::Result<()> {
let mut servers = load_mcp_servers().await?;
let server = servers
.get_mut(&name)
.ok_or_else(|| anyhow::anyhow!("Server '{}' not found", name))?;
let new_state = if enable {
true
} else if disable {
false
} else {
!server.enabled // Toggle if neither specified
};
server.enabled = new_state;
save_mcp_servers(&servers).await?;
let status = if new_state { "enabled" } else { "disabled" };
println!();
println!(" ✓ Server '{}' is now {}.", name, status);
println!();
Ok(())
}
/// Initialize and return the secrets store.
async fn get_secrets_store() -> anyhow::Result<Arc<dyn SecretsStore + Send + Sync>> {
let config = Config::from_env()?;
let master_key = config.secrets.master_key().ok_or_else(|| {
anyhow::anyhow!("SECRETS_MASTER_KEY not set. Run 'ironclaw setup' first or set it in .env")
})?;
let store = Store::new(&config.database).await?;
store.run_migrations().await?;
let crypto = SecretsCrypto::new(master_key.clone())?;
Ok(Arc::new(PostgresSecretsStore::new(
store.pool(),
Arc::new(crypto),
)))
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_mcp_command_parsing() {
// Just verify the command structure is valid
use clap::CommandFactory;
// Create a dummy parent command to test subcommand parsing
#[derive(clap::Parser)]
struct TestCli {
#[command(subcommand)]
cmd: McpCommand,
}
TestCli::command().debug_assert();
}
}
+7
View File
@@ -5,11 +5,14 @@
//! - Interactive setup wizard (`setup`)
//! - Managing configuration (`config list`, `config get`, `config set`)
//! - Managing WASM tools (`tool install`, `tool list`, `tool remove`)
//! - Managing MCP servers (`mcp add`, `mcp auth`, `mcp list`, `mcp test`)
mod config;
mod mcp;
mod tool;
pub use config::{ConfigCommand, run_config_command};
pub use mcp::{McpCommand, run_mcp_command};
pub use tool::{ToolCommand, run_tool_command};
use clap::{Parser, Subcommand};
@@ -72,6 +75,10 @@ pub enum Command {
/// Manage WASM tools
#[command(subcommand)]
Tool(ToolCommand),
/// Manage MCP servers (hosted tool providers)
#[command(subcommand)]
Mcp(McpCommand),
}
impl Cli {
+586 -1
View File
@@ -1,13 +1,18 @@
//! Tool management CLI commands.
//!
//! Commands for installing, listing, and removing WASM tools.
//! Commands for installing, listing, removing, and authenticating WASM tools.
use std::io::Write;
use std::path::{Path, PathBuf};
use std::process::Command as ProcessCommand;
use std::sync::Arc;
use clap::Subcommand;
use tokio::fs;
use crate::config::Config;
use crate::history::Store;
use crate::secrets::{CreateSecretParams, PostgresSecretsStore, SecretsCrypto, SecretsStore};
use crate::tools::wasm::{CapabilitiesFile, compute_binary_hash};
/// Default tools directory.
@@ -79,6 +84,20 @@ pub enum ToolCommand {
#[arg(short, long)]
dir: Option<PathBuf>,
},
/// Configure authentication for a tool
Auth {
/// Name of the tool
name: String,
/// Directory to look for tool (default: ~/.ironclaw/tools/)
#[arg(short, long)]
dir: Option<PathBuf>,
/// User ID for storing the secret (default: "default")
#[arg(short, long, default_value = "default")]
user: String,
},
}
/// Run a tool command.
@@ -96,6 +115,7 @@ pub async fn run_tool_command(cmd: ToolCommand) -> anyhow::Result<()> {
ToolCommand::List { dir, verbose } => list_tools(dir, verbose).await,
ToolCommand::Remove { name, dir } => remove_tool(name, dir).await,
ToolCommand::Info { name_or_path, dir } => show_tool_info(name_or_path, dir).await,
ToolCommand::Auth { name, dir, user } => auth_tool(name, dir, user).await,
}
}
@@ -658,6 +678,571 @@ fn print_capabilities_detail(caps: &CapabilitiesFile) {
}
}
/// Configure authentication for a tool.
async fn auth_tool(name: String, dir: Option<PathBuf>, user_id: String) -> anyhow::Result<()> {
let tools_dir = dir.unwrap_or_else(default_tools_dir);
let caps_path = tools_dir.join(format!("{}.capabilities.json", name));
if !caps_path.exists() {
anyhow::bail!(
"Tool '{}' not found or has no capabilities file at {}",
name,
caps_path.display()
);
}
// Parse capabilities
let content = fs::read_to_string(&caps_path).await?;
let caps = CapabilitiesFile::from_json(&content)
.map_err(|e| anyhow::anyhow!("Invalid capabilities file: {}", e))?;
// Check for auth section
let auth = caps.auth.ok_or_else(|| {
anyhow::anyhow!(
"Tool '{}' has no auth configuration.\n\
The tool may not require authentication, or auth setup is not defined.",
name
)
})?;
let display_name = auth.display_name.as_deref().unwrap_or(&name);
let header = format!("{} Authentication", display_name);
println!();
println!("╔════════════════════════════════════════════════════════════════╗");
println!("{:^62}", header);
println!("╚════════════════════════════════════════════════════════════════╝");
println!();
// Initialize secrets store
let config = Config::from_env()?;
let master_key = config.secrets.master_key().ok_or_else(|| {
anyhow::anyhow!("SECRETS_MASTER_KEY not set. Run 'ironclaw setup' first or set it in .env")
})?;
let store = Store::new(&config.database).await?;
store.run_migrations().await?;
let crypto = SecretsCrypto::new(master_key.clone())?;
let secrets_store = Arc::new(PostgresSecretsStore::new(store.pool(), Arc::new(crypto)));
// Check if already configured
let already_configured = secrets_store
.exists(&user_id, &auth.secret_name)
.await
.unwrap_or(false);
if already_configured {
println!(" {} is already configured.", display_name);
println!();
print!(" Replace existing credentials? [y/N]: ");
std::io::stdout().flush()?;
let mut input = String::new();
std::io::stdin().read_line(&mut input)?;
if !input.trim().eq_ignore_ascii_case("y") {
println!();
println!(" Keeping existing credentials.");
return Ok(());
}
println!();
}
// Check for environment variable
if let Some(ref env_var) = auth.env_var {
if let Ok(token) = std::env::var(env_var) {
if !token.is_empty() {
println!(" Found {} in environment.", env_var);
println!();
// Validate if endpoint is provided
if let Some(ref validation) = auth.validation_endpoint {
print!(" Validating token...");
std::io::stdout().flush()?;
match validate_token(&token, validation, &auth.secret_name).await {
Ok(()) => {
println!("");
}
Err(e) => {
println!("");
println!(" Validation failed: {}", e);
println!();
println!(" Falling back to manual entry...");
return auth_tool_manual(&secrets_store, &user_id, &auth).await;
}
}
}
// Save the token
save_token(&secrets_store, &user_id, &auth, &token).await?;
print_success(display_name);
return Ok(());
}
}
}
// Check for OAuth configuration
if let Some(ref oauth) = auth.oauth {
return auth_tool_oauth(&secrets_store, &user_id, &auth, oauth).await;
}
// Fall back to manual entry
auth_tool_manual(&secrets_store, &user_id, &auth).await
}
/// OAuth browser-based login flow.
async fn auth_tool_oauth(
store: &PostgresSecretsStore,
user_id: &str,
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 tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
use tokio::net::TcpListener;
let display_name = auth.display_name.as_deref().unwrap_or(&auth.secret_name);
// Get client_id from config or env
let client_id = oauth
.client_id
.clone()
.or_else(|| {
oauth
.client_id_env
.as_ref()
.and_then(|env| std::env::var(env).ok())
})
.ok_or_else(|| {
anyhow::anyhow!(
"OAuth client_id not configured.\n\
Set it in the capabilities file or via environment variable."
)
})?;
// Get client_secret if provided
let client_secret = oauth.client_secret.clone().or_else(|| {
oauth
.client_secret_env
.as_ref()
.and_then(|env| std::env::var(env).ok())
});
println!(" Starting OAuth authentication...");
println!();
// Find an available port for the callback
let mut listener = None;
let mut port = 0;
for p in 9876..=9886 {
match TcpListener::bind(format!("127.0.0.1:{}", p)).await {
Ok(l) => {
listener = Some(l);
port = p;
break;
}
Err(_) => continue,
}
}
let listener = listener.ok_or_else(|| anyhow::anyhow!("Could not find available port"))?;
let redirect_uri = format!("http://localhost:{}/callback", port);
// 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)
);
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)
));
}
println!(" Opening browser for {} login...", display_name);
println!();
if let Err(e) = open::that(&auth_url) {
println!(" Could not open browser: {}", e);
println!(" Please open this URL manually:");
println!(" {}", auth_url);
}
println!(" Waiting for authorization...");
// Wait for callback with timeout
let timeout = std::time::Duration::from_secs(300);
let code = tokio::time::timeout(timeout, async {
loop {
let (mut socket, _) = listener.accept().await?;
let mut reader = BufReader::new(&mut socket);
let mut request_line = String::new();
reader.read_line(&mut request_line).await?;
// Parse GET /callback?code=xxx HTTP/1.1
if let Some(path) = request_line.split_whitespace().nth(1) {
if path.starts_with("/callback") {
if let Some(query) = path.split('?').nth(1) {
for param in query.split('&') {
let parts: Vec<&str> = param.splitn(2, '=').collect();
if parts.len() == 2 && parts[0] == "code" {
let code = urlencoding::decode(parts[1])
.unwrap_or_else(|_| parts[1].into())
.into_owned();
// Send success response
let response = format!(
"HTTP/1.1 200 OK\r\n\
Content-Type: text/html\r\n\
\r\n\
<!DOCTYPE html><html><body style=\"font-family: sans-serif; \
display: flex; justify-content: center; align-items: center; \
height: 100vh; margin: 0; background: #191919; color: white;\">\
<div style=\"text-align: center;\">\
<h1>✓ {} Connected!</h1>\
<p>You can close this window.</p>\
</div></body></html>",
display_name
);
let _ = socket.write_all(response.as_bytes()).await;
let _ = socket.shutdown().await;
return Ok::<_, anyhow::Error>(code);
}
}
// Check for error
if query.contains("error=") {
let response =
"HTTP/1.1 400 Bad Request\r\n\r\nAuthorization denied";
let _ = socket.write_all(response.as_bytes()).await;
return Err(anyhow::anyhow!("Authorization denied by user"));
}
}
}
}
let response = "HTTP/1.1 404 Not Found\r\n\r\n";
let _ = socket.write_all(response.as_bytes()).await;
}
})
.await
.map_err(|_| anyhow::anyhow!("Timed out waiting for authorization"))??;
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
)
})?;
// Save the token
save_token(store, user_id, auth, access_token).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()));
println!();
println!("{} connected!", display_name);
if let Some(workspace) = workspace_name {
println!(" Workspace: {}", workspace);
}
println!();
println!(" The tool can now access the API.");
println!();
Ok(())
}
/// Manual token entry flow.
async fn auth_tool_manual(
store: &PostgresSecretsStore,
user_id: &str,
auth: &crate::tools::wasm::AuthCapabilitySchema,
) -> anyhow::Result<()> {
let display_name = auth.display_name.as_deref().unwrap_or(&auth.secret_name);
// Show instructions
if let Some(ref instructions) = auth.instructions {
println!(" Setup instructions:");
println!();
for line in instructions.lines() {
println!(" {}", line);
}
println!();
}
// Offer to open setup URL
if let Some(ref url) = auth.setup_url {
print!(" Press Enter to open setup page (or 's' to skip): ");
std::io::stdout().flush()?;
let mut input = String::new();
std::io::stdin().read_line(&mut input)?;
if !input.trim().eq_ignore_ascii_case("s") {
if let Err(e) = open::that(url) {
println!(" Could not open browser: {}", e);
println!(" Please open manually: {}", url);
} else {
println!(" Opening browser...");
}
}
println!();
}
// Show token hint
if let Some(ref hint) = auth.token_hint {
println!(" Token format: {}", hint);
println!();
}
// Prompt for token
print!(" Paste your token: ");
std::io::stdout().flush()?;
let token = read_hidden_input()?;
println!();
if token.is_empty() {
println!(" No token provided. Aborting.");
return Ok(());
}
// Validate if endpoint is provided
if let Some(ref validation) = auth.validation_endpoint {
print!(" Validating token...");
std::io::stdout().flush()?;
match validate_token(&token, validation, &auth.secret_name).await {
Ok(()) => {
println!("");
}
Err(e) => {
println!("");
println!(" Validation failed: {}", e);
println!();
print!(" Save anyway? [y/N]: ");
std::io::stdout().flush()?;
let mut confirm = String::new();
std::io::stdin().read_line(&mut confirm)?;
if !confirm.trim().eq_ignore_ascii_case("y") {
println!(" Aborting.");
return Ok(());
}
}
}
}
// Save the token
save_token(store, user_id, auth, &token).await?;
print_success(display_name);
Ok(())
}
/// Read input with hidden characters.
fn read_hidden_input() -> anyhow::Result<String> {
use crossterm::{
event::{self, Event, KeyCode, KeyModifiers},
terminal,
};
let mut input = String::new();
terminal::enable_raw_mode()?;
loop {
if let Event::Key(key_event) = event::read()? {
match key_event.code {
KeyCode::Enter => {
break;
}
KeyCode::Backspace => {
if !input.is_empty() {
input.pop();
print!("\x08 \x08");
std::io::stdout().flush()?;
}
}
KeyCode::Char('c') if key_event.modifiers.contains(KeyModifiers::CONTROL) => {
terminal::disable_raw_mode()?;
return Err(anyhow::anyhow!("Interrupted"));
}
KeyCode::Char(c) => {
input.push(c);
print!("*");
std::io::stdout().flush()?;
}
_ => {}
}
}
}
terminal::disable_raw_mode()?;
Ok(input)
}
/// Validate a token against the validation endpoint.
async fn validate_token(
token: &str,
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
}
))
}
}
/// Save token to secrets store.
async fn save_token(
store: &PostgresSecretsStore,
user_id: &str,
auth: &crate::tools::wasm::AuthCapabilitySchema,
token: &str,
) -> anyhow::Result<()> {
let mut params = CreateSecretParams::new(&auth.secret_name, token);
if let Some(ref provider) = auth.provider {
params = params.with_provider(provider);
}
store
.create(user_id, params)
.await
.map_err(|e| anyhow::anyhow!("Failed to save token: {}", e))?;
Ok(())
}
/// Print success message.
fn print_success(display_name: &str) {
println!();
println!("{} connected!", display_name);
println!();
println!(" The tool can now access the API.");
println!();
}
#[cfg(test)]
mod tests {
use super::*;