From 2421f28436bafab223bf224892adfc2cbeaa8b69 Mon Sep 17 00:00:00 2001 From: Aaron Schlesinger <70865+arschles@users.noreply.github.com> Date: Mon, 10 Aug 2026 21:44:58 -0700 Subject: [PATCH 1/3] Adding `qt add` command --- README.md | 9 ++ cli/README.md | 9 ++ cli/src/cli.rs | 10 ++ cli/src/commands/add.rs | 246 +++++++++++++++++++++++++++++++++++++++ cli/src/commands/mod.rs | 2 + cli/src/main.rs | 10 +- cli/tests/add.rs | 248 ++++++++++++++++++++++++++++++++++++++++ 7 files changed, 533 insertions(+), 1 deletion(-) create mode 100644 cli/src/commands/add.rs create mode 100644 cli/tests/add.rs diff --git a/README.md b/README.md index 8cdfebf..8d64a8a 100644 --- a/README.md +++ b/README.md @@ -86,6 +86,7 @@ Common commands: ```bash qt --version +qt add qt run qt list qt show @@ -117,6 +118,14 @@ Registry benchmarks are ready-to-run evaluations with predefined datasets, scori qt run simpleqa-verified --input '{"model":"random","limit":10}' ``` +To save a registry benchmark in the local configuration, add it by name: + +```bash +qt add simpleqa-verified +``` + +This command downloads the benchmark definition and prompt template, appends the benchmark to an existing `quantiles.toml` or `.quantiles.toml`, or creates `quantiles.toml` in the current directory. It returns an error if the benchmark is already configured or is not present in the registry. Pass `--json` for machine-readable output. + The [benchmark hub](https://quantiles.io/benchmark-hub) describes available benchmarks, their evaluation setup, and common metrics used across AI evaluation workflows. > To request another registry benchmark, [file an issue](https://github.com/quantiles-evals/quantiles/issues) with its name, source dataset or repository, and any available reference implementation. diff --git a/cli/README.md b/cli/README.md index c3ff939..bb519b5 100644 --- a/cli/README.md +++ b/cli/README.md @@ -74,6 +74,15 @@ See the [configuration guide](https://quantiles.io/documentation/configuration) ### Remote benchmark fallback +Use `qt add ` to download a registry benchmark and save it in the local configuration. If `quantiles.toml` or `.quantiles.toml` exists in the current directory, the command appends the benchmark section without rewriting the existing content. Otherwise, it creates `quantiles.toml`. The downloaded prompt template is stored under `.quantiles/registry/` and referenced by the added configuration. + +```bash +qt add simpleqa-verified +qt add simpleqa-verified --json +``` + +The command fails without modifying the configuration if the benchmark is already configured or the registry does not contain it. If you pass `--json` to this command, output of any kind (successful or unsuccessful), returns machine-readable JSON. Resolving and downloading the benchmark requires network access. + When you run `qt run `, the CLI first looks in the local configuration file for an evaluation called `eval_name`. If one is found, the CLI runs it immediately. If none is found, `qt` looks in the Quantiles remote benchmark service for an evaluation of the same name. If a match is found, the CLI downloads the benchmark definition and runs it. >If you want to override the location of the remote benchmark service, use the `--remote-url` flag or the `QUANTILES_REMOTE_URL` environment variable. diff --git a/cli/src/cli.rs b/cli/src/cli.rs index 7cc8e03..835cbac 100644 --- a/cli/src/cli.rs +++ b/cli/src/cli.rs @@ -28,6 +28,16 @@ pub struct Cli { #[derive(Debug, Subcommand)] pub enum Command { + /// Add a benchmark from the remote registry to the local configuration. + Add { + benchmark_name: String, + /// Override the remote benchmark service URL. + #[arg(long)] + remote_url: Option, + /// Emit machine-readable JSON. + #[arg(long)] + json: bool, + }, /// Initialize or update a local Quantiles workspace. Init, /// Show a list of all eval runs. diff --git a/cli/src/commands/add.rs b/cli/src/commands/add.rs new file mode 100644 index 0000000..d9d7603 --- /dev/null +++ b/cli/src/commands/add.rs @@ -0,0 +1,246 @@ +use std::fs::{self, OpenOptions}; +use std::io::Write as _; +use std::path::{Path, PathBuf}; + +use anyhow::{Context, Result, bail}; +use serde::Serialize; + +pub async fn add(benchmark_name: &str, cli_remote_url: Option<&str>, json: bool) -> Result<()> { + let cwd = std::env::current_dir().context("failed to determine current directory")?; + let config_path = config_path(&cwd)?; + let existing = read_config(config_path.as_deref())?; + + if existing + .as_ref() + .is_some_and(|config| config.benchmarks.contains_key(benchmark_name)) + { + bail!("benchmark `{benchmark_name}` is already present in the local configuration"); + } + + let remote_url = qt::benchmark_registry::select_remote_url( + cli_remote_url, + std::env::var_os("QUANTILES_REMOTE_URL"), + )?; + let remote = qt::benchmark_registry::resolve_and_download(benchmark_name, None, &remote_url) + .await? + .with_context(|| { + format!("benchmark `{benchmark_name}` was not found in the remote registry") + })?; + + let config_path = config_path.unwrap_or_else(|| cwd.join("quantiles.toml")); + let version = remote.version.clone(); + persist_remote_benchmark(benchmark_name, remote, &config_path, &cwd)?; + + if json { + println!( + "{}", + serde_json::to_string(&AddOutput { + benchmark_name, + version: &version, + config_path: config_path.display().to_string(), + })? + ); + } else { + println!( + "Added benchmark `{benchmark_name}` (version {version}) to {}", + config_path.display() + ); + } + Ok(()) +} + +#[derive(Serialize)] +struct AddOutput<'a> { + benchmark_name: &'a str, + version: &'a str, + config_path: String, +} + +fn config_path(cwd: &Path) -> Result> { + let plain = cwd.join("quantiles.toml"); + let dot = cwd.join(".quantiles.toml"); + match (plain.exists(), dot.exists()) { + (true, true) => bail!( + "both `quantiles.toml` and `.quantiles.toml` found in {}. remove one to avoid ambiguity", + cwd.display() + ), + (true, false) => Ok(Some(plain)), + (false, true) => Ok(Some(dot)), + (false, false) => Ok(None), + } +} + +fn read_config(path: Option<&Path>) -> Result> { + let Some(path) = path else { + return Ok(None); + }; + let contents = + fs::read_to_string(path).with_context(|| format!("failed to read {}", path.display()))?; + let config = + toml::from_str(&contents).with_context(|| format!("failed to parse {}", path.display()))?; + Ok(Some(config)) +} + +fn persist_remote_benchmark( + benchmark_name: &str, + mut remote: qt::benchmark_registry::RemoteBenchmark, + config_path: &Path, + cwd: &Path, +) -> Result<()> { + let prompt_relative = PathBuf::from(".quantiles") + .join("registry") + .join(&remote.manifest_sha256) + .join("prompt.txt"); + let prompt_path = cwd.join(&prompt_relative); + persist_prompt(&prompt_path, remote.prompt_template.as_bytes())?; + remote.config.params.prompt_template_file = path_for_toml(&prompt_relative); + + let section = render_benchmark_section(benchmark_name, &remote.config)?; + append_section(config_path, §ion) +} + +fn persist_prompt(path: &Path, contents: &[u8]) -> Result<()> { + let parent = path + .parent() + .context("prompt path has no parent directory")?; + fs::create_dir_all(parent).with_context(|| format!("failed to create {}", parent.display()))?; + match OpenOptions::new().write(true).create_new(true).open(path) { + Ok(mut file) => file + .write_all(contents) + .with_context(|| format!("failed to write {}", path.display())), + Err(error) if error.kind() == std::io::ErrorKind::AlreadyExists => { + let existing = fs::read(path) + .with_context(|| format!("failed to read existing {}", path.display()))?; + if existing == contents { + Ok(()) + } else { + bail!( + "existing registry prompt {} has unexpected contents", + path.display() + ) + } + } + Err(error) => Err(error).with_context(|| format!("failed to create {}", path.display())), + } +} + +fn render_benchmark_section( + benchmark_name: &str, + config: &qt::config::CustomNoCodeBenchmarkConfig, +) -> Result { + let mut benchmark = toml::Value::try_from(&config.params) + .context("failed to serialize remote benchmark configuration")? + .as_table() + .cloned() + .context("remote benchmark configuration did not serialize to a TOML table")?; + benchmark.insert( + "type".to_owned(), + toml::Value::String("custom_nocode".to_owned()), + ); + + let mut benchmarks = toml::map::Map::new(); + benchmarks.insert(benchmark_name.to_owned(), toml::Value::Table(benchmark)); + let mut root = toml::map::Map::new(); + root.insert("benchmarks".to_owned(), toml::Value::Table(benchmarks)); + toml::to_string_pretty(&toml::Value::Table(root)) + .context("failed to render remote benchmark configuration as TOML") +} + +fn append_section(path: &Path, section: &str) -> Result<()> { + let existing = if path.exists() { + fs::read_to_string(path).with_context(|| format!("failed to read {}", path.display()))? + } else { + String::new() + }; + let separator = if existing.is_empty() || existing.ends_with("\n\n") { + "" + } else if existing.ends_with('\n') { + "\n" + } else { + "\n\n" + }; + + let mut file = OpenOptions::new() + .create(true) + .append(true) + .open(path) + .with_context(|| format!("failed to open {}", path.display()))?; + write!(file, "{separator}{section}") + .with_context(|| format!("failed to append benchmark to {}", path.display())) +} + +fn path_for_toml(path: &Path) -> String { + path.components() + .map(|component| component.as_os_str().to_string_lossy()) + .collect::>() + .join("/") +} + +#[cfg(test)] +mod tests { + use super::*; + use qt::config::{CustomNoCodeDatasetConfig, CustomNoCodeParams, CustomNoCodeStyleConfig}; + + fn remote() -> qt::benchmark_registry::RemoteBenchmark { + qt::benchmark_registry::RemoteBenchmark { + config: qt::config::CustomNoCodeBenchmarkConfig { + type_: "custom_nocode".to_owned(), + params: CustomNoCodeParams { + dataset: CustomNoCodeDatasetConfig { + name: "quantiles/example".to_owned(), + config_name: None, + split: None, + revision: None, + }, + model: None, + prompt_template_file: "prompts/qa.txt".to_owned(), + limit: None, + max_workers: None, + metrics: vec![], + style: CustomNoCodeStyleConfig::ExactMatch { + golden_column: "answer".to_owned(), + }, + }, + }, + prompt_template: "{{ row.question }}".to_owned(), + version: "v1".to_owned(), + manifest_sha256: "a".repeat(64), + } + } + + #[test] + fn creates_runnable_configuration_and_prompt() { + let temp = tempfile::tempdir().unwrap(); + let config_path = temp.path().join("quantiles.toml"); + + persist_remote_benchmark("remote-test", remote(), &config_path, temp.path()).unwrap(); + + let contents = fs::read_to_string(&config_path).unwrap(); + let parsed: qt::config::WorkspaceConfig = toml::from_str(&contents).unwrap(); + assert!(parsed.benchmarks.contains_key("remote-test")); + assert_eq!( + fs::read_to_string( + temp.path() + .join(".quantiles/registry") + .join("a".repeat(64)) + .join("prompt.txt") + ) + .unwrap(), + "{{ row.question }}" + ); + } + + #[test] + fn appends_after_existing_configuration_without_rewriting_it() { + let temp = tempfile::tempdir().unwrap(); + let config_path = temp.path().join("quantiles.toml"); + let original = "# keep this comment\n[benchmarks.local]\ntype = \"custom_code\"\ncommand = [\"echo\"]\n"; + fs::write(&config_path, original).unwrap(); + + persist_remote_benchmark("remote-test", remote(), &config_path, temp.path()).unwrap(); + + let contents = fs::read_to_string(config_path).unwrap(); + assert!(contents.starts_with(original)); + assert!(contents.contains("[benchmarks.remote-test]")); + } +} diff --git a/cli/src/commands/mod.rs b/cli/src/commands/mod.rs index 9225f49..710bae8 100644 --- a/cli/src/commands/mod.rs +++ b/cli/src/commands/mod.rs @@ -1,3 +1,4 @@ +mod add; mod compare; mod custom_nocode_metrics; mod init; @@ -7,6 +8,7 @@ mod run; mod serve; mod show; +pub use add::add; pub use compare::compare; pub use init::init; pub use list::list; diff --git a/cli/src/main.rs b/cli/src/main.rs index c06251b..2316ac8 100644 --- a/cli/src/main.rs +++ b/cli/src/main.rs @@ -9,7 +9,10 @@ use clap::Parser; fn main() -> ExitCode { let cli = cli::Cli::parse(); - let json_errors = matches!(&cli.command, Some(cli::Command::Resume { json: true, .. })); + let json_errors = matches!( + &cli.command, + Some(cli::Command::Add { json: true, .. } | cli::Command::Resume { json: true, .. }) + ); match try_main(cli) { Ok(()) => ExitCode::SUCCESS, @@ -50,6 +53,11 @@ async fn async_main(cli: cli::Cli, process_start: Instant) -> Result<()> { } match cli.command.expect("clap requires a subcommand") { + cli::Command::Add { + benchmark_name, + remote_url, + json, + } => commands::add(&benchmark_name, remote_url.as_deref(), json).await, cli::Command::Init => commands::init().await, cli::Command::List { json } => commands::list(json).await, cli::Command::Compare { run_a, run_b, json } => commands::compare(run_a, run_b, json).await, diff --git a/cli/tests/add.rs b/cli/tests/add.rs new file mode 100644 index 0000000..b8e933d --- /dev/null +++ b/cli/tests/add.rs @@ -0,0 +1,248 @@ +use assert_cmd::Command; +use predicates::prelude::*; +use sha2::{Digest as _, Sha256}; +use wiremock::matchers::{method, path}; +use wiremock::{Mock, MockServer, ResponseTemplate}; + +#[expect( + clippy::allow_attributes, + clippy::pedantic, + reason = "ConnectRPC and Buffa generated code uses allow attributes" +)] +mod registry_proto { + connectrpc::include_generated!(); +} + +const CONFIG: &str = r#"[benchmarks.existing] +type = "custom_code" +command = ["echo"] +"#; + +#[test] +fn duplicate_benchmark_error_is_json_when_requested() { + let temp = tempfile::tempdir().unwrap(); + std::fs::write(temp.path().join("quantiles.toml"), CONFIG).unwrap(); + + let expected = format!( + "{}\n", + serde_json::json!({ + "error": "benchmark `existing` is already present in the local configuration" + }) + ); + Command::new(assert_cmd::cargo::cargo_bin!("qt")) + .current_dir(temp.path()) + .args(["add", "existing", "--json"]) + .assert() + .failure() + .stdout(predicate::eq(expected)) + .stderr(predicate::str::is_empty()); + + assert_eq!( + std::fs::read_to_string(temp.path().join("quantiles.toml")).unwrap(), + CONFIG + ); +} + +#[test] +fn duplicate_benchmark_error_is_human_readable_by_default() { + let temp = tempfile::tempdir().unwrap(); + std::fs::write(temp.path().join("quantiles.toml"), CONFIG).unwrap(); + + Command::new(assert_cmd::cargo::cargo_bin!("qt")) + .current_dir(temp.path()) + .args(["add", "existing"]) + .assert() + .failure() + .stdout(predicate::str::is_empty()) + .stderr(predicate::eq( + "Error: benchmark `existing` is already present in the local configuration\n", + )); +} + +#[tokio::test(flavor = "multi_thread")] +async fn missing_remote_benchmark_returns_json_and_does_not_create_config() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path( + "/quantiles.benchmark.v1.BenchmarkRegistryService/ResolveBenchmark", + )) + .respond_with( + ResponseTemplate::new(404) + .insert_header("content-type", "application/json") + .set_body_raw( + r#"{"code":"not_found","message":"benchmark does not exist"}"#, + "application/json", + ), + ) + .mount(&server) + .await; + let temp = tempfile::tempdir().unwrap(); + + let expected = format!( + "{}\n", + serde_json::json!({ + "error": "benchmark `missing` was not found in the remote registry" + }) + ); + Command::new(assert_cmd::cargo::cargo_bin!("qt")) + .current_dir(temp.path()) + .env("QUANTILES_REMOTE_URL", server.uri()) + .args(["add", "missing", "--json"]) + .assert() + .failure() + .stdout(predicate::eq(expected)) + .stderr(predicate::str::is_empty()); + + assert!(!temp.path().join("quantiles.toml").exists()); + assert!(!temp.path().join(".quantiles").exists()); +} + +#[tokio::test(flavor = "multi_thread")] +async fn successful_add_creates_config_and_emits_json_only() { + let server = MockServer::start().await; + let manifest_sha256 = mock_successful_registry(&server).await; + let temp = tempfile::tempdir().unwrap(); + let config_path = temp.path().canonicalize().unwrap().join("quantiles.toml"); + + let output = Command::new(assert_cmd::cargo::cargo_bin!("qt")) + .current_dir(temp.path()) + .env("QUANTILES_REMOTE_URL", server.uri()) + .args(["add", "remote-test", "--json"]) + .output() + .unwrap(); + + assert!(output.status.success()); + assert!(output.stderr.is_empty()); + let stdout: serde_json::Value = serde_json::from_slice(&output.stdout).unwrap(); + assert_eq!( + stdout, + serde_json::json!({ + "benchmark_name": "remote-test", + "version": "v1", + "config_path": config_path.display().to_string(), + }) + ); + assert_added_files(temp.path(), &manifest_sha256, None); +} + +#[tokio::test(flavor = "multi_thread")] +async fn successful_add_appends_config_and_emits_human_output_only() { + let server = MockServer::start().await; + let manifest_sha256 = mock_successful_registry(&server).await; + let temp = tempfile::tempdir().unwrap(); + let config_path = temp.path().canonicalize().unwrap().join("quantiles.toml"); + let original = "# preserve this comment\n[benchmarks.existing]\ntype = \"custom_code\"\ncommand = [\"echo\"]\n"; + std::fs::write(&config_path, original).unwrap(); + + Command::new(assert_cmd::cargo::cargo_bin!("qt")) + .current_dir(temp.path()) + .env("QUANTILES_REMOTE_URL", server.uri()) + .args(["add", "remote-test"]) + .assert() + .success() + .stdout(predicate::eq(format!( + "Added benchmark `remote-test` (version v1) to {}\n", + config_path.display() + ))) + .stderr(predicate::str::is_empty()); + + assert_added_files(temp.path(), &manifest_sha256, Some(original)); +} + +async fn mock_successful_registry(server: &MockServer) -> String { + use buffa::Message as _; + use registry_proto::quantiles::benchmark::v1::{ + BenchmarkResource, ResolveBenchmarkResponse, ResourceKind, + }; + + let definition = br#"[benchmarks.remote-test] +type = "custom_nocode" +dataset = { name = "quantiles/example" } +prompt_template_file = "prompts/qa.txt" +style = { type = "exact_match", golden_column = "answer" } +"#; + let prompt = b"{{ row.question }}\nAnswer:"; + let manifest_sha256 = "a".repeat(64); + let resource = + |id: &str, logical_path: &str, kind: ResourceKind, route: &str, contents: &[u8]| { + BenchmarkResource { + resource_id: id.to_owned(), + logical_path: logical_path.to_owned(), + kind: kind.into(), + download_url: format!("{}{route}", server.uri()), + sha256: format!("{:x}", Sha256::digest(contents)), + size_bytes: u64::try_from(contents.len()).unwrap(), + content_type: "application/octet-stream".to_owned(), + ..Default::default() + } + }; + let response = ResolveBenchmarkResponse { + benchmark_name: "remote-test".to_owned(), + version: "v1".to_owned(), + manifest_sha256: manifest_sha256.clone(), + resources: vec![ + resource( + "definition", + "bundle/quantiles.toml", + ResourceKind::Definition, + "/definition", + definition, + ), + resource( + "prompt", + "bundle/prompts/qa.txt", + ResourceKind::PromptTemplate, + "/prompt", + prompt, + ), + ], + ..Default::default() + }; + + Mock::given(method("POST")) + .and(path( + "/quantiles.benchmark.v1.BenchmarkRegistryService/ResolveBenchmark", + )) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "application/proto") + .set_body_bytes(response.encode_to_vec()), + ) + .expect(1) + .mount(server) + .await; + for (route, contents) in [ + ("/definition", definition.as_slice()), + ("/prompt", prompt.as_slice()), + ] { + Mock::given(method("GET")) + .and(path(route)) + .respond_with(ResponseTemplate::new(200).set_body_bytes(contents)) + .expect(1) + .mount(server) + .await; + } + + manifest_sha256 +} + +fn assert_added_files(root: &std::path::Path, manifest_sha256: &str, prefix: Option<&str>) { + let config_contents = std::fs::read_to_string(root.join("quantiles.toml")).unwrap(); + if let Some(prefix) = prefix { + assert!(config_contents.starts_with(prefix)); + } + let config: qt::config::WorkspaceConfig = toml::from_str(&config_contents).unwrap(); + assert!(config.benchmarks.contains_key("remote-test")); + if prefix.is_some() { + assert!(config.benchmarks.contains_key("existing")); + } + + let prompt_path = root + .join(".quantiles/registry") + .join(manifest_sha256) + .join("prompt.txt"); + assert_eq!( + std::fs::read_to_string(prompt_path).unwrap(), + "{{ row.question }}\nAnswer:" + ); +} From 363ab9a78a559273111341f18a686f43333cc054 Mon Sep 17 00:00:00 2001 From: Aaron Schlesinger <70865+arschles@users.noreply.github.com> Date: Mon, 10 Aug 2026 21:50:59 -0700 Subject: [PATCH 2/3] Saving prompt templates next to the quantiles.toml file --- README.md | 2 +- cli/README.md | 2 +- cli/src/commands/add.rs | 55 +++++++++++++++++++++++++++++------------ cli/tests/add.rs | 29 +++++++++++++--------- 4 files changed, 58 insertions(+), 30 deletions(-) diff --git a/README.md b/README.md index 8d64a8a..78d3ac6 100644 --- a/README.md +++ b/README.md @@ -124,7 +124,7 @@ To save a registry benchmark in the local configuration, add it by name: qt add simpleqa-verified ``` -This command downloads the benchmark definition and prompt template, appends the benchmark to an existing `quantiles.toml` or `.quantiles.toml`, or creates `quantiles.toml` in the current directory. It returns an error if the benchmark is already configured or is not present in the registry. Pass `--json` for machine-readable output. +This command downloads the benchmark definition and prompt template, appends the benchmark to an existing `quantiles.toml` or `.quantiles.toml`, or creates `quantiles.toml` in the current directory. The prompt template is stored beside the configuration at `-prompt/prompt.txt`. The command returns an error if the benchmark is already configured or is not present in the registry. Pass `--json` for machine-readable output. The [benchmark hub](https://quantiles.io/benchmark-hub) describes available benchmarks, their evaluation setup, and common metrics used across AI evaluation workflows. diff --git a/cli/README.md b/cli/README.md index bb519b5..3176bba 100644 --- a/cli/README.md +++ b/cli/README.md @@ -74,7 +74,7 @@ See the [configuration guide](https://quantiles.io/documentation/configuration) ### Remote benchmark fallback -Use `qt add ` to download a registry benchmark and save it in the local configuration. If `quantiles.toml` or `.quantiles.toml` exists in the current directory, the command appends the benchmark section without rewriting the existing content. Otherwise, it creates `quantiles.toml`. The downloaded prompt template is stored under `.quantiles/registry/` and referenced by the added configuration. +Use `qt add ` to download a registry benchmark and save it in the local configuration. If `quantiles.toml` or `.quantiles.toml` exists in the current directory, the command appends the benchmark section without rewriting the existing content. Otherwise, it creates `quantiles.toml`. The downloaded prompt template is stored beside the configuration at `-prompt/prompt.txt` and referenced by the added configuration. ```bash qt add simpleqa-verified diff --git a/cli/src/commands/add.rs b/cli/src/commands/add.rs index d9d7603..9fcdb69 100644 --- a/cli/src/commands/add.rs +++ b/cli/src/commands/add.rs @@ -29,7 +29,7 @@ pub async fn add(benchmark_name: &str, cli_remote_url: Option<&str>, json: bool) let config_path = config_path.unwrap_or_else(|| cwd.join("quantiles.toml")); let version = remote.version.clone(); - persist_remote_benchmark(benchmark_name, remote, &config_path, &cwd)?; + persist_remote_benchmark(benchmark_name, remote, &config_path)?; if json { println!( @@ -85,13 +85,12 @@ fn persist_remote_benchmark( benchmark_name: &str, mut remote: qt::benchmark_registry::RemoteBenchmark, config_path: &Path, - cwd: &Path, ) -> Result<()> { - let prompt_relative = PathBuf::from(".quantiles") - .join("registry") - .join(&remote.manifest_sha256) - .join("prompt.txt"); - let prompt_path = cwd.join(&prompt_relative); + let prompt_relative = prompt_relative_path(benchmark_name)?; + let config_dir = config_path + .parent() + .context("configuration path has no parent directory")?; + let prompt_path = config_dir.join(&prompt_relative); persist_prompt(&prompt_path, remote.prompt_template.as_bytes())?; remote.config.params.prompt_template_file = path_for_toml(&prompt_relative); @@ -99,6 +98,16 @@ fn persist_remote_benchmark( append_section(config_path, §ion) } +fn prompt_relative_path(benchmark_name: &str) -> Result { + let unsafe_character = benchmark_name + .chars() + .any(|character| character.is_control() || r#"/\:*?"<>|"#.contains(character)); + if benchmark_name.is_empty() || matches!(benchmark_name, "." | "..") || unsafe_character { + bail!("benchmark name `{benchmark_name}` cannot be used for a local prompt directory"); + } + Ok(PathBuf::from(format!("{benchmark_name}-prompt")).join("prompt.txt")) +} + fn persist_prompt(path: &Path, contents: &[u8]) -> Result<()> { let parent = path .parent() @@ -213,21 +222,24 @@ mod tests { let temp = tempfile::tempdir().unwrap(); let config_path = temp.path().join("quantiles.toml"); - persist_remote_benchmark("remote-test", remote(), &config_path, temp.path()).unwrap(); + persist_remote_benchmark("remote-test", remote(), &config_path).unwrap(); let contents = fs::read_to_string(&config_path).unwrap(); let parsed: qt::config::WorkspaceConfig = toml::from_str(&contents).unwrap(); assert!(parsed.benchmarks.contains_key("remote-test")); assert_eq!( - fs::read_to_string( - temp.path() - .join(".quantiles/registry") - .join("a".repeat(64)) - .join("prompt.txt") - ) - .unwrap(), + fs::read_to_string(temp.path().join("remote-test-prompt/prompt.txt")).unwrap(), "{{ row.question }}" ); + let qt::config::BenchmarkConfig::CustomNoCode(config) = + parsed.benchmarks.get("remote-test").unwrap() + else { + panic!("expected custom_nocode benchmark"); + }; + assert_eq!( + config.params.prompt_template_file, + "remote-test-prompt/prompt.txt" + ); } #[test] @@ -237,10 +249,21 @@ mod tests { let original = "# keep this comment\n[benchmarks.local]\ntype = \"custom_code\"\ncommand = [\"echo\"]\n"; fs::write(&config_path, original).unwrap(); - persist_remote_benchmark("remote-test", remote(), &config_path, temp.path()).unwrap(); + persist_remote_benchmark("remote-test", remote(), &config_path).unwrap(); let contents = fs::read_to_string(config_path).unwrap(); assert!(contents.starts_with(original)); assert!(contents.contains("[benchmarks.remote-test]")); } + + #[test] + fn rejects_benchmark_names_that_cannot_form_a_safe_prompt_directory() { + for name in ["../outside", "nested/name", r"nested\name"] { + let temp = tempfile::tempdir().unwrap(); + let config_path = temp.path().join("quantiles.toml"); + let error = persist_remote_benchmark(name, remote(), &config_path).unwrap_err(); + assert!(error.to_string().contains("local prompt directory")); + assert!(!config_path.exists()); + } + } } diff --git a/cli/tests/add.rs b/cli/tests/add.rs index b8e933d..7dbc8e6 100644 --- a/cli/tests/add.rs +++ b/cli/tests/add.rs @@ -95,12 +95,13 @@ async fn missing_remote_benchmark_returns_json_and_does_not_create_config() { assert!(!temp.path().join("quantiles.toml").exists()); assert!(!temp.path().join(".quantiles").exists()); + assert!(!temp.path().join("missing-prompt").exists()); } #[tokio::test(flavor = "multi_thread")] async fn successful_add_creates_config_and_emits_json_only() { let server = MockServer::start().await; - let manifest_sha256 = mock_successful_registry(&server).await; + mock_successful_registry(&server).await; let temp = tempfile::tempdir().unwrap(); let config_path = temp.path().canonicalize().unwrap().join("quantiles.toml"); @@ -122,13 +123,13 @@ async fn successful_add_creates_config_and_emits_json_only() { "config_path": config_path.display().to_string(), }) ); - assert_added_files(temp.path(), &manifest_sha256, None); + assert_added_files(temp.path(), None); } #[tokio::test(flavor = "multi_thread")] async fn successful_add_appends_config_and_emits_human_output_only() { let server = MockServer::start().await; - let manifest_sha256 = mock_successful_registry(&server).await; + mock_successful_registry(&server).await; let temp = tempfile::tempdir().unwrap(); let config_path = temp.path().canonicalize().unwrap().join("quantiles.toml"); let original = "# preserve this comment\n[benchmarks.existing]\ntype = \"custom_code\"\ncommand = [\"echo\"]\n"; @@ -146,10 +147,10 @@ async fn successful_add_appends_config_and_emits_human_output_only() { ))) .stderr(predicate::str::is_empty()); - assert_added_files(temp.path(), &manifest_sha256, Some(original)); + assert_added_files(temp.path(), Some(original)); } -async fn mock_successful_registry(server: &MockServer) -> String { +async fn mock_successful_registry(server: &MockServer) { use buffa::Message as _; use registry_proto::quantiles::benchmark::v1::{ BenchmarkResource, ResolveBenchmarkResponse, ResourceKind, @@ -222,11 +223,9 @@ style = { type = "exact_match", golden_column = "answer" } .mount(server) .await; } - - manifest_sha256 } -fn assert_added_files(root: &std::path::Path, manifest_sha256: &str, prefix: Option<&str>) { +fn assert_added_files(root: &std::path::Path, prefix: Option<&str>) { let config_contents = std::fs::read_to_string(root.join("quantiles.toml")).unwrap(); if let Some(prefix) = prefix { assert!(config_contents.starts_with(prefix)); @@ -236,11 +235,17 @@ fn assert_added_files(root: &std::path::Path, manifest_sha256: &str, prefix: Opt if prefix.is_some() { assert!(config.benchmarks.contains_key("existing")); } + let qt::config::BenchmarkConfig::CustomNoCode(remote) = + config.benchmarks.get("remote-test").unwrap() + else { + panic!("expected custom_nocode benchmark"); + }; + assert_eq!( + remote.params.prompt_template_file, + "remote-test-prompt/prompt.txt" + ); - let prompt_path = root - .join(".quantiles/registry") - .join(manifest_sha256) - .join("prompt.txt"); + let prompt_path = root.join("remote-test-prompt/prompt.txt"); assert_eq!( std::fs::read_to_string(prompt_path).unwrap(), "{{ row.question }}\nAnswer:" From d2f65ef977c2c1d75ab84bad9280c57de2ee999f Mon Sep 17 00:00:00 2001 From: Aaron Schlesinger <70865+arschles@users.noreply.github.com> Date: Mon, 10 Aug 2026 22:27:17 -0700 Subject: [PATCH 3/3] making error output json when requested --- cli/src/main.rs | 28 +++++++++++++++++++++++++++- cli/tests/add.rs | 29 +++++++++++++++++++++++++++++ 2 files changed, 56 insertions(+), 1 deletion(-) diff --git a/cli/src/main.rs b/cli/src/main.rs index 2316ac8..7588f4a 100644 --- a/cli/src/main.rs +++ b/cli/src/main.rs @@ -1,6 +1,7 @@ mod cli; mod commands; +use std::ffi::OsStr; use std::process::ExitCode; use std::time::Instant; @@ -8,7 +9,10 @@ use anyhow::Result; use clap::Parser; fn main() -> ExitCode { - let cli = cli::Cli::parse(); + let cli = match cli::Cli::try_parse() { + Ok(cli) => cli, + Err(error) => return handle_parse_error(&error), + }; let json_errors = matches!( &cli.command, Some(cli::Command::Add { json: true, .. } | cli::Command::Resume { json: true, .. }) @@ -27,6 +31,28 @@ fn main() -> ExitCode { } } +fn handle_parse_error(error: &clap::Error) -> ExitCode { + let exit_code = ExitCode::from(u8::try_from(error.exit_code()).unwrap_or(1)); + if add_json_requested() { + let message = error.to_string(); + if error.use_stderr() { + println!("{}", serde_json::json!({ "error": message.trim_end() })); + } else { + println!("{}", serde_json::json!({ "output": message.trim_end() })); + } + } else if let Err(print_error) = error.print() { + eprintln!("Error: failed to print command-line error: {print_error}"); + } + exit_code +} + +fn add_json_requested() -> bool { + let mut args = std::env::args_os().skip(1); + args.next() + .is_some_and(|argument| argument == OsStr::new("add")) + && args.any(|argument| argument == OsStr::new("--json")) +} + fn try_main(cli: cli::Cli) -> Result<()> { // `fastembed`'s dependency tree enables another Rustls crypto provider alongside AWS-LC. // Install AWS-LC explicitly, to avoid Rustls-related panics when multiple providers are diff --git a/cli/tests/add.rs b/cli/tests/add.rs index 7dbc8e6..5192415 100644 --- a/cli/tests/add.rs +++ b/cli/tests/add.rs @@ -18,6 +18,22 @@ type = "custom_code" command = ["echo"] "#; +#[test] +fn missing_benchmark_argument_error_is_json_when_requested() { + assert_json_parse_error( + &["add", "--json"], + "the following required arguments were not provided", + ); +} + +#[test] +fn unknown_option_error_is_json_when_requested() { + assert_json_parse_error( + &["add", "remote-test", "--json", "--unknown"], + "unexpected argument '--unknown'", + ); +} + #[test] fn duplicate_benchmark_error_is_json_when_requested() { let temp = tempfile::tempdir().unwrap(); @@ -98,6 +114,19 @@ async fn missing_remote_benchmark_returns_json_and_does_not_create_config() { assert!(!temp.path().join("missing-prompt").exists()); } +fn assert_json_parse_error(args: &[&str], expected_message: &str) { + let output = Command::new(assert_cmd::cargo::cargo_bin!("qt")) + .args(args) + .output() + .unwrap(); + + assert_eq!(output.status.code(), Some(2)); + assert!(output.stderr.is_empty()); + let stdout: serde_json::Value = serde_json::from_slice(&output.stdout).unwrap(); + assert_eq!(stdout.as_object().unwrap().len(), 1); + assert!(stdout["error"].as_str().unwrap().contains(expected_message)); +} + #[tokio::test(flavor = "multi_thread")] async fn successful_add_creates_config_and_emits_json_only() { let server = MockServer::start().await;