feat: goose bench framework for functional and regression testing

Co-authored-by: Zaki Ali <zaki@squareup.com>

crates/goose-cli/Cargo.toml

@@ -13,6 +13,7 @@ path = "src/main.rs"
 
 [dependencies]
 goose = { path = "../goose" }
+goose-bench = { path = "../goose-bench" }
 goose-mcp = { path = "../goose-mcp" }
 mcp-client = { path = "../mcp-client" }
 mcp-server = { path = "../mcp-server" }
@@ -48,6 +49,7 @@ tracing-subscriber = { version = "0.3", features = ["env-filter", "fmt", "json",
 tracing-appender = "0.2"
 once_cell = "1.20.2"
 shlex = "1.3.0"
+async-trait = "0.1.86"
 
 [target.'cfg(target_os = "windows")'.dependencies]
 winapi = { version = "0.3", features = ["wincred"] }

crates/goose-cli/src/commands/bench.rs

@@ -0,0 +1,171 @@
+use crate::session::build_session;
+use crate::Session;
+use async_trait::async_trait;
+use chrono::Local;
+use goose::config::Config;
+use goose::message::Message;
+use goose_bench::error_capture::ErrorCaptureLayer;
+use goose_bench::eval_suites::{BenchAgent, BenchAgentError, Evaluation, EvaluationSuiteFactory};
+use goose_bench::reporting::{BenchmarkResults, EvaluationResult, SuiteResult};
+use goose_bench::work_dir::WorkDir;
+use std::collections::HashMap;
+use std::path::PathBuf;
+use std::sync::Arc;
+use std::sync::Once;
+use tokio::sync::Mutex;
+use tracing_subscriber::layer::SubscriberExt;
+
+// Used to ensure we only set up tracing once
+static INIT: Once = Once::new();
+
+pub struct BenchSession {
+    session: Session,
+    errors: Arc<Mutex<Vec<BenchAgentError>>>,
+}
+
+impl BenchSession {
+    pub fn new(session: Session) -> Self {
+        let errors = Arc::new(Mutex::new(Vec::new()));
+
+        // Create and register the error capture layer only once
+        INIT.call_once(|| {
+            let error_layer = ErrorCaptureLayer::new(errors.clone());
+            let subscriber = tracing_subscriber::Registry::default().with(error_layer);
+
+            tracing::subscriber::set_global_default(subscriber)
+                .expect("Failed to set tracing subscriber");
+        });
+
+        Self { session, errors }
+    }
+}
+
+#[async_trait]
+impl BenchAgent for BenchSession {
+    async fn prompt(&mut self, p: String) -> anyhow::Result<Vec<Message>> {
+        // Clear previous errors
+        {
+            let mut errors = self.errors.lock().await;
+            errors.clear();
+        }
+
+        self.session.headless(p).await?;
+        Ok(self.session.message_history())
+    }
+
+    async fn get_errors(&self) -> Vec<BenchAgentError> {
+        let errors = self.errors.lock().await;
+        errors.clone()
+    }
+}
+
+// Wrapper struct to implement BenchAgent for Arc<Mutex<BenchSession>>
+struct BenchAgentWrapper(Arc<Mutex<BenchSession>>);
+
+#[async_trait]
+impl BenchAgent for BenchAgentWrapper {
+    async fn prompt(&mut self, p: String) -> anyhow::Result<Vec<Message>> {
+        let mut session = self.0.lock().await;
+        session.prompt(p).await
+    }
+
+    async fn get_errors(&self) -> Vec<BenchAgentError> {
+        let session = self.0.lock().await;
+        session.get_errors().await
+    }
+}
+
+async fn run_eval(
+    evaluation: Box<dyn Evaluation>,
+    work_dir: &mut WorkDir,
+) -> anyhow::Result<EvaluationResult> {
+    let mut result = EvaluationResult::new(evaluation.name().to_string());
+
+    if let Ok(work_dir) = work_dir.move_to(format!("./{}", &evaluation.name())) {
+        let required_extensions = evaluation.required_extensions();
+
+        // Create session with error capture
+        let base_session = build_session(None, false, Vec::new(), required_extensions).await;
+
+        let bench_session = Arc::new(Mutex::new(BenchSession::new(base_session)));
+        let bench_session_clone = bench_session.clone();
+
+        if let Ok(metrics) = evaluation
+            .run(Box::new(BenchAgentWrapper(bench_session)), work_dir)
+            .await
+        {
+            for (name, metric) in metrics {
+                result.add_metric(name, metric);
+            }
+
+            // Add any errors that occurred
+            let agent = BenchAgentWrapper(bench_session_clone);
+            for error in agent.get_errors().await {
+                result.add_error(error);
+            }
+        }
+    }
+
+    Ok(result)
+}
+
+async fn run_suite(suite: &str, work_dir: &mut WorkDir) -> anyhow::Result<SuiteResult> {
+    let mut suite_result = SuiteResult::new(suite.to_string());
+
+    if let Ok(work_dir) = work_dir.move_to(format!("./{}", &suite)) {
+        if let Some(evals) = EvaluationSuiteFactory::create(suite) {
+            for eval in evals {
+                let eval_result = run_eval(eval, work_dir).await?;
+                suite_result.add_evaluation(eval_result);
+            }
+        }
+    }
+
+    Ok(suite_result)
+}
+
+pub async fn run_benchmark(
+    suites: Vec<String>,
+    include_dirs: Vec<PathBuf>,
+) -> anyhow::Result<BenchmarkResults> {
+    let suites = EvaluationSuiteFactory::available_evaluations()
+        .into_iter()
+        .filter(|&s| suites.contains(&s.to_string()))
+        .collect::<Vec<_>>();
+
+    let config = Config::global();
+    let provider_name: String = config
+        .get("GOOSE_PROVIDER")
+        .expect("No provider configured. Run 'goose configure' first");
+
+    let mut results = BenchmarkResults::new(provider_name.clone());
+
+    let current_time = Local::now().format("%H:%M:%S").to_string();
+    let current_date = Local::now().format("%Y-%m-%d").to_string();
+    if let Ok(mut work_dir) = WorkDir::at(
+        format!("./benchmark-{}", &provider_name),
+        include_dirs.clone(),
+    ) {
+        if let Ok(work_dir) = work_dir.move_to(format!("./{}-{}", &current_date, current_time)) {
+            for suite in suites {
+                let suite_result = run_suite(suite, work_dir).await?;
+                results.add_suite(suite_result);
+            }
+        }
+    }
+
+    Ok(results)
+}
+
+pub async fn list_suites() -> anyhow::Result<HashMap<String, usize>> {
+    let suites = EvaluationSuiteFactory::available_evaluations();
+    let mut suite_counts = HashMap::new();
+
+    for suite in suites {
+        if let Some(evals) = EvaluationSuiteFactory::create(suite) {
+            suite_counts.insert(suite.to_string(), evals.len());
+        }
+    }
+
+    Ok(suite_counts)
+}

crates/goose-cli/src/commands/mod.rs

@@ -1,4 +1,5 @@
 pub mod agent_version;
+pub mod bench;
 pub mod configure;
 pub mod info;
 pub mod mcp;

crates/goose-cli/src/main.rs

@@ -2,12 +2,15 @@ use anyhow::Result;
 use clap::{Args, Parser, Subcommand};
 
 use goose::config::Config;
+
+use goose_cli::commands::agent_version::AgentCommand;
+use goose_cli::commands::bench::{list_suites, run_benchmark};
 use goose_cli::commands::configure::handle_configure;
 use goose_cli::commands::info::handle_info;
 use goose_cli::commands::mcp::run_server;
 use goose_cli::logging::setup_logging;
+use goose_cli::session;
 use goose_cli::session::build_session;
-use goose_cli::{commands::agent_version::AgentCommand, session};
 use std::io::{self, Read};
 use std::path::PathBuf;
 
@@ -194,6 +197,66 @@ enum Command {
         #[arg(short, long, help = "Enforce to re-configure goose during update")]
         reconfigure: bool,
     },
+
+    Bench {
+        #[arg(
+            short = 's',
+            long = "suites",
+            value_name = "BENCH_SUITE_NAME",
+            help = "Run this list of bench-suites.",
+            long_help = "Specify a comma-separated list of evaluation-suite names to be run.",
+            value_delimiter = ','
+        )]
+        suites: Vec<String>,
+
+        #[arg(
+            short = 'i',
+            long = "include-dir",
+            value_name = "DIR_NAME",
+            action = clap::ArgAction::Append,
+            long_help = "Make one or more dirs available to all bench suites. Specify either a single dir-name, a comma-separated list of dir-names, or use this multiple instances of this flag to specify multiple dirs.",
+            value_delimiter = ','
+        )]
+        include_dirs: Vec<PathBuf>,
+
+        #[arg(
+            long = "repeat",
+            value_name = "QUANTITY",
+            long_help = "Number of times to repeat the benchmark run.",
+            default_value = "1"
+        )]
+        repeat: usize,
+
+        #[arg(
+            long = "list",
+            value_name = "LIST",
+            help = "List all available bench suites."
+        )]
+        list: bool,
+
+        #[arg(
+            long = "output",
+            short = 'o',
+            value_name = "FILE",
+            help = "Save benchmark results to a file"
+        )]
+        output: Option<PathBuf>,
+
+        #[arg(
+            long = "format",
+            value_name = "FORMAT",
+            help = "Output format (text, json)",
+            default_value = "text"
+        )]
+        format: String,
+
+        #[arg(
+            long = "summary",
+            help = "Show only summary results",
+            action = clap::ArgAction::SetTrue
+        )]
+        summary: bool,
+    },
 }
 
 #[derive(clap::ValueEnum, Clone, Debug)]
