gitops/optimclaw
1
0
Fork 0
mirror of https://github.com/outbackdingo/optimclaw.git synced 2026-08-26 23:50:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(extensions): fix lifecycle bugs + comprehensive E2E tests (#1070)

Browse Source 
* 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]>

* fix(extensions): fix 5 extension lifecycle bugs found during E2E testing

Bug fixes in src/extensions/manager.rs:
- Add auth guard to activate_wasm_tool() blocking activation when secrets
  are missing (NeedsSetup), matching activate_wasm_channel() behavior
- Evict WasmToolRuntime module cache on remove() so reinstall uses fresh binary
- Clear activation_errors on remove() for both WasmTool and WasmChannel
- Clean up in-progress OAuth flows on remove() (abort TCP listener, purge
  pending flow entries)

Bug fix in src/channels/web/server.rs:
- Broadcast AuthCompleted SSE event on expired OAuth callback so web UI
  doesn't stay stuck showing "auth required"

E2E test coverage:
- test_wasm_lifecycle.py: 35 tests covering install/configure/activate/
  remove/reinstall lifecycle with regression tests for bugs 1 and 3
- test_extension_oauth.py: 9 tests covering OAuth round-trip flow
- test_tool_execution.py: 5 tests for tool invocation via chat
- test_pairing.py: 4 tests for pairing request lifecycle
- Enhanced conftest.py, helpers.py, mock_llm.py for OAuth mock support

[skip-regression-check]

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

* fix(web): unify extension auth UX and add lifecycle regressions

* test: fix pending oauth flow fixtures after rebase

* test(e2e): fix playwright route ordering for extensions reloads

* test: address e2e review follow-ups

* test: address remaining PR review comments

---------

