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

.github/workflows/e2e.yml

@@ -52,7 +52,7 @@ jobs:
           - group: features
             files: "tests/e2e/scenarios/test_skills.py tests/e2e/scenarios/test_tool_approval.py"
           - group: extensions
-            files: "tests/e2e/scenarios/test_extensions.py"
+            files: "tests/e2e/scenarios/test_extensions.py tests/e2e/scenarios/test_extension_oauth.py tests/e2e/scenarios/test_wasm_lifecycle.py tests/e2e/scenarios/test_tool_execution.py tests/e2e/scenarios/test_pairing.py"
     steps:
       - uses: actions/checkout@v6
 

src/channels/web/server.rs

@@ -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 = 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());
+        let Some(created_at) = expired_flow_created_at() else {
+            eprintln!("Skipping expired OAuth flow test: monotonic uptime below expiry window");
+            return;
+        };
 
-        // 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);