@@ -232,6 +295,7 @@ async fn main() -> Result<()> {
                 builtin,
             )
             .await;
+
             setup_logging(session.session_file().file_stem().and_then(|s| s.to_str()))?;
             let _ = session.interactive(None).await;
             return Ok(());
@@ -290,6 +354,56 @@ async fn main() -> Result<()> {
             goose_cli::commands::update::update(canary, reconfigure)?;
             return Ok(());
         }
+        Some(Command::Bench {
+            suites,
+            include_dirs,
+            repeat,
+            list,
+            output,
+            format,
+            summary,
+        }) => {
+            if list {
+                let suites = list_suites().await?;
+                for suite in suites.keys() {
+                    println!("{}: {}", suite, suites.get(suite).unwrap());
+                }
+                return Ok(());
+            }
+            let suites = if suites.is_empty() {
+                vec!["core".to_string()]
+            } else {
+                suites
+            };
+            let current_dir = std::env::current_dir()?;
+
+            for i in 0..repeat {
+                if repeat > 1 {
+                    println!("\nRun {} of {}:", i + 1, repeat);
+                }
+                let results = run_benchmark(suites.clone(), include_dirs.clone()).await?;
+
+                // Handle output based on format
+                let output_str = match format.as_str() {
+                    "json" => serde_json::to_string_pretty(&results)?,
+                    _ => results.to_string(), // Uses Display impl
+                };
+
+                // Save to file if specified
+                if let Some(path) = &output {
+                    std::fs::write(current_dir.join(path), &output_str)?;
+                    println!("Results saved to: {}", path.display());
+                } else {
+                    // Print to console
+                    if summary {
+                        println!("{}", results.summary());
+                    } else {
+                        println!("{}", output_str);
+                    }
+                }
+            }
+            return Ok(());
+        }
         None => {
             if !Config::global().exists() {
                 let _ = handle_configure().await;

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

@@ -622,4 +622,8 @@ impl Session {
         cache.prompt_info.clear();
         cache.last_updated = Instant::now();
     }
+
+    pub fn message_history(&self) -> Vec<Message> {
+        self.messages.clone()
+    }
 }