Co-authored-by: Claude Opus 4.6 (1M context) <[email protected]>
This commit is contained in:
Henry Park
2026-03-12 16:36:08 -07:00
committed by GitHub
co-authored by Claude Opus 4.6
parent 8a60fa2d37
commit 9fbdd42988
13 changed files with 1848 additions and 190 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/channels/web/server.rs
+120 -72
View File
@@ -573,6 +573,14 @@ async fn oauth_callback_handler(
extension = %flow.extension_name,
"OAuth flow expired"
);
// Notify UI so auth card can show error instead of staying stuck
if let Some(ref sender) = flow.sse_sender {
let _ = sender.send(SseEvent::AuthCompleted {
extension_name: flow.extension_name.clone(),
success: false,
message: "OAuth flow expired. Please try again.".to_string(),
});
}
return oauth_error_page(&flow.display_name);
}
@@ -2706,6 +2714,7 @@ struct GatewayStatusResponse {
#[cfg(test)]
mod tests {
use super::*;
use crate::cli::oauth_defaults;
use crate::testing::credentials::TEST_GATEWAY_CRYPTO_KEY;
#[test]
@@ -2823,6 +2832,11 @@ mod tests {
.with_state(state)
}
fn expired_flow_created_at() -> Option<std::time::Instant> {
std::time::Instant::now()
.checked_sub(oauth_defaults::OAUTH_FLOW_EXPIRY + std::time::Duration::from_secs(1))
}
#[tokio::test]
async fn test_csp_header_present_on_responses() {
use std::net::SocketAddr;
@@ -2929,29 +2943,14 @@ mod tests {
use tower::ServiceExt;
// Build an ExtensionManager so the handler can look up flows
let secrets = Arc::new(crate::secrets::InMemorySecretsStore::new(Arc::new(
crate::secrets::SecretsCrypto::new(secrecy::SecretString::from(
TEST_GATEWAY_CRYPTO_KEY.to_string(),
))
.expect("crypto"),
)));
let tool_registry = Arc::new(ToolRegistry::new());
let mcp_sm = Arc::new(crate::tools::mcp::session::McpSessionManager::new());
let ext_mgr = Arc::new(ExtensionManager::new(
mcp_sm,
Arc::new(crate::tools::mcp::process::McpProcessManager::new()),
secrets,
tool_registry,
None,
None,
std::path::PathBuf::from("/tmp/wasm_tools"),
std::path::PathBuf::from("/tmp/wasm_channels"),
None,
"test".to_string(),
None,
vec![],
));
let secrets: Arc<dyn crate::secrets::SecretsStore + Send + Sync> =
Arc::new(crate::secrets::InMemorySecretsStore::new(Arc::new(
crate::secrets::SecretsCrypto::new(secrecy::SecretString::from(
TEST_GATEWAY_CRYPTO_KEY.to_string(),
))
.expect("crypto"),
)));
let (ext_mgr, _wasm_tools_dir, _wasm_channels_dir) = test_ext_mgr(secrets);
let state = test_gateway_state(Some(ext_mgr));
let app = test_oauth_router(state);
@@ -2985,25 +2984,13 @@ mod tests {
))
.expect("crypto"),
)));
let tool_registry = Arc::new(ToolRegistry::new());
let mcp_sm = Arc::new(crate::tools::mcp::session::McpSessionManager::new());
let (ext_mgr, _wasm_tools_dir, _wasm_channels_dir) = test_ext_mgr(secrets.clone());
let Some(created_at) = expired_flow_created_at() else {
eprintln!("Skipping expired OAuth flow test: monotonic uptime below expiry window");
return;
};
let ext_mgr = Arc::new(ExtensionManager::new(
mcp_sm,
Arc::new(crate::tools::mcp::process::McpProcessManager::new()),
secrets.clone(),
tool_registry,
None,
None,
std::path::PathBuf::from("/tmp/wasm_tools"),
std::path::PathBuf::from("/tmp/wasm_channels"),
None,
"test".to_string(),
None,
vec![],
));
// Insert an expired flow (created 10 minutes ago)
// Insert an expired flow.
let flow = crate::cli::oauth_defaults::PendingOAuthFlow {
extension_name: "test_tool".to_string(),
display_name: "Test Tool".to_string(),
@@ -3023,9 +3010,7 @@ mod tests {
gateway_token: None,
resource: None,
client_id_secret_name: None,
created_at: std::time::Instant::now()
.checked_sub(std::time::Duration::from_secs(600))
.expect("System uptime is too low to run expired flow test"),
created_at,
};
ext_mgr
@@ -3055,6 +3040,80 @@ mod tests {
assert!(html.contains("Authorization Failed"));
}
#[tokio::test]
async fn test_oauth_callback_expired_flow_broadcasts_auth_completed_failure() {
use axum::body::Body;
use tower::ServiceExt;
let secrets: Arc<dyn crate::secrets::SecretsStore + Send + Sync> =
Arc::new(crate::secrets::InMemorySecretsStore::new(Arc::new(
crate::secrets::SecretsCrypto::new(secrecy::SecretString::from(
TEST_GATEWAY_CRYPTO_KEY.to_string(),
))
.expect("crypto"),
)));
let (ext_mgr, _wasm_tools_dir, _wasm_channels_dir) = test_ext_mgr(secrets.clone());
let (sender, mut receiver) = tokio::sync::broadcast::channel(4);
let Some(created_at) = expired_flow_created_at() else {
eprintln!("Skipping expired OAuth flow SSE test: monotonic uptime below expiry window");
return;
};
let flow = crate::cli::oauth_defaults::PendingOAuthFlow {
extension_name: "test_tool".to_string(),
display_name: "Test Tool".to_string(),
token_url: "https://example.com/token".to_string(),
client_id: "client123".to_string(),
client_secret: None,
redirect_uri: "https://example.com/oauth/callback".to_string(),
code_verifier: None,
access_token_field: "access_token".to_string(),
secret_name: "test_token".to_string(),
provider: None,
validation_endpoint: None,
scopes: vec![],
user_id: "test".to_string(),
secrets,
sse_sender: Some(sender),
gateway_token: None,
resource: None,
client_id_secret_name: None,
created_at,
};
ext_mgr
.pending_oauth_flows()
.write()
.await
.insert("expired_state".to_string(), flow);
let state = test_gateway_state(Some(ext_mgr));
let app = test_oauth_router(state);
let req = axum::http::Request::builder()
.uri("/oauth/callback?code=test_code&state=expired_state")
.body(Body::empty())
.expect("request");
let resp = ServiceExt::<axum::http::Request<Body>>::oneshot(app, req)
.await
.expect("response");
assert_eq!(resp.status(), StatusCode::OK);
match receiver.recv().await.expect("auth_completed event") {
crate::channels::web::types::SseEvent::AuthCompleted {
extension_name,
success,
message,
} => {
assert_eq!(extension_name, "test_tool");
assert!(!success, "expired OAuth flow should broadcast failure");
assert_eq!(message, "OAuth flow expired. Please try again.");
}
event => panic!("expected AuthCompleted event, got {event:?}"),
}
}
#[tokio::test]
async fn test_oauth_callback_no_extension_manager() {
use axum::body::Body;
@@ -3093,28 +3152,16 @@ mod tests {
))
.expect("crypto"),
)));
let tool_registry = Arc::new(ToolRegistry::new());
let mcp_sm = Arc::new(crate::tools::mcp::session::McpSessionManager::new());
let ext_mgr = Arc::new(ExtensionManager::new(
mcp_sm,
Arc::new(crate::tools::mcp::process::McpProcessManager::new()),
secrets.clone(),
tool_registry,
None,
None,
std::path::PathBuf::from("/tmp/wasm_tools"),
std::path::PathBuf::from("/tmp/wasm_channels"),
None,
"test".to_string(),
None,
vec![],
));
let (ext_mgr, _wasm_tools_dir, _wasm_channels_dir) = test_ext_mgr(secrets.clone());
// Insert a flow keyed by raw nonce "test_nonce" (without instance prefix).
// Use an expired flow so the handler exits before attempting a real HTTP
// token exchange — we only need to verify that the instance prefix was
// stripped and the flow was found by the raw nonce.
let Some(created_at) = expired_flow_created_at() else {
eprintln!("Skipping OAuth state-prefix test: monotonic uptime below expiry window");
return;
};
let flow = crate::cli::oauth_defaults::PendingOAuthFlow {
extension_name: "test_tool".to_string(),
display_name: "Test Tool".to_string(),
@@ -3135,9 +3182,7 @@ mod tests {
resource: None,
client_id_secret_name: None,
// Expired — handler will reject after lookup (no network I/O)
created_at: std::time::Instant::now()
.checked_sub(std::time::Duration::from_secs(600))
.expect("System uptime is too low to run expired flow test"),
created_at,
};
ext_mgr
@@ -3208,24 +3253,27 @@ mod tests {
fn test_ext_mgr(
secrets: Arc<dyn crate::secrets::SecretsStore + Send + Sync>,
) -> Arc<ExtensionManager> {
) -> (Arc<ExtensionManager>, tempfile::TempDir, tempfile::TempDir) {
let tool_registry = Arc::new(ToolRegistry::new());
let mcp_sm = Arc::new(crate::tools::mcp::session::McpSessionManager::new());
let mcp_pm = Arc::new(crate::tools::mcp::process::McpProcessManager::new());
Arc::new(ExtensionManager::new(
let wasm_tools_dir = tempfile::tempdir().expect("temp wasm tools dir");
let wasm_channels_dir = tempfile::tempdir().expect("temp wasm channels dir");
let ext_mgr = Arc::new(ExtensionManager::new(
mcp_sm,
mcp_pm,
secrets,
tool_registry,
None,
None,
std::path::PathBuf::from("/tmp/wasm_tools"),
std::path::PathBuf::from("/tmp/wasm_channels"),
wasm_tools_dir.path().to_path_buf(),
wasm_channels_dir.path().to_path_buf(),
None,
"test".to_string(),
None,
vec![],
))
));
(ext_mgr, wasm_tools_dir, wasm_channels_dir)
}
#[tokio::test]
@@ -3234,7 +3282,7 @@ mod tests {
use tower::ServiceExt;
let secrets = test_secrets_store();
let ext_mgr = test_ext_mgr(secrets);
let (ext_mgr, _wasm_tools_dir, _wasm_channels_dir) = test_ext_mgr(secrets);
let state = test_gateway_state(Some(ext_mgr));
let app = test_relay_oauth_router(state);
@@ -3278,7 +3326,7 @@ mod tests {
.await
.expect("store nonce");
let ext_mgr = test_ext_mgr(secrets);
let (ext_mgr, _wasm_tools_dir, _wasm_channels_dir) = test_ext_mgr(secrets);
let state = test_gateway_state(Some(ext_mgr));
let app = test_relay_oauth_router(state);
@@ -3323,7 +3371,7 @@ mod tests {
.await
.expect("store nonce");
let ext_mgr = test_ext_mgr(secrets.clone());
let (ext_mgr, _wasm_tools_dir, _wasm_channels_dir) = test_ext_mgr(secrets.clone());
let state = test_gateway_state(Some(ext_mgr));
let app = test_relay_oauth_router(state);
src/channels/web/static/app.js
+104 -36
View File
@@ -358,29 +358,11 @@ function connectSSE() {
});
eventSource.addEventListener('auth_required', (e) => {
const data = JSON.parse(e.data);
if (data.auth_url) {
// OAuth flow: show the auth card with an OAuth button + optional token paste field.
showAuthCard(data);
} else {
// Setup flow: fetch the extension's credential schema and show the multi-field
// configure modal (the same UI used by the Extensions tab "Setup" button).
showConfigureModal(data.extension_name);
}
handleAuthRequired(JSON.parse(e.data));
});
eventSource.addEventListener('auth_completed', (e) => {
const data = JSON.parse(e.data);
// Dismiss whichever UI path was active: auth card (OAuth) or configure modal (setup).
removeAuthCard(data.extension_name);
closeConfigureModal();
showToast(data.message, data.success ? 'success' : 'error');
if (shouldShowChannelConnectedMessage(data.extension_name, data.success)) {
addMessage('system', 'Telegram is now connected. You can message me there and I can send you notifications.');
}
// Refresh extensions list so status indicators update
if (currentTab === 'extensions') loadExtensions();
enableChatInput();
handleAuthCompleted(JSON.parse(e.data));
});
eventSource.addEventListener('extension_status', (e) => {
@@ -1139,13 +1121,71 @@ function showJobCard(data) {
// --- Auth card ---
function handleAuthRequired(data) {
if (data.auth_url) {
// OAuth flow: show the global auth prompt with an OAuth button + optional token paste field.
showAuthCard(data);
} else {
// Setup flow: fetch the extension's credential schema and show the multi-field
// configure modal (the same UI used by the Extensions tab "Setup" button).
showConfigureModal(data.extension_name);
}
}
function handleAuthCompleted(data) {
// Dismiss only the matching extension's UI so unrelated setup work is not interrupted.
removeAuthCard(data.extension_name);
closeConfigureModal(data.extension_name);
showToast(data.message, data.success ? 'success' : 'error');
if (shouldShowChannelConnectedMessage(data.extension_name, data.success)) {
addMessage('system', 'Telegram is now connected. You can message me there and I can send you notifications.');
}
if (currentTab === 'extensions') loadExtensions();
enableChatInput();
}
function queryByDataAttribute(selector, attributeName, attributeValue) {
if (typeof attributeValue !== 'string') return document.querySelector(selector);
if (window.CSS && typeof window.CSS.escape === 'function') {
return document.querySelector(
selector + '[' + attributeName + '="' + window.CSS.escape(attributeValue) + '"]'
);
}
const candidates = document.querySelectorAll(selector);
for (const candidate of candidates) {
if (candidate.getAttribute(attributeName) === attributeValue) return candidate;
}
return null;
}
function getAuthOverlay(extensionName) {
return queryByDataAttribute('.auth-overlay', 'data-extension-name', extensionName);
}
function getAuthCard(extensionName) {
return queryByDataAttribute('.auth-card', 'data-extension-name', extensionName);
}
function getConfigureOverlay(extensionName) {
return queryByDataAttribute('.configure-overlay', 'data-extension-name', extensionName);
}
function showAuthCard(data) {
// Remove any existing card for this extension first
removeAuthCard(data.extension_name);
// Keep a single global auth prompt so the experience is consistent across tabs.
const existing = getAuthOverlay();
if (existing) existing.remove();
const overlay = document.createElement('div');
overlay.className = 'auth-overlay';
overlay.setAttribute('data-extension-name', data.extension_name);
overlay.addEventListener('click', (e) => {
if (e.target === overlay) cancelAuth(data.extension_name);
});
const container = document.getElementById('chat-messages');
const card = document.createElement('div');
card.className = 'auth-card';
card.className = 'auth-card auth-modal';
card.setAttribute('data-extension-name', data.extension_name);
const header = document.createElement('div');
@@ -1224,21 +1264,30 @@ function showAuthCard(data) {
actions.appendChild(cancelBtn);
card.appendChild(actions);
container.appendChild(card);
container.scrollTop = container.scrollHeight;
overlay.appendChild(card);
document.body.appendChild(overlay);
tokenInput.focus();
}
function removeAuthCard(extensionName) {
const card = document.querySelector('.auth-card[data-extension-name="' + extensionName + '"]');
if (card) card.remove();
const overlay = getAuthOverlay(extensionName);
if (overlay) {
overlay.remove();
return;
}
const card = getAuthCard(extensionName);
if (card) {
const parentOverlay = card.closest('.auth-overlay');
if (parentOverlay) parentOverlay.remove();
else card.remove();
}
}
function submitAuthToken(extensionName, tokenValue) {
if (!tokenValue || !tokenValue.trim()) return;
// Disable submit button while in flight
const card = document.querySelector('.auth-card[data-extension-name="' + extensionName + '"]');
const card = getAuthCard(extensionName);
if (card) {
const btns = card.querySelectorAll('button');
btns.forEach((b) => { b.disabled = true; });
@@ -1249,8 +1298,10 @@ function submitAuthToken(extensionName, tokenValue) {
body: { extension_name: extensionName, token: tokenValue.trim() },
}).then((result) => {
if (result.success) {
// Close immediately for responsiveness; the authoritative success UX
// (toast + extensions refresh) still comes from auth_completed SSE.
removeAuthCard(extensionName);
addMessage('system', result.message);
enableChatInput();
} else {
showAuthCardError(extensionName, result.message);
}
@@ -1269,7 +1320,7 @@ function cancelAuth(extensionName) {
}
function showAuthCardError(extensionName, message) {
const card = document.querySelector('.auth-card[data-extension-name="' + extensionName + '"]');
const card = getAuthCard(extensionName);
if (!card) return;
// Re-enable buttons
const btns = card.querySelectorAll('button');
@@ -2199,6 +2250,10 @@ function renderAvailableExtensionCard(entry) {
showToast(I18n.t('extensions.installedSuccess', {name: entry.display_name}), 'success');
// OAuth popup if auth started during install (builtin creds)
if (res.auth_url) {
showAuthCard({
extension_name: entry.name,
auth_url: res.auth_url,
});
showToast('Opening authentication for ' + entry.display_name, 'info');
openOAuthUrl(res.auth_url);
}
@@ -2464,6 +2519,10 @@ function activateExtension(name) {
if (res.success) {
// Even on success, the tool may need OAuth (e.g., WASM loaded but no token yet)
if (res.auth_url) {
showAuthCard({
extension_name: name,
auth_url: res.auth_url,
});
showToast('Opening authentication for ' + name, 'info');
openOAuthUrl(res.auth_url);
}
@@ -2472,6 +2531,10 @@ function activateExtension(name) {
}
if (res.auth_url) {
showAuthCard({
extension_name: name,
auth_url: res.auth_url,
});
showToast('Opening authentication for ' + name, 'info');
openOAuthUrl(res.auth_url);
} else if (res.awaiting_token) {
@@ -2514,6 +2577,7 @@ function renderConfigureModal(name, secrets) {
closeConfigureModal();
const overlay = document.createElement('div');
overlay.className = 'configure-overlay';
overlay.setAttribute('data-extension-name', name);
overlay.addEventListener('click', (e) => {
if (e.target === overlay) closeConfigureModal();
});
@@ -2607,7 +2671,8 @@ function submitConfigureModal(name, fields) {
}
// Disable buttons to prevent double-submit
var btns = document.querySelectorAll('.configure-actions button');
const overlay = getConfigureOverlay(name) || document.querySelector('.configure-overlay');
var btns = overlay ? overlay.querySelectorAll('.configure-actions button') : [];
btns.forEach(function(b) { b.disabled = true; });
apiFetch('/api/extensions/' + encodeURIComponent(name) + '/setup', {
@@ -2618,8 +2683,10 @@ function submitConfigureModal(name, fields) {
if (res.success) {
closeConfigureModal();
if (res.auth_url) {
// OAuth flow started — open consent popup. The auth_completed SSE will
// not arrive immediately (it fires after OAuth callback), so show a toast now.
showAuthCard({
extension_name: name,
auth_url: res.auth_url,
});
showToast('Opening OAuth authorization for ' + name, 'info');
openOAuthUrl(res.auth_url);
loadExtensions();
@@ -2638,8 +2705,9 @@ function submitConfigureModal(name, fields) {
});
}
function closeConfigureModal() {
const existing = document.querySelector('.configure-overlay');
function closeConfigureModal(extensionName) {
if (typeof extensionName !== 'string') extensionName = null;
const existing = getConfigureOverlay(extensionName);
if (existing) existing.remove();
}
src/channels/web/static/style.css
+25 -1
View File
@@ -1219,7 +1219,21 @@ body {
color: var(--danger);
}
/* Auth card (inline in chat) */
/* Auth prompt */
.auth-overlay {
position: fixed;
top: 0;
left: 0;
width: 100%;
height: 100%;
background: rgba(0, 0, 0, 0.6);
z-index: 1001;
display: flex;
align-items: center;
justify-content: center;
padding: 16px;
}
.auth-card {
align-self: flex-start;
max-width: 80%;
@@ -1234,6 +1248,16 @@ body {
transition: border-color 0.2s;
}
.auth-overlay .auth-card {
width: 460px;
max-width: min(460px, 90vw);
margin: 0;
align-self: auto;
background: var(--bg);
border-color: rgba(52, 211, 153, 0.35);
box-shadow: 0 24px 48px rgba(0, 0, 0, 0.35);
}
.auth-card .auth-header {
font-weight: 600;
color: var(--accent);
src/extensions/manager.rs
+189 -6
View File
@@ -786,6 +786,19 @@ impl ExtensionManager {
Self::validate_extension_name(name)?;
let kind = self.determine_installed_kind(name).await?;
// Clean up any in-progress OAuth flows for this extension.
// TCP mode: abort the listener task so port 9876 is freed immediately.
// Gateway mode: remove stale pending flow entries.
if let Some(pending) = self.pending_auth.write().await.remove(name)
&& let Some(handle) = pending.task_handle
{
handle.abort();
}
self.pending_oauth_flows
.write()
.await
.retain(|_, flow| flow.extension_name != name);
match kind {
ExtensionKind::McpServer => {
// Unregister tools with this server's prefix
@@ -819,6 +832,14 @@ impl ExtensionManager {
// Unregister from tool registry
self.tool_registry.unregister(name).await;
// Evict compiled module from runtime cache so reinstall uses fresh binary
if let Some(ref rt) = self.wasm_tool_runtime {
rt.remove(name).await;
}
// Clear stale activation errors so reinstall starts clean
self.activation_errors.write().await.remove(name);
// Revoke credential mappings from the shared registry
let cap_path = self
.wasm_tools_dir
@@ -859,6 +880,9 @@ impl ExtensionManager {
self.active_channel_names.write().await.remove(name);
self.persist_active_channels().await;
// Clear stale activation errors so reinstall starts clean
self.activation_errors.write().await.remove(name);
// Delete channel files
let wasm_path = self.wasm_channels_dir.join(format!("{}.wasm", name));
let cap_path = self
@@ -2860,6 +2884,17 @@ impl ExtensionManager {
});
}
// Check auth status — block activation if required secrets are missing.
// NeedsAuth (OAuth not yet completed) is allowed because configure() loads
// the tool first, then starts the OAuth flow to obtain the token.
let auth_state = self.check_tool_auth_status(name).await;
if auth_state == ToolAuthState::NeedsSetup {
return Err(ExtensionError::ActivationFailed(format!(
"Tool '{}' requires configuration. Use the setup form to provide credentials.",
name
)));
}
let runtime = self.wasm_tool_runtime.as_ref().ok_or_else(|| {
ExtensionError::ActivationFailed("WASM runtime not available".to_string())
})?;
@@ -4495,14 +4530,18 @@ mod tests {
// available" because the ExtensionManager had `wasm_tool_runtime: None`.
/// Build a minimal ExtensionManager suitable for unit tests.
fn make_test_manager(
fn make_test_manager_with_dirs(
wasm_runtime: Option<Arc<crate::tools::wasm::WasmToolRuntime>>,
tools_dir: std::path::PathBuf,
channels_dir: std::path::PathBuf,
) -> crate::extensions::manager::ExtensionManager {
use crate::secrets::{InMemorySecretsStore, SecretsCrypto};
use crate::tools::mcp::process::McpProcessManager;
use crate::tools::mcp::session::McpSessionManager;
std::fs::create_dir_all(&tools_dir).ok();
std::fs::create_dir_all(&channels_dir).ok();
let key = secrecy::SecretString::from(crate::secrets::keychain::generate_master_key_hex());
let crypto = Arc::new(SecretsCrypto::new(key).expect("crypto"));
let secrets: Arc<dyn crate::secrets::SecretsStore + Send + Sync> =
@@ -4517,15 +4556,22 @@ mod tests {
tools,
None, // hooks
wasm_runtime,
tools_dir.clone(),
tools_dir, // channels dir (unused here)
None, // tunnel_url
tools_dir,
channels_dir,
None, // tunnel_url
"test".to_string(),
None, // db
vec![],
)
}
fn make_test_manager(
wasm_runtime: Option<Arc<crate::tools::wasm::WasmToolRuntime>>,
tools_dir: std::path::PathBuf,
) -> crate::extensions::manager::ExtensionManager {
make_test_manager_with_dirs(wasm_runtime, tools_dir.clone(), tools_dir)
}
#[tokio::test]
async fn test_activate_wasm_tool_with_runtime_passes_runtime_check() {
// When the ExtensionManager has a WASM runtime, activation should get
@@ -4878,6 +4924,145 @@ mod tests {
);
}
#[tokio::test]
async fn test_remove_wasm_tool_clears_pending_oauth_state_and_activation_error() {
let dir = tempfile::tempdir().expect("temp dir");
let mgr = make_test_manager(None, dir.path().to_path_buf());
std::fs::write(dir.path().join("gmail.wasm"), b"fake-tool").expect("write tool");
let listener = tokio::spawn(async {
std::future::pending::<()>().await;
});
let abort_handle = listener.abort_handle();
mgr.pending_auth.write().await.insert(
"gmail".to_string(),
super::PendingAuth {
_name: "gmail".to_string(),
_kind: ExtensionKind::WasmTool,
created_at: std::time::Instant::now(),
task_handle: Some(listener),
},
);
mgr.activation_errors
.write()
.await
.insert("gmail".to_string(), "cached failure".to_string());
let secrets = Arc::clone(&mgr.secrets);
mgr.pending_oauth_flows().write().await.insert(
"gmail-state".to_string(),
crate::cli::oauth_defaults::PendingOAuthFlow {
extension_name: "gmail".to_string(),
display_name: "Gmail".to_string(),
token_url: "https://example.com/token".to_string(),
client_id: "client123".to_string(),
client_secret: None,
redirect_uri: "https://example.com/oauth/callback".to_string(),
code_verifier: None,
access_token_field: "access_token".to_string(),
secret_name: "google_oauth_token".to_string(),
provider: None,
validation_endpoint: None,
scopes: vec![],
user_id: "test".to_string(),
secrets: Arc::clone(&secrets),
sse_sender: None,
gateway_token: None,
resource: None,
client_id_secret_name: None,
created_at: std::time::Instant::now(),
},
);
mgr.pending_oauth_flows().write().await.insert(
"other-state".to_string(),
crate::cli::oauth_defaults::PendingOAuthFlow {
extension_name: "web-search".to_string(),
display_name: "Web Search".to_string(),
token_url: "https://example.com/token".to_string(),
client_id: "client456".to_string(),
client_secret: None,
redirect_uri: "https://example.com/oauth/callback".to_string(),
code_verifier: None,
access_token_field: "access_token".to_string(),
secret_name: "other_token".to_string(),
provider: None,
validation_endpoint: None,
scopes: vec![],
user_id: "test".to_string(),
secrets,
sse_sender: None,
gateway_token: None,
resource: None,
client_id_secret_name: None,
created_at: std::time::Instant::now(),
},
);
let result = mgr.remove("gmail").await;
assert!(result.is_ok(), "remove should succeed: {:?}", result.err());
tokio::task::yield_now().await;
assert!(
mgr.pending_auth.read().await.get("gmail").is_none(),
"pending auth entry should be removed"
);
assert!(
abort_handle.is_finished(),
"pending auth listener should be aborted"
);
assert!(
!mgr.activation_errors.read().await.contains_key("gmail"),
"stale activation error should be cleared"
);
let flows = mgr.pending_oauth_flows().read().await;
assert!(
!flows.contains_key("gmail-state"),
"gateway OAuth flow for removed extension should be cleared"
);
assert!(
flows.contains_key("other-state"),
"unrelated pending OAuth flows should be retained"
);
}
#[tokio::test]
async fn test_remove_wasm_channel_clears_activation_error_and_deletes_files() {
let dir = tempfile::tempdir().expect("temp dir");
let tools_dir = dir.path().join("tools");
let channels_dir = dir.path().join("channels");
let mgr = make_test_manager_with_dirs(None, tools_dir, channels_dir.clone());
let wasm_path = channels_dir.join("telegram.wasm");
let cap_path = channels_dir.join("telegram.capabilities.json");
std::fs::write(&wasm_path, b"fake-channel").expect("write channel");
std::fs::write(&cap_path, b"{}").expect("write capabilities");
mgr.activation_errors
.write()
.await
.insert("telegram".to_string(), "channel failed".to_string());
let result = mgr.remove("telegram").await;
assert!(result.is_ok(), "remove should succeed: {:?}", result.err());
assert!(
!mgr.activation_errors.read().await.contains_key("telegram"),
"channel activation error should be cleared on remove"
);
assert!(
!wasm_path.exists(),
"channel wasm file should be deleted on remove"
);
assert!(
!cap_path.exists(),
"channel capabilities file should be deleted on remove"
);
}
#[test]
fn test_sanitize_url_with_query_params() {
let url = "https://api.example.com/path?api_key=secret123&token=abc";
@@ -5153,7 +5338,6 @@ mod tests {
Some("https://my-gateway.example.com/oauth/callback".to_string()),
);
}
// ── Regression tests for PR #677 (unify-extension-lifecycle) ─────────
#[tokio::test]
@@ -5303,7 +5487,6 @@ mod tests {
"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