Start extensions concurrently (#4234)

Co-authored-by: Michael Neale <michael.neale@gmail.com>

crates/goose-cli/src/scenario_tests/scenario_runner.rs

@@ -136,6 +136,9 @@ async fn run_provider_scenario_with_validation<F>(
 where
     F: Fn(&ScenarioResult) -> Result<()>,
 {
+    use goose::config::ExtensionConfig;
+    use tokio::sync::Mutex;
+
     if let Ok(path) = dotenv() {
         println!("Loaded environment from {:?}", path);
     }
@@ -193,10 +196,23 @@ where
     let mock_client = weather_client();
 
     let agent = Agent::new();
-    {
-        let mut extension_manager = agent.extension_manager.write().await;
-        extension_manager.add_client("weather_extension".to_string(), Box::new(mock_client));
-    }
+    agent
+        .extension_manager
+        .add_client(
+            "weather_extension".to_string(),
+            ExtensionConfig::Builtin {
+                name: "".to_string(),
+                display_name: None,
+                description: None,
+                timeout: None,
+                bundled: None,
+                available_tools: vec![],
+            },
+            Arc::new(Mutex::new(Box::new(mock_client))),
+            None,
+            None,
+        )
+        .await;
 
     agent
         .update_provider(provider_arc as Arc<dyn goose::providers::base::Provider>)

crates/goose-cli/src/session/builder.rs

@@ -7,8 +7,10 @@ use goose::recipe::{Response, SubRecipe};
 use goose::session;
 use goose::session::Identifier;
 use rustyline::EditMode;
+use std::collections::HashSet;
 use std::process;
 use std::sync::Arc;
+use tokio::task::JoinSet;
 
 use super::output;
 use super::Session;
@@ -355,38 +357,54 @@ pub async fn build_session(session_config: SessionBuilderConfig) -> Session {
             .collect()
     };
 
-    for extension in extensions_to_run {
-        if let Err(e) = agent.add_extension(extension.clone()).await {
-            let err = e.to_string();
-            eprintln!(
-                "{}",
-                style(format!(
-                    "Warning: Failed to start extension '{}': {}",
-                    extension.name(),
-                    err
-                ))
-                .yellow()
-            );
-            eprintln!(
-                "{}",
-                style(format!(
-                    "Continuing without extension '{}'",
-                    extension.name()
-                ))
-                .yellow()
-            );
+    let mut set = JoinSet::new();
+    let agent_ptr = Arc::new(agent);
 
-            // Offer debugging help
-            if let Err(debug_err) = offer_extension_debugging_help(
-                &extension.name(),
-                &err,
-                Arc::clone(&provider_for_display),
-                session_config.interactive,
+    let mut waiting_on = HashSet::new();
+    for extension in extensions_to_run {
+        waiting_on.insert(extension.name());
+        let agent_ptr = agent_ptr.clone();
+        set.spawn(async move {
+            (
+                extension.name(),
+                agent_ptr.add_extension(extension.clone()).await,
             )
-            .await
-            {
-                eprintln!("Note: Could not start debugging session: {}", debug_err);
+        });
+    }
+
+    let get_message = |waiting_on: &HashSet<String>| {
+        let mut names: Vec<_> = waiting_on.iter().cloned().collect();
+        names.sort();
+        format!("starting {} extensions: {}", names.len(), names.join(", "))
+    };
+
+    let spinner = cliclack::spinner();
+    spinner.start(get_message(&waiting_on));
+
+    let mut offer_debug = Vec::new();
+    while let Some(result) = set.join_next().await {
+        match result {
+            Ok((name, Ok(_))) => {
+                waiting_on.remove(&name);
+                spinner.set_message(get_message(&waiting_on));
             }
+            Ok((name, Err(e))) => offer_debug.push((name, e)),
+            Err(e) => tracing::error!("failed to add extension: {}", e),
+        }
+    }
+
+    spinner.clear();
+
+    for (name, err) in offer_debug {
+        if let Err(debug_err) = offer_extension_debugging_help(
+            &name,
+            &err.to_string(),
+            Arc::clone(&provider_for_display),
+            session_config.interactive,
+        )
+        .await
+        {
+            eprintln!("Note: Could not start debugging session: {}", debug_err);
         }
     }
 
@@ -405,7 +423,7 @@ pub async fn build_session(session_config: SessionBuilderConfig) -> Session {
 
     // Create new session
     let mut session = Session::new(
-        agent,
+        Arc::try_unwrap(agent_ptr).unwrap_or_else(|_| panic!("There should be no more references")),
         session_file.clone(),
         session_config.debug,
         session_config.scheduled_job_id.clone(),