feat(extensions): unify auth and configure into single entrypoint (#677)

* feat(extensions): unify auth and configure into single entrypoint

Refactors the extension lifecycle to eliminate the divergence between
chat and gateway paths that caused Telegram setup via chat to fail
(missing webhook secret auto-generation, no token validation).

Key changes:
- Rename save_setup_secrets() → configure(): single entrypoint for
  providing secrets to any extension (WasmChannel, WasmTool, MCP).
  Validates, stores, auto-generates, and activates.
- Add configure_token(): convenience wrapper for single-token callers
  (chat auth card, WebSocket, agent auth mode).
- Refactor auth() to pure status check: remove token parameter,
  delete token-storing branches from auth_mcp/auth_wasm_tool,
  rename auth_wasm_channel → auth_wasm_channel_status.
- Add ConfigureResult/MissingSecret types for structured responses.
- Replace hardcoded Telegram token validation with generic
  validation_endpoint from capabilities.json.
- Update all callers (9 files) to use the new interface.

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

* fix: use ValidationFailed error variant instead of string matching

Replace brittle msg.contains("Invalid token") checks with a proper
ExtensionError::ValidationFailed variant. configure() now returns
this variant for token validation failures, and callers match on it
directly instead of parsing error message strings.

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

* fix: address review — SSRF protection, error typing, missing-secret selection, WS auth

1. SSRF: call validate_fetch_url() before validation_endpoint HTTP request
2. Transport errors map to ExtensionError::Other (not ValidationFailed)
3. configure_token() picks first *missing* secret, not first non-optional
4. WebSocket error path re-emits AuthRequired on ValidationFailed

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

* test: add regression tests for extension lifecycle refactoring

- test_configure_token_picks_first_missing_secret: verifies multi-secret
  channels can be configured one secret at a time (commit ce106f4)
- test_auth_is_read_only_for_wasm_channel: verifies auth() has no side
  effects and doesn't store secrets (commit 47f8eb6)
- test_validation_failed_is_distinct_error_variant: verifies the typed
  error variant can be pattern-matched (commit a318161)

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* fix: address review comments — activation dispatch, dead code, caps consolidation

- Fix configure() fallthrough bug: dispatch activation by ExtensionKind
  instead of unconditionally calling activate_wasm_channel() for all
  non-WasmTool types (MCP servers and channel relays now use their
  correct activation methods)
- Remove dead MissingSecret struct and missing_secrets field (never
  populated, flagged by reviewer)
- Consolidate capabilities file parsing in configure(): parse once
  and reuse for allowed names, validation_endpoint, and auto-generation
- Fix auth() doc comment: note MCP OAuth side effects
- Fix stale save_setup_secrets reference in server.rs comment
- Add regression test for activation dispatch bug

Co-Authored-By: Claude Opus 4.6 <[email protected]>

---------

Co-authored-by: Claude Opus 4.6 (1M context) <[email protected]>
This commit is contained in:
Henry Park
2026-03-11 16:01:41 -07:00
committed by GitHub
co-authored by Claude Opus 4.6
parent 28a22f2a59
commit 19d9562b4f
8 changed files with 555 additions and 368 deletions
@@ -20,7 +20,8 @@
"optional": false
}
],
"setup_url": "https://t.me/BotFather"
"setup_url": "https://t.me/BotFather",
"validation_endpoint": "https://api.telegram.org/bot{telegram_bot_token}/getMe"
},
"capabilities": {
"http": {
+38 -82
View File
@@ -1355,100 +1355,56 @@ impl Agent {
None => return Ok(Some("Extension manager not available.".to_string())),
};
match ext_mgr.auth(&pending.extension_name, Some(token)).await {
Ok(result) if result.is_authenticated() => {
tracing::info!(
"Extension '{}' authenticated via auth mode",
pending.extension_name
);
// Auto-activate so tools are available immediately after auth
match ext_mgr.activate(&pending.extension_name).await {
Ok(activate_result) => {
let tool_count = activate_result.tools_loaded.len();
let tool_list = if activate_result.tools_loaded.is_empty() {
String::new()
} else {
format!("\n\nTools: {}", activate_result.tools_loaded.join(", "))
};
let msg = format!(
"{} authenticated and activated ({} tools loaded).{}",
pending.extension_name, tool_count, tool_list
);
let _ = self
.channels
.send_status(
&message.channel,
StatusUpdate::AuthCompleted {
extension_name: pending.extension_name.clone(),
success: true,
message: msg.clone(),
},
&message.metadata,
)
.await;
Ok(Some(msg))
}
Err(e) => {
tracing::warn!(
"Extension '{}' authenticated but activation failed: {}",
pending.extension_name,
e
);
let msg = format!(
"{} authenticated successfully, but activation failed: {}. \
Try activating manually.",
pending.extension_name, e
);
let _ = self
.channels
.send_status(
&message.channel,
StatusUpdate::AuthCompleted {
extension_name: pending.extension_name.clone(),
success: true,
message: msg.clone(),
},
&message.metadata,
)
.await;
Ok(Some(msg))
}
}
}
match ext_mgr
.configure_token(&pending.extension_name, token)
.await
{
Ok(result) => {
// Invalid token, re-enter auth mode
{
let mut sess = session.lock().await;
if let Some(thread) = sess.threads.get_mut(&thread_id) {
thread.enter_auth_mode(pending.extension_name.clone());
}
}
let msg = result
.instructions()
.map(String::from)
.unwrap_or_else(|| "Invalid token. Please try again.".to_string());
// Re-emit AuthRequired so web UI re-shows the card
tracing::info!(
"Extension '{}' configured via auth mode: {}",
pending.extension_name,
result.message
);
let _ = self
.channels
.send_status(
&message.channel,
StatusUpdate::AuthRequired {
StatusUpdate::AuthCompleted {
extension_name: pending.extension_name.clone(),
instructions: Some(msg.clone()),
auth_url: result.auth_url().map(String::from),
setup_url: result.setup_url().map(String::from),
success: true,
message: result.message.clone(),
},
&message.metadata,
)
.await;
Ok(Some(msg))
Ok(Some(result.message))
}
Err(e) => {
let msg = format!(
"Authentication failed for {}: {}",
pending.extension_name, e
);
let msg = e.to_string();
// Token validation errors: re-enter auth mode and re-prompt
if matches!(e, crate::extensions::ExtensionError::ValidationFailed(_)) {
{
let mut sess = session.lock().await;
if let Some(thread) = sess.threads.get_mut(&thread_id) {
thread.enter_auth_mode(pending.extension_name.clone());
}
}
let _ = self
.channels
.send_status(
&message.channel,
StatusUpdate::AuthRequired {
extension_name: pending.extension_name.clone(),
instructions: Some(msg.clone()),
auth_url: None,
setup_url: None,
},
&message.metadata,
)
.await;
return Ok(Some(msg));
}
// Infrastructure errors
let _ = self
.channels
.send_status(
+24 -40
View File
@@ -145,49 +145,33 @@ pub async fn chat_auth_token_handler(
"Extension manager not available".to_string(),
))?;
let result = ext_mgr
.auth(&req.extension_name, Some(&req.token))
match ext_mgr
.configure_token(&req.extension_name, &req.token)
.await
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
{
Ok(result) => {
clear_auth_mode(&state).await;
if result.is_authenticated() {
// Auto-activate so tools are available immediately
let msg = match ext_mgr.activate(&req.extension_name).await {
Ok(r) => format!(
"{} authenticated ({} tools loaded)",
req.extension_name,
r.tools_loaded.len()
),
Err(e) => format!(
"{} authenticated but activation failed: {}",
req.extension_name, e
),
};
state.sse.broadcast(SseEvent::AuthCompleted {
extension_name: req.extension_name.clone(),
success: true,
message: result.message.clone(),
});
// Clear auth mode on the active thread
clear_auth_mode(&state).await;
state.sse.broadcast(SseEvent::AuthCompleted {
extension_name: req.extension_name,
success: true,
message: msg.clone(),
});
Ok(Json(ActionResponse::ok(msg)))
} else {
// Re-emit auth_required for retry
state.sse.broadcast(SseEvent::AuthRequired {
extension_name: req.extension_name.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()
.map(String::from)
.unwrap_or_else(|| "Invalid token".to_string()),
)))
Ok(Json(ActionResponse::ok(result.message)))
}
Err(e) => {
let msg = e.to_string();
if matches!(e, crate::extensions::ExtensionError::ValidationFailed(_)) {
state.sse.broadcast(SseEvent::AuthRequired {
extension_name: req.extension_name.clone(),
instructions: Some(msg.clone()),
auth_url: None,
setup_url: None,
});
}
Ok(Json(ActionResponse::fail(msg)))
}
}
}
+31 -45
View File
@@ -1062,49 +1062,35 @@ async fn chat_auth_token_handler(
"Extension manager not available".to_string(),
))?;
let result = ext_mgr
.auth(&req.extension_name, Some(&req.token))
match ext_mgr
.configure_token(&req.extension_name, &req.token)
.await
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
{
Ok(result) => {
// Clear auth mode on the active thread
clear_auth_mode(&state).await;
if result.is_authenticated() {
// Auto-activate so tools are available immediately
let msg = match ext_mgr.activate(&req.extension_name).await {
Ok(r) => format!(
"{} authenticated ({} tools loaded)",
req.extension_name,
r.tools_loaded.len()
),
Err(e) => format!(
"{} authenticated but activation failed: {}",
req.extension_name, e
),
};
state.sse.broadcast(SseEvent::AuthCompleted {
extension_name: req.extension_name.clone(),
success: true,
message: result.message.clone(),
});
// Clear auth mode on the active thread
clear_auth_mode(&state).await;
state.sse.broadcast(SseEvent::AuthCompleted {
extension_name: req.extension_name,
success: true,
message: msg.clone(),
});
Ok(Json(ActionResponse::ok(msg)))
} else {
// Re-emit auth_required for retry
state.sse.broadcast(SseEvent::AuthRequired {
extension_name: req.extension_name.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()
.map(String::from)
.unwrap_or_else(|| "Invalid token".to_string()),
)))
Ok(Json(ActionResponse::ok(result.message)))
}
Err(e) => {
let msg = e.to_string();
// Re-emit auth_required for retry on validation errors
if matches!(e, crate::extensions::ExtensionError::ValidationFailed(_)) {
state.sse.broadcast(SseEvent::AuthRequired {
extension_name: req.extension_name.clone(),
instructions: Some(msg.clone()),
auth_url: None,
setup_url: None,
});
}
Ok(Json(ActionResponse::fail(msg)))
}
}
}
@@ -1853,7 +1839,7 @@ async fn extensions_install_handler(
// 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 {
match ext_mgr.auth(&req.name).await {
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().map(String::from);
@@ -1882,9 +1868,9 @@ async fn extensions_activate_handler(
// 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.
// Initial OAuth setup is triggered via configure.
let mut resp = ActionResponse::ok(result.message);
if let Ok(auth_result) = ext_mgr.auth(&name, None).await
if let Ok(auth_result) = ext_mgr.auth(&name).await
&& auth_result.auth_url().is_some()
{
resp.auth_url = auth_result.auth_url().map(String::from);
@@ -1902,7 +1888,7 @@ async fn extensions_activate_handler(
}
// Activation failed due to auth; try authenticating first.
match ext_mgr.auth(&name, None).await {
match ext_mgr.auth(&name).await {
Ok(auth_result) if auth_result.is_authenticated() => {
// Auth succeeded, retry activation.
match ext_mgr.activate(&name).await {
@@ -2109,7 +2095,7 @@ async fn extensions_setup_submit_handler(
"Extension manager not available (secrets store required)".to_string(),
))?;
match ext_mgr.save_setup_secrets(&name, &req.secrets).await {
match ext_mgr.configure(&name, &req.secrets).await {
Ok(result) => {
// Broadcast auth_completed so the chat UI can dismiss any in-progress
// auth card or setup modal that was triggered by tool_auth/tool_activate.
+15 -27
View File
@@ -255,43 +255,31 @@ async fn handle_client_message(
token,
} => {
if let Some(ref ext_mgr) = state.extension_manager {
match ext_mgr.auth(&extension_name, Some(&token)).await {
Ok(result) if result.is_authenticated() => {
let msg = match ext_mgr.activate(&extension_name).await {
Ok(r) => format!(
"{} authenticated ({} tools loaded)",
extension_name,
r.tools_loaded.len()
),
Err(e) => format!(
"{} authenticated but activation failed: {}",
extension_name, e
),
};
match ext_mgr.configure_token(&extension_name, &token).await {
Ok(result) => {
crate::channels::web::server::clear_auth_mode(state).await;
state
.sse
.broadcast(crate::channels::web::types::SseEvent::AuthCompleted {
extension_name,
success: true,
message: msg,
});
}
Ok(result) => {
state
.sse
.broadcast(crate::channels::web::types::SseEvent::AuthRequired {
extension_name,
instructions: result.instructions().map(String::from),
auth_url: result.auth_url().map(String::from),
setup_url: result.setup_url().map(String::from),
message: result.message,
});
}
Err(e) => {
let msg = format!("Auth failed: {}", e);
if matches!(e, crate::extensions::ExtensionError::ValidationFailed(_)) {
state.sse.broadcast(
crate::channels::web::types::SseEvent::AuthRequired {
extension_name: extension_name.clone(),
instructions: Some(msg.clone()),
auth_url: None,
setup_url: None,
},
);
}
let _ = direct_tx
.send(WsServerMessage::Error {
message: format!("Auth failed: {}", e),
})
.send(WsServerMessage::Error { message: msg })
.await;
}
}
+426 -171
View File
@@ -17,9 +17,9 @@ use crate::channels::wasm::{
use crate::extensions::discovery::OnlineDiscovery;
use crate::extensions::registry::ExtensionRegistry;
use crate::extensions::{
ActivateResult, AuthResult, ExtensionError, ExtensionKind, ExtensionSource, InstallResult,
InstalledExtension, RegistryEntry, ResultSource, SearchResult, ToolAuthState, UpgradeOutcome,
UpgradeResult,
ActivateResult, AuthResult, ConfigureResult, ExtensionError, ExtensionKind, ExtensionSource,
InstallResult, InstalledExtension, RegistryEntry, ResultSource, SearchResult, ToolAuthState,
UpgradeOutcome, UpgradeResult,
};
use crate::hooks::HookRegistry;
use crate::pairing::PairingStore;
@@ -56,16 +56,6 @@ struct ChannelRuntimeState {
wasm_channel_owner_ids: std::collections::HashMap<String, i64>,
}
/// Result of saving setup secrets and attempting activation.
pub struct SetupResult {
/// Human-readable status message.
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.
pub struct ExtensionManager {
registry: ExtensionRegistry,
@@ -440,12 +430,11 @@ impl ExtensionManager {
Err(err)
}
/// Authenticate an installed extension.
pub async fn auth(
&self,
name: &str,
token: Option<&str>,
) -> Result<AuthResult, ExtensionError> {
/// Check auth status for an installed extension.
///
/// Read-only for WASM extensions; may initiate OAuth for MCP servers.
/// To provide secrets, use [`configure()`] instead.
pub async fn auth(&self, name: &str) -> Result<AuthResult, ExtensionError> {
// Clean up expired pending auths
self.cleanup_expired_auths().await;
@@ -453,10 +442,10 @@ impl ExtensionManager {
let kind = self.determine_installed_kind(name).await?;
match kind {
ExtensionKind::McpServer => self.auth_mcp(name, token).await,
ExtensionKind::WasmTool => self.auth_wasm_tool(name, token).await,
ExtensionKind::WasmChannel => self.auth_wasm_channel(name, token).await,
ExtensionKind::ChannelRelay => self.auth_channel_relay(name, token).await,
ExtensionKind::McpServer => self.auth_mcp(name).await,
ExtensionKind::WasmTool => self.auth_wasm_tool(name).await,
ExtensionKind::WasmChannel => self.auth_wasm_channel_status(name).await,
ExtensionKind::ChannelRelay => self.auth_channel_relay(name).await,
}
}
@@ -1684,30 +1673,12 @@ impl ExtensionManager {
})
}
async fn auth_mcp(
&self,
name: &str,
token: Option<&str>,
) -> Result<AuthResult, ExtensionError> {
async fn auth_mcp(&self, name: &str) -> Result<AuthResult, ExtensionError> {
let server = self
.get_mcp_server(name)
.await
.map_err(|e| ExtensionError::NotInstalled(e.to_string()))?;
// If a token was provided directly, store it and we're done.
if let Some(token_value) = token {
let secret_name = server.token_secret_name();
let params =
CreateSecretParams::new(&secret_name, token_value).with_provider(name.to_string());
self.secrets
.create(&self.user_id, params)
.await
.map_err(|e| ExtensionError::AuthFailed(e.to_string()))?;
tracing::info!("MCP server '{}' authenticated via manual token", name);
return Ok(AuthResult::authenticated(name, ExtensionKind::McpServer));
}
// Check if already authenticated
if is_authenticated(&server, &self.secrets, &self.user_id).await {
return Ok(AuthResult::authenticated(name, ExtensionKind::McpServer));
@@ -1820,11 +1791,7 @@ impl ExtensionManager {
))
}
async fn auth_wasm_tool(
&self,
name: &str,
token: Option<&str>,
) -> Result<AuthResult, ExtensionError> {
async fn auth_wasm_tool(&self, name: &str) -> Result<AuthResult, ExtensionError> {
// Read the capabilities file to get auth config
let cap_path = self
.wasm_tools_dir
@@ -1895,18 +1862,6 @@ impl ExtensionManager {
// Fall through to OAuth branch for scope expansion
}
// If a token was provided, store it
if let Some(token_value) = token {
let params = CreateSecretParams::new(&auth.secret_name, token_value)
.with_provider(name.to_string());
self.secrets
.create(&self.user_id, params)
.await
.map_err(|e| ExtensionError::AuthFailed(e.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.
@@ -2554,11 +2509,8 @@ impl ExtensionManager {
}
}
async fn auth_wasm_channel(
&self,
name: &str,
token: Option<&str>,
) -> Result<AuthResult, ExtensionError> {
/// Check auth status for a WASM channel (read-only).
async fn auth_wasm_channel_status(&self, name: &str) -> Result<AuthResult, ExtensionError> {
let cap_path = self
.wasm_channels_dir
.join(format!("{}.capabilities.json", name));
@@ -2577,7 +2529,6 @@ impl ExtensionManager {
let cap_file = crate::channels::wasm::ChannelCapabilitiesFile::from_bytes(&cap_bytes)
.map_err(|e| ExtensionError::Other(e.to_string()))?;
// Get required secrets from the setup section
let required_secrets = &cap_file.setup.required_secrets;
if required_secrets.is_empty() {
return Ok(AuthResult::no_auth_required(
@@ -2586,7 +2537,7 @@ impl ExtensionManager {
));
}
// Find the first non-optional secret that isn't yet stored
// Find non-optional secrets that aren't yet stored
let mut missing = Vec::new();
for secret in required_secrets {
if secret.optional {
@@ -2606,31 +2557,6 @@ impl ExtensionManager {
return Ok(AuthResult::authenticated(name, ExtensionKind::WasmChannel));
}
// If a token was provided, store it for the first missing secret
if let Some(token_value) = token {
let secret = &missing[0];
let params =
CreateSecretParams::new(&secret.name, token_value).with_provider(name.to_string());
self.secrets
.create(&self.user_id, params)
.await
.map_err(|e| ExtensionError::AuthFailed(e.to_string()))?;
// Check if there are more missing secrets
if missing.len() <= 1 {
return Ok(AuthResult::authenticated(name, ExtensionKind::WasmChannel));
}
// More secrets needed; prompt for the next one
let next = &missing[1];
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::awaiting_token(
@@ -3218,11 +3144,7 @@ impl ExtensionManager {
/// For Slack: initiates OAuth flow (redirect-based).
/// For Telegram: accepts a bot token, registers it with channel-relay,
/// and stores the returned stream token.
async fn auth_channel_relay(
&self,
name: &str,
_token: Option<&str>,
) -> Result<AuthResult, ExtensionError> {
async fn auth_channel_relay(&self, name: &str) -> Result<AuthResult, ExtensionError> {
// Check if already authenticated (stream token exists)
let token_key = format!("relay:{}:stream_token", name);
if self
@@ -3525,16 +3447,26 @@ impl ExtensionManager {
/// Save setup secrets for an extension, validating names against the capabilities schema.
///
/// After saving, attempts to hot-activate the channel. Returns a [`SetupResult`]
/// indicating whether activation succeeded (so the frontend can show appropriate UI).
pub async fn save_setup_secrets(
/// Configure secrets for an extension: validate, store, auto-generate, and activate.
///
/// This is the single entrypoint for providing secrets to any extension.
/// Both the chat auth flow and the Extensions tab setup form call this method.
///
/// - Validates tokens against `validation_endpoint` (if declared in capabilities)
/// - Stores secrets in the encrypted secrets store
/// - Auto-generates missing secrets (e.g., webhook keys)
/// - Activates the extension after configuration
pub async fn configure(
&self,
name: &str,
secrets: &std::collections::HashMap<String, String>,
) -> Result<SetupResult, ExtensionError> {
) -> Result<ConfigureResult, ExtensionError> {
let kind = self.determine_installed_kind(name).await?;
// Load allowed secret names from the extension's capabilities file
// Load allowed secret names and (for channels) the parsed capabilities file.
// The capabilities file is parsed once here and reused for validation_endpoint
// and auto-generation below, avoiding redundant I/O + JSON parsing.
let mut channel_cap_file: Option<crate::channels::wasm::ChannelCapabilitiesFile> = None;
let allowed: std::collections::HashSet<String> = match kind {
ExtensionKind::WasmChannel => {
let cap_path = self
@@ -3552,45 +3484,71 @@ impl ExtensionManager {
let cap_file =
crate::channels::wasm::ChannelCapabilitiesFile::from_bytes(&cap_bytes)
.map_err(|e| ExtensionError::Other(e.to_string()))?;
cap_file
let names = cap_file
.setup
.required_secrets
.iter()
.map(|s| s.name.clone())
.collect()
.collect();
channel_cap_file = Some(cap_file);
names
}
ExtensionKind::WasmTool => {
let cap_file = self.load_tool_capabilities(name).await.ok_or_else(|| {
ExtensionError::Other(format!("Capabilities file not found for '{}'", name))
})?;
match cap_file.setup {
Some(s) => s.required_secrets.iter().map(|s| s.name.clone()).collect(),
None => {
return Err(ExtensionError::Other(format!(
"Tool '{}' has no setup schema — no secrets to configure",
name
)));
}
let mut names: std::collections::HashSet<String> = std::collections::HashSet::new();
if let Some(ref s) = cap_file.setup {
names.extend(s.required_secrets.iter().map(|s| s.name.clone()));
}
// Also allow storing the auth token secret directly
if let Some(ref auth) = cap_file.auth {
names.insert(auth.secret_name.clone());
}
if names.is_empty() {
return Err(ExtensionError::Other(format!(
"Tool '{}' has no setup or auth schema — no secrets to configure",
name
)));
}
names
}
_ => {
return Err(ExtensionError::Other(
"Setup is only supported for WASM channels and tools".to_string(),
));
ExtensionKind::McpServer => {
let server = self
.get_mcp_server(name)
.await
.map_err(|e| ExtensionError::NotInstalled(e.to_string()))?;
let mut names = std::collections::HashSet::new();
names.insert(server.token_secret_name());
names
}
ExtensionKind::ChannelRelay => {
let mut names = std::collections::HashSet::new();
names.insert(format!("relay:{}:stream_token", name));
names
}
};
// For Telegram, validate the bot token against the API before storing it.
// This catches bad tokens immediately (both on first setup and reconfigure),
// before the channel activates and potentially shows as active with a bad token.
if name == "telegram"
&& let Some(token_value) = secrets.get("telegram_bot_token")
// Validate secrets against the validation_endpoint if declared in capabilities.
// The endpoint URL template uses {secret_name} placeholders that are
// substituted with the provided secret value before making the request.
if let Some(ref cap_file) = channel_cap_file
&& let Some(ref endpoint_template) = cap_file.setup.validation_endpoint
&& let Some(secret_def) = cap_file
.setup
.required_secrets
.iter()
.find(|s| !s.optional && secrets.contains_key(&s.name))
&& let Some(token_value) = secrets.get(&secret_def.name)
{
let token = token_value.trim();
if !token.is_empty() {
let encoded_token =
let encoded =
url::form_urlencoded::byte_serialize(token.as_bytes()).collect::<String>();
let url = format!("https://api.telegram.org/bot{}/getMe", encoded_token);
let url = endpoint_template.replace(&format!("{{{}}}", secret_def.name), &encoded);
// SSRF defense: block private IPs, localhost, cloud metadata endpoints
crate::tools::builtin::skill_tools::validate_fetch_url(&url)
.map_err(|e| ExtensionError::Other(format!("SSRF blocked: {}", e)))?;
let resp = reqwest::Client::builder()
.timeout(std::time::Duration::from_secs(10))
.build()
@@ -3598,12 +3556,13 @@ impl ExtensionManager {
.get(&url)
.send()
.await
// Transport errors are infrastructure failures, not token issues
.map_err(|e| {
ExtensionError::Other(format!("Failed to validate bot token: {}", e))
ExtensionError::Other(format!("Token validation request failed: {}", e))
})?;
if !resp.status().is_success() {
return Err(ExtensionError::Other(format!(
"Invalid bot token (Telegram API returned {})",
return Err(ExtensionError::ValidationFailed(format!(
"Invalid token (API returned {})",
resp.status()
)));
}
@@ -3630,43 +3589,34 @@ impl ExtensionManager {
}
// Auto-generate any missing secrets (channel-only feature)
if kind == ExtensionKind::WasmChannel {
let cap_path = self
.wasm_channels_dir
.join(format!("{}.capabilities.json", name));
if let Ok(cap_bytes) = tokio::fs::read(&cap_path).await
&& let Ok(cap_file) =
crate::channels::wasm::ChannelCapabilitiesFile::from_bytes(&cap_bytes)
{
for secret_def in &cap_file.setup.required_secrets {
if let Some(ref auto_gen) = secret_def.auto_generate {
let already_provided = secrets
.get(&secret_def.name)
.is_some_and(|v| !v.trim().is_empty());
let already_stored = self
.secrets
.exists(&self.user_id, &secret_def.name)
if let Some(ref cap_file) = channel_cap_file {
for secret_def in &cap_file.setup.required_secrets {
if let Some(ref auto_gen) = secret_def.auto_generate {
let already_provided = secrets
.get(&secret_def.name)
.is_some_and(|v| !v.trim().is_empty());
let already_stored = self
.secrets
.exists(&self.user_id, &secret_def.name)
.await
.unwrap_or(false);
if !already_provided && !already_stored {
use rand::RngCore;
use rand::rngs::OsRng;
let mut bytes = vec![0u8; auto_gen.length];
OsRng.fill_bytes(&mut bytes);
let hex_value: String = bytes.iter().map(|b| format!("{b:02x}")).collect();
let params = CreateSecretParams::new(&secret_def.name, &hex_value)
.with_provider(name.to_string());
self.secrets
.create(&self.user_id, params)
.await
.unwrap_or(false);
if !already_provided && !already_stored {
use rand::RngCore;
use rand::rngs::OsRng;
let mut bytes = vec![0u8; auto_gen.length];
OsRng.fill_bytes(&mut bytes);
let hex_value: String =
bytes.iter().map(|b| format!("{b:02x}")).collect();
let params = CreateSecretParams::new(&secret_def.name, &hex_value)
.with_provider(name.to_string());
self.secrets
.create(&self.user_id, params)
.await
.map_err(|e| ExtensionError::AuthFailed(e.to_string()))?;
tracing::info!(
"Auto-generated secret '{}' for channel '{}'",
secret_def.name,
name
);
}
.map_err(|e| ExtensionError::AuthFailed(e.to_string()))?;
tracing::info!(
"Auto-generated secret '{}' for channel '{}'",
secret_def.name,
name
);
}
}
}
@@ -3703,7 +3653,9 @@ impl ExtensionManager {
// 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 {
// Box::pin breaks the async recursion cycle:
// auth() → auth_wasm_tool() → (OAuth) → configure() → auth()
if let Ok(auth_result) = Box::pin(self.auth(name)).await {
auth_url = auth_result.auth_url().map(String::from);
}
let message = if auth_url.is_some() {
@@ -3717,7 +3669,7 @@ impl ExtensionManager {
name, result.message
)
};
return Ok(SetupResult {
return Ok(ConfigureResult {
message,
activated: true,
auth_url,
@@ -3729,7 +3681,7 @@ impl ExtensionManager {
name,
e
);
return Ok(SetupResult {
return Ok(ConfigureResult {
message: format!("Configuration saved for '{}'.", name),
activated: false,
auth_url: None,
@@ -3738,14 +3690,29 @@ impl ExtensionManager {
}
}
// Try to hot-activate the channel now that secrets are saved
match self.activate_wasm_channel(name).await {
// Activate the extension now that secrets are saved.
// Dispatch by kind — WasmTool was already handled above with an early return.
let activate_result = match kind {
ExtensionKind::WasmChannel => self.activate_wasm_channel(name).await,
ExtensionKind::McpServer => self.activate_mcp(name).await,
ExtensionKind::ChannelRelay => self.activate_channel_relay(name).await,
ExtensionKind::WasmTool => {
// WasmTool is handled above and returns early; this branch is unreachable.
return Ok(ConfigureResult {
message: format!("Configuration saved for '{}'.", name),
activated: false,
auth_url: None,
});
}
};
match activate_result {
Ok(result) => {
self.activation_errors.write().await.remove(name);
self.broadcast_extension_status(name, "active", None).await;
Ok(SetupResult {
Ok(ConfigureResult {
message: format!(
"Configuration saved and channel '{}' activated. {}",
"Configuration saved and '{}' activated. {}",
name, result.message
),
activated: true,
@@ -3755,9 +3722,9 @@ impl ExtensionManager {
Err(e) => {
let error_msg = e.to_string();
tracing::warn!(
channel = name,
extension = name,
error = %e,
"Saved configuration but hot-activation failed"
"Saved configuration but activation failed"
);
self.activation_errors
.write()
@@ -3765,7 +3732,7 @@ impl ExtensionManager {
.insert(name.to_string(), error_msg.clone());
self.broadcast_extension_status(name, "failed", Some(&error_msg))
.await;
Ok(SetupResult {
Ok(ConfigureResult {
message: format!(
"Configuration saved for '{}'. Activation failed: {}",
name, e
@@ -3777,6 +3744,118 @@ impl ExtensionManager {
}
}
/// Convenience wrapper: configure a single token for an extension.
///
/// Determines the primary secret name from the extension's capabilities,
/// then delegates to [`configure()`]. Use this when the caller only has
/// a bare token value (e.g., from the chat auth card or WebSocket auth).
pub async fn configure_token(
&self,
name: &str,
token: &str,
) -> Result<ConfigureResult, ExtensionError> {
let kind = self.determine_installed_kind(name).await?;
let secret_name = match kind {
ExtensionKind::WasmChannel => {
let cap_path = self
.wasm_channels_dir
.join(format!("{}.capabilities.json", name));
let cap_bytes = tokio::fs::read(&cap_path)
.await
.map_err(|e| ExtensionError::Other(e.to_string()))?;
let cap_file =
crate::channels::wasm::ChannelCapabilitiesFile::from_bytes(&cap_bytes)
.map_err(|e| ExtensionError::Other(e.to_string()))?;
// Pick the first *missing* non-optional secret so re-configure
// of a second secret works for multi-secret channels.
let mut target = None;
for s in &cap_file.setup.required_secrets {
if s.optional {
continue;
}
if !self
.secrets
.exists(&self.user_id, &s.name)
.await
.unwrap_or(false)
{
target = Some(s.name.clone());
break;
}
}
// Fall back to first non-optional if all exist (overwrite)
target
.or_else(|| {
cap_file
.setup
.required_secrets
.iter()
.find(|s| !s.optional)
.map(|s| s.name.clone())
})
.ok_or_else(|| {
ExtensionError::Other(format!("Channel '{}' has no required secrets", name))
})?
}
ExtensionKind::WasmTool => {
let cap = self.load_tool_capabilities(name).await.ok_or_else(|| {
ExtensionError::Other(format!("Capabilities not found for '{}'", name))
})?;
// Prefer auth secret, then first missing setup secret
if let Some(ref auth) = cap.auth {
if !self
.secrets
.exists(&self.user_id, &auth.secret_name)
.await
.unwrap_or(false)
{
auth.secret_name.clone()
} else if let Some(ref setup) = cap.setup {
// Auth secret exists, find first missing setup secret
let mut found = None;
for s in &setup.required_secrets {
if !self
.secrets
.exists(&self.user_id, &s.name)
.await
.unwrap_or(false)
{
found = Some(s.name.clone());
break;
}
}
found.unwrap_or_else(|| auth.secret_name.clone())
} else {
auth.secret_name.clone()
}
} else {
cap.setup
.as_ref()
.and_then(|s| s.required_secrets.first())
.map(|s| s.name.clone())
.ok_or_else(|| {
ExtensionError::Other(format!(
"Tool '{}' has no auth or setup secrets",
name
))
})?
}
}
ExtensionKind::McpServer => {
let server = self
.get_mcp_server(name)
.await
.map_err(|e| ExtensionError::NotInstalled(e.to_string()))?;
server.token_secret_name()
}
ExtensionKind::ChannelRelay => format!("relay:{}:stream_token", name),
};
let mut secrets = std::collections::HashMap::new();
secrets.insert(secret_name, token.to_string());
self.configure(name, &secrets).await
}
/// Read a capabilities.json file and revoke its credential mappings from
/// the shared credential registry, so removed extensions lose injection
/// authority immediately.
@@ -4686,4 +4765,180 @@ mod tests {
assert_eq!(result, url);
assert!(result.contains("/v1/users/123/profile"));
}
// ── Regression tests for PR #677 (unify-extension-lifecycle) ─────────
#[tokio::test]
async fn test_configure_token_picks_first_missing_secret() {
// Regression: configure_token() must pick the first *missing* secret,
// not the first non-optional one. This allows multi-secret channels
// to be configured one secret at a time.
let dir = tempfile::tempdir().expect("temp dir");
let channels_dir = dir.path().join("channels");
std::fs::create_dir_all(&channels_dir).unwrap();
// Write a fake channel WASM + capabilities with two required secrets
std::fs::write(channels_dir.join("multi.wasm"), b"\0asm fake").unwrap();
let caps = serde_json::json!({
"type": "channel",
"name": "multi",
"setup": {
"required_secrets": [
{"name": "SECRET_A", "prompt": "Enter secret A (at least 30 chars for validation)"},
{"name": "SECRET_B", "prompt": "Enter secret B (at least 30 chars for validation)"}
]
}
});
std::fs::write(
channels_dir.join("multi.capabilities.json"),
serde_json::to_string(&caps).unwrap(),
)
.unwrap();
let mgr = make_manager_custom_dirs(dir.path().join("tools"), channels_dir);
// Pre-store SECRET_A so it's no longer missing
mgr.secrets
.create(
"test",
crate::secrets::CreateSecretParams::new("SECRET_A", "value-a"),
)
.await
.expect("store SECRET_A");
// configure_token should target SECRET_B (the first missing one)
let _result = mgr.configure_token("multi", "value-b").await;
// configure will fail at activation (no real WASM runtime), but the
// secret should still have been stored before activation was attempted.
// Check that SECRET_B was stored.
assert!(
mgr.secrets
.exists("test", "SECRET_B")
.await
.unwrap_or(false),
"configure_token should have stored SECRET_B (the first missing secret)"
);
}
#[tokio::test]
async fn test_auth_is_read_only_for_wasm_channel() {
// Regression: auth() must be a pure status check — it must not store
// any secrets or modify state. The old API accepted a token parameter.
let dir = tempfile::tempdir().expect("temp dir");
let channels_dir = dir.path().join("channels");
std::fs::create_dir_all(&channels_dir).unwrap();
std::fs::write(channels_dir.join("test-ch.wasm"), b"\0asm fake").unwrap();
let caps = serde_json::json!({
"type": "channel",
"name": "test-ch",
"setup": {
"required_secrets": [
{"name": "BOT_TOKEN", "prompt": "Enter bot token (at least 30 chars for prompt validation)"}
]
}
});
std::fs::write(
channels_dir.join("test-ch.capabilities.json"),
serde_json::to_string(&caps).unwrap(),
)
.unwrap();
let mgr = make_manager_custom_dirs(dir.path().join("tools"), channels_dir);
// auth() should return a result without storing anything
let result = mgr.auth("test-ch").await;
assert!(result.is_ok(), "auth should succeed: {:?}", result.err());
// No secrets should have been created
assert!(
!mgr.secrets
.exists("test", "BOT_TOKEN")
.await
.unwrap_or(true),
"auth() must not create any secrets — it should be read-only"
);
}
#[tokio::test]
async fn test_configure_dispatches_activation_by_kind() {
// Regression: configure() must dispatch to the correct activation method
// by kind. Previously it unconditionally called activate_wasm_channel()
// for all non-WasmTool types, which would fail with a channel-specific
// error for MCP servers and channel relays.
let dir = tempfile::tempdir().expect("temp dir");
let channels_dir = dir.path().join("channels");
std::fs::create_dir_all(&channels_dir).unwrap();
let mgr = make_manager_custom_dirs(dir.path().join("tools"), channels_dir);
// Register a channel relay extension (in-memory)
mgr.installed_relay_extensions
.write()
.await
.insert("test-relay".to_string());
// configure() should dispatch to activate_channel_relay(), not
// activate_wasm_channel(). Both will fail (no runtime configured),
// but the error should be about relay config, not WASM channels.
let mut secrets = std::collections::HashMap::new();
secrets.insert(
"relay:test-relay:stream_token".to_string(),
"tok".to_string(),
);
let result = mgr.configure("test-relay", &secrets).await;
assert!(
result.is_ok(),
"configure should return Ok: {:?}",
result.err()
);
let result = result.unwrap();
// Activation will fail (no relay config), but secrets should still be stored
assert!(
!result.activated,
"activation should fail without relay config"
);
assert!(
!result.message.contains("WASM"),
"error should not mention WASM — got: {}",
result.message
);
// Verify the secret was stored
assert!(
mgr.secrets
.exists("test", "relay:test-relay:stream_token")
.await
.unwrap_or(false),
"configure should have stored the relay stream token"
);
}
#[test]
fn test_validation_failed_is_distinct_error_variant() {
// Regression: ValidationFailed must be a distinct error variant so
// callers can match on it instead of parsing error message strings.
let err = ExtensionError::ValidationFailed("Invalid token".to_string());
assert!(
matches!(err, ExtensionError::ValidationFailed(_)),
"Should match ValidationFailed variant"
);
assert!(
!matches!(err, ExtensionError::Other(_)),
"Must NOT match Other variant"
);
assert!(
!matches!(err, ExtensionError::AuthFailed(_)),
"Must NOT match AuthFailed variant"
);
let msg = err.to_string();
assert!(
msg.contains("validation failed"),
"Display should contain 'validation failed', got: {msg}"
);
}
}
+17
View File
@@ -449,6 +449,20 @@ pub struct ActivateResult {
pub message: String,
}
/// Result of configuring secrets for an extension.
///
/// Returned by `ExtensionManager::configure()`, the single entrypoint
/// for providing secrets to any extension (chat auth, gateway setup, etc.).
#[derive(Debug, Clone)]
pub struct ConfigureResult {
/// Human-readable status message.
pub message: String,
/// Whether the extension was successfully activated after configuration.
pub activated: bool,
/// OAuth authorization URL (if OAuth flow was started).
pub auth_url: Option<String>,
}
fn default_true() -> bool {
true
}
@@ -530,6 +544,9 @@ pub enum ExtensionError {
fallback: Box<ExtensionError>,
},
#[error("Token validation failed: {0}")]
ValidationFailed(String),
#[error("{0}")]
Other(String),
}
+2 -2
View File
@@ -213,7 +213,7 @@ impl Tool for ToolAuthTool {
let result = self
.manager
.auth(name, None)
.auth(name)
.await
.map_err(|e| ToolError::ExecutionFailed(e.to_string()))?;
@@ -323,7 +323,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 {
match self.manager.auth(name).await {
Ok(auth_result) if auth_result.is_authenticated() => {
// Auth succeeded (e.g. env var was set); retry activation.
let result = self