From d27b0fd303cf792a8113ea6bccf79cfdc17898c9 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 1 Jul 2026 09:32:21 +0000 Subject: [PATCH 1/2] Harden qr init secret handling and config recovery Co-authored-by: Aanish Bhirud --- src/atomic.rs | 79 +++++++++++++++++++++++++++++++++----- src/config.rs | 5 ++- src/main.rs | 63 +++++++++++++++---------------- src/secret.rs | 13 +++++++ tests/cli.rs | 102 ++++++++++++++++++++++++++++++++++++++++++++++++-- 5 files changed, 213 insertions(+), 49 deletions(-) diff --git a/src/atomic.rs b/src/atomic.rs index 7121a62..8018f11 100644 --- a/src/atomic.rs +++ b/src/atomic.rs @@ -7,6 +7,8 @@ //! concurrently: the project cache (read by `qr go` while the hourly cron //! rewrites it) and the user's shell rc file (corrupting it is a 5-alarm fire). +#[cfg(unix)] +use std::os::unix::fs::{OpenOptionsExt, PermissionsExt}; use std::{ ffi::OsString, fs, @@ -19,6 +21,21 @@ use anyhow::{Context, Result}; /// Atomically replace `path` with `contents`. pub fn write(path: &Path, contents: &[u8]) -> Result<()> { + write_impl(path, contents, WriteMode::PreserveExisting) +} + +/// Atomically replace `path` with `contents`, creating the file with private +/// permissions (`0600`) from the start on Unix. +pub fn write_private(path: &Path, contents: &[u8]) -> Result<()> { + write_impl(path, contents, WriteMode::Private) +} + +enum WriteMode { + PreserveExisting, + Private, +} + +fn write_impl(path: &Path, contents: &[u8], mode: WriteMode) -> Result<()> { // Follow a final symlink so we replace its target — as `fs::write` did — // rather than clobbering the symlink itself with a regular file. This keeps // dotfiles-managed symlinked rc files (e.g. `~/.zshrc` -> a versioned repo) @@ -32,8 +49,7 @@ pub fn write(path: &Path, contents: &[u8]) -> Result<()> { } let tmp = temp_path(target); - let mut file = fs::File::create(&tmp) - .with_context(|| format!("Failed to create temp file {}", tmp.display()))?; + let mut file = create_temp_file(&tmp, &mode)?; let result = file.write_all(contents).and_then(|()| file.sync_all()); drop(file); if let Err(error) = result { @@ -45,12 +61,24 @@ pub fn write(path: &Path, contents: &[u8]) -> Result<()> { // installs a fresh inode, so without this an existing rc/config file's mode // (e.g. a 0600 config) would reset to the default. Fail loudly rather than // silently downgrade the mode. - if let Ok(metadata) = fs::metadata(target) { - if let Err(error) = fs::set_permissions(&tmp, metadata.permissions()) { - let _ = fs::remove_file(&tmp); - return Err(error).with_context(|| { - format!("Failed to preserve permissions on {}", target.display()) - }); + match mode { + WriteMode::PreserveExisting => { + if let Ok(metadata) = fs::metadata(target) { + if let Err(error) = fs::set_permissions(&tmp, metadata.permissions()) { + let _ = fs::remove_file(&tmp); + return Err(error).with_context(|| { + format!("Failed to preserve permissions on {}", target.display()) + }); + } + } + } + WriteMode::Private => { + #[cfg(unix)] + if let Err(error) = fs::set_permissions(&tmp, fs::Permissions::from_mode(0o600)) { + let _ = fs::remove_file(&tmp); + return Err(error) + .with_context(|| format!("Failed to secure {}", target.display())); + } } } @@ -61,6 +89,27 @@ pub fn write(path: &Path, contents: &[u8]) -> Result<()> { Ok(()) } +fn create_temp_file(path: &Path, mode: &WriteMode) -> Result { + #[cfg(unix)] + { + let mut options = fs::OpenOptions::new(); + options.write(true).create(true).truncate(true); + if matches!(mode, WriteMode::Private) { + options.mode(0o600); + } + options + .open(path) + .with_context(|| format!("Failed to create temp file {}", path.display())) + } + + #[cfg(not(unix))] + { + let _ = mode; + fs::File::create(path) + .with_context(|| format!("Failed to create temp file {}", path.display())) + } +} + /// Resolve a final symlink to its real target so an atomic replacement writes /// through the link (matching `fs::write`) instead of clobbering it — including a /// dangling symlink, whose not-yet-existing target is created rather than the @@ -136,8 +185,6 @@ mod tests { #[cfg(unix)] #[test] fn write_preserves_existing_file_permissions() { - use std::os::unix::fs::PermissionsExt; - let dir = tempfile::tempdir().unwrap(); let path = dir.path().join("config.toml"); write(&path, b"v1").unwrap(); @@ -149,6 +196,18 @@ mod tests { assert_eq!(mode, 0o600, "atomic replace must preserve the file mode"); } + #[cfg(unix)] + #[test] + fn write_private_creates_file_with_private_permissions() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("config.toml"); + + write_private(&path, b"secret = true").unwrap(); + + let mode = fs::metadata(&path).unwrap().permissions().mode() & 0o777; + assert_eq!(mode, 0o600, "private writes must create 0600 files"); + } + #[test] fn concurrent_writes_to_same_path_never_interleave() { // Two threads writing the same target must each use a private temp file, diff --git a/src/config.rs b/src/config.rs index d587a86..2541483 100644 --- a/src/config.rs +++ b/src/config.rs @@ -8,6 +8,7 @@ use anyhow::{Context, Result, anyhow}; use serde::{Deserialize, Serialize}; use crate::ai::providers::{AiProtocol, ProviderConfig}; +use crate::atomic; const DEFAULT_CONFIG: &str = include_str!("../config/default.toml"); @@ -267,7 +268,7 @@ fn rewrite_legacy_paths_in_config( }; if let Some(updated) = rewrite_stats_db_path_in_toml(&raw, &new_value)? { - fs::write(config_path, updated)?; + atomic::write_private(config_path, updated.as_bytes())?; } Ok(()) @@ -369,7 +370,7 @@ pub fn write_default_config_if_missing(path: &Path) -> Result { if let Some(parent) = path.parent() { fs::create_dir_all(parent)?; } - fs::write(path, DEFAULT_CONFIG)?; + atomic::write_private(path, DEFAULT_CONFIG.as_bytes())?; Ok(true) } diff --git a/src/main.rs b/src/main.rs index b988cb0..2ebca2c 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,5 +1,5 @@ use std::{ - env, fs, + env, io::{self, Write}, process::{Command, ExitCode}, time::Instant, @@ -10,7 +10,7 @@ use clap::{Args, Parser, Subcommand}; use quick_runner::{ ai, ai::providers::AiProtocol, - commands, + atomic, commands, commands::{alias::AliasCommand, config_cmd::ConfigArgs, go::GoResult}, config::{ AgentConfig, AiConfig, AppConfig, DoConfig, FallbackAiConfig, GeneralConfig, @@ -19,8 +19,6 @@ use quick_runner::{ pricing, scanner, secret, shell, stats_db::{CommandStats, StatsDb}, }; -#[cfg(unix)] -use std::os::unix::fs::PermissionsExt; #[derive(Parser)] #[command(name = "qr", version, about = "QuickRunner")] @@ -143,6 +141,10 @@ fn run() -> Result { commands::learn::print_summary(&result); Ok(ExitCode::SUCCESS) } + Commands::Init(args) => { + execute_init(args)?; + Ok(ExitCode::SUCCESS) + } command => run_with_config(command), } } @@ -217,10 +219,8 @@ fn run_with_config(command: Commands) -> Result { ); Ok(ExitCode::SUCCESS) } - Commands::Init(args) => { - execute_init(&config, args)?; - stats.command_type = "__skip_stats__".into(); - Ok(ExitCode::SUCCESS) + Commands::Init(_) => { + unreachable!("config-independent commands are dispatched before config load") } Commands::Cost { refresh } => { run_cost(&config, refresh)?; @@ -277,11 +277,22 @@ fn record_stats(config: &AppConfig, stats: &CommandStats) -> Result<()> { db.record(stats) } -fn execute_init(config: &AppConfig, args: InitArgs) -> Result<()> { +fn execute_init(args: InitArgs) -> Result<()> { let config_path = config_file_path(); - let is_new = !config_path.exists(); + let existing = if config_path.exists() { + Some(AppConfig::load_from_env_with_path(config_path.clone())) + } else { + None + }; + let needs_recreate = !config_path.exists() || existing.as_ref().is_some_and(Result::is_err); - if is_new { + let effective_config = if needs_recreate { + if let Some(Err(error)) = existing { + println!( + "config at {} is invalid; recreating it ({error:#})", + config_path.display() + ); + } // Ask for project roots interactively println!("Welcome to QuickRunner! Let's set up your project roots."); println!("Enter directories to scan for projects (one per line, empty line to finish):"); @@ -328,21 +339,14 @@ fn execute_init(config: &AppConfig, args: InitArgs) -> Result<()> { }, })?; - if let Some(parent) = config_path.parent() { - fs::create_dir_all(parent)?; - } - fs::write(&config_path, &config_content)?; - restrict_config_permissions(&config_path)?; + atomic::write_private(&config_path, config_content.as_bytes())?; println!("created {}", config_path.display()); + AppConfig::load_from_env_with_path(config_path.clone())? } else { println!("config already present at {}", config_path.display()); - } - - // Reload config in case we just wrote a new one with custom roots - let effective_config = if is_new { - AppConfig::load()? - } else { - config.clone() + existing + .transpose()? + .expect("existing config must be loaded when no recreation is needed") }; if !args.no_shell_wrapper { @@ -449,6 +453,7 @@ fn maybe_store_key_in_keychain( "Store the API key in your OS keychain (recommended; keeps it out of config.toml)?", true, )? { + println!("API key will be stored in config.toml"); return Ok(api_key); } let account = secret::account_for(api_key_env, protocol.default_api_key_env()); @@ -461,6 +466,7 @@ fn maybe_store_key_in_keychain( println!( "⚠ could not use the keychain ({error}); storing the key in config.toml instead" ); + println!("API key will be stored in config.toml"); Ok(api_key) } } @@ -468,6 +474,7 @@ fn maybe_store_key_in_keychain( fn prompt_fallback_ai_config() -> Result { let (protocol, base_url, model, api_key, api_key_env) = prompt_provider_fields("fallback")?; + let api_key = maybe_store_key_in_keychain(protocol, &api_key_env, api_key)?; Ok(FallbackAiConfig { protocol, base_url, @@ -553,16 +560,6 @@ fn prompt(label: &str) -> Result> { Ok(Some(line.trim().to_string())) } -fn restrict_config_permissions(path: &std::path::Path) -> Result<()> { - #[cfg(unix)] - { - let permissions = fs::Permissions::from_mode(0o600); - fs::set_permissions(path, permissions)?; - } - - Ok(()) -} - fn install_cron() -> Result<()> { let exe = env::current_exe().context("Could not resolve qr binary path")?; let cron_line = shell::cron_line(&exe); diff --git a/src/secret.rs b/src/secret.rs index ae5f523..9e5e151 100644 --- a/src/secret.rs +++ b/src/secret.rs @@ -6,6 +6,8 @@ //! and writes return an `Err` the caller can surface — nothing panics, and the //! rest of `qr` keeps working from environment variables or config. +use std::sync::Once; + const SERVICE: &str = "quick-runner"; /// Keychain account name for an API key: the custom env-var name when one is @@ -22,6 +24,7 @@ pub fn account_for(api_key_env: &str, default: &str) -> String { /// Fetch a stored secret by account name. Returns `None` if there is no entry or /// the keychain is unavailable for any reason. pub fn get(account: &str) -> Option { + configure_test_backend(); keyring::Entry::new(SERVICE, account) .ok()? .get_password() @@ -32,7 +35,17 @@ pub fn get(account: &str) -> Option { /// panicking) when the keychain is unavailable, so `qr init` can fall back to /// storing the key in the config file. pub fn set(account: &str, secret: &str) -> anyhow::Result<()> { + configure_test_backend(); keyring::Entry::new(SERVICE, account) .and_then(|entry| entry.set_password(secret)) .map_err(|error| anyhow::anyhow!("keychain unavailable: {error}")) } + +fn configure_test_backend() { + static ONCE: Once = Once::new(); + ONCE.call_once(|| { + if std::env::var_os("QR_TEST_USE_MOCK_KEYCHAIN").is_some() { + keyring::set_default_credential_builder(keyring::mock::default_credential_builder()); + } + }); +} diff --git a/tests/cli.rs b/tests/cli.rs index 4e41956..b012b4d 100644 --- a/tests/cli.rs +++ b/tests/cli.rs @@ -11,7 +11,7 @@ use std::{ thread, }; -use quick_runner::stats_db::StatsDb; +use quick_runner::{config::AppConfig, stats_db::StatsDb}; fn env_lock() -> &'static Mutex<()> { static LOCK: OnceLock> = OnceLock::new(); @@ -19,7 +19,11 @@ fn env_lock() -> &'static Mutex<()> { } fn clear_test_env() { - for key in ["QR_CONFIG_DIR", "QR_TEST_AI_KEY"] { + for key in [ + "QR_CONFIG_DIR", + "QR_TEST_AI_KEY", + "QR_TEST_USE_MOCK_KEYCHAIN", + ] { unsafe { std::env::remove_var(key); } @@ -260,7 +264,7 @@ db_path = "__default__" } #[test] -fn init_collects_ai_provider_settings_and_writes_secure_config() { +fn init_stores_keys_in_config_with_explicit_notice_and_private_permissions() { let _guard = env_lock().lock().unwrap(); clear_test_env(); let tmp = tempfile::tempdir().unwrap(); @@ -280,11 +284,12 @@ fn init_collects_ai_provider_settings_and_writes_secure_config() { .arg("--no-prices") .write_stdin(format!( "{root}\n\nopenai-compatible\n\nopenai-model\nconfig-primary-key\n\n\ -n\ny\nanthropic-compatible\n\nclaude-fallback\nconfig-fallback-key\nFALLBACK_SECRET\n", +n\ny\nanthropic-compatible\n\nclaude-fallback\nconfig-fallback-key\nFALLBACK_SECRET\nn\n", root = project_root.display() )) .assert() .success() + .stdout(contains("API key will be stored in config.toml")) .stdout(contains("created")) .stdout(contains("initial scan found")); @@ -311,6 +316,95 @@ n\ny\nanthropic-compatible\n\nclaude-fallback\nconfig-fallback-key\nFALLBACK_SEC } } +#[test] +fn init_recovers_from_corrupt_config_and_rewrites_valid_config() { + let _guard = env_lock().lock().unwrap(); + clear_test_env(); + let tmp = tempfile::tempdir().unwrap(); + let cfg_dir = tmp.path().join("cfg"); + let project_root = tmp.path().join("projects"); + fs::create_dir_all(&cfg_dir).unwrap(); + fs::create_dir_all(&project_root).unwrap(); + fs::write(cfg_dir.join("config.toml"), "[broken").unwrap(); + + unsafe { + std::env::set_var("QR_CONFIG_DIR", &cfg_dir); + } + + Command::cargo_bin("qr") + .unwrap() + .arg("init") + .arg("--no-shell-wrapper") + .arg("--no-cron") + .arg("--no-prices") + .write_stdin(format!( + "{root}\n\nopenai-compatible\n\nopenai-model\nrecovered-primary-key\n\nn\nn\n", + root = project_root.display() + )) + .assert() + .success() + .stdout(contains("created")) + .stdout(contains("initial scan found")); + + let config_path = cfg_dir.join("config.toml"); + let config = AppConfig::load_from_env_with_path(config_path.clone()).unwrap(); + assert_eq!( + config.projects.roots, + vec![project_root.display().to_string()] + ); + assert_eq!(config.ai.model, "openai-model"); + assert_eq!(config.ai.api_key, "recovered-primary-key"); + assert!(config.ai.fallback.is_none()); + assert!(!fs::read_to_string(config_path).unwrap().contains("[broken")); + + unsafe { + std::env::remove_var("QR_CONFIG_DIR"); + } +} + +#[test] +fn init_stores_fallback_key_in_keychain_when_requested() { + let _guard = env_lock().lock().unwrap(); + clear_test_env(); + let tmp = tempfile::tempdir().unwrap(); + let cfg_dir = tmp.path().join("cfg"); + let project_root = tmp.path().join("projects"); + fs::create_dir_all(&cfg_dir).unwrap(); + fs::create_dir_all(&project_root).unwrap(); + + unsafe { + std::env::set_var("QR_CONFIG_DIR", &cfg_dir); + std::env::set_var("QR_TEST_USE_MOCK_KEYCHAIN", "1"); + } + + Command::cargo_bin("qr") + .unwrap() + .arg("init") + .arg("--no-shell-wrapper") + .arg("--no-cron") + .arg("--no-prices") + .write_stdin(format!( + "{root}\n\nopenai-compatible\n\nopenai-model\nprimary-secret\n\n\ +y\ny\nanthropic-compatible\n\nclaude-fallback\nfallback-secret\nFALLBACK_SECRET\ny\n", + root = project_root.display() + )) + .assert() + .success() + .stdout(contains("stored API key in the OS keychain")); + + let raw = fs::read_to_string(cfg_dir.join("config.toml")).unwrap(); + assert_eq!(raw.matches("api_key = \"\"").count(), 2); + assert!(!raw.contains("primary-secret")); + assert!(!raw.contains("fallback-secret")); + assert!(raw.contains("api_key_env = \"OPENAI_API_KEY\"")); + assert!(raw.contains("api_key_env = \"FALLBACK_SECRET\"")); + + unsafe { + std::env::remove_var("QR_CONFIG_DIR"); + std::env::remove_var("QR_TEST_USE_MOCK_KEYCHAIN"); + } +} + #[test] fn init_fails_fast_on_closed_stdin_instead_of_hanging() { let _guard = env_lock().lock().unwrap(); From 3baf7d7cc9d1c61edffd534d339ba8de20a233d3 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 1 Jul 2026 10:37:19 +0000 Subject: [PATCH 2/2] Address Codex review: env-safe init check and distinct keychain roles - Parse config.toml without QR_* overrides when deciding if init should recreate - Fail clearly on invalid env overrides instead of treating them as corrupt config - Prefix keychain accounts with primary/fallback role; legacy lookup for primary - Add tests for invalid-env init behavior and shared-default-env keychain storage Co-authored-by: Aanish Bhirud --- src/ai/client.rs | 65 +++++++++++++++++++++++++++------------------- src/config.rs | 20 +++++++++------ src/main.rs | 13 +++++----- src/secret.rs | 49 ++++++++++++++++++++++++++++++++--- tests/cli.rs | 67 +++++++++++++++++++++++++++++++++++++++++++++--- 5 files changed, 165 insertions(+), 49 deletions(-) diff --git a/src/ai/client.rs b/src/ai/client.rs index 6e34c6c..e967deb 100644 --- a/src/ai/client.rs +++ b/src/ai/client.rs @@ -57,11 +57,11 @@ impl AiClient { } pub fn execute_prompt(&self, system_prompt: &str, user_message: &str) -> Result { - match self.send_prompt(&self.primary, system_prompt, user_message) { + match self.send_prompt(&self.primary, "primary", system_prompt, user_message) { Ok(response) => Ok(response), Err(primary_error) => { if let Some(fallback) = &self.fallback { - self.send_prompt(fallback, system_prompt, user_message) + self.send_prompt(fallback, "fallback", system_prompt, user_message) .map_err(|fallback_error| { anyhow!( "Primary provider failed: {primary_error}. Fallback provider failed: {fallback_error}" @@ -77,10 +77,11 @@ impl AiClient { fn send_prompt( &self, provider: &ProviderConfig, + keychain_role: &str, system_prompt: &str, user_message: &str, ) -> Result { - let api_key = resolve_api_key(provider)?; + let api_key = resolve_api_key(provider, keychain_role)?; let url = endpoint_url(provider); let request = match provider.protocol { @@ -126,7 +127,7 @@ impl AiClient { } } -fn resolve_api_key(provider: &ProviderConfig) -> Result { +fn resolve_api_key(provider: &ProviderConfig, keychain_role: &str) -> Result { // 1. explicit per-provider env var (highest-precedence override) if !provider.api_key_env.trim().is_empty() { if let Ok(value) = env::var(&provider.api_key_env) { @@ -152,8 +153,9 @@ fn resolve_api_key(provider: &ProviderConfig) -> Result { // 4. OS keychain — where `qr init` stores the key when you opt in. Checked // last because it is only populated when the key is NOT in env or config, // so the common paths never touch the keychain backend. - let account = crate::secret::account_for(&provider.api_key_env, well_known_env); - if let Some(value) = crate::secret::get(&account) { + if let Some(value) = + crate::secret::get_for_role(keychain_role, &provider.api_key_env, well_known_env) + { if !value.trim().is_empty() { return Ok(value); } @@ -437,13 +439,16 @@ mod tests { std::env::set_var("OPENAI_API_KEY", "well-known-token"); } - let api_key = resolve_api_key(&ProviderConfig { - protocol: AiProtocol::OpenAi, - base_url: "https://example.test/v1".into(), - model: "demo".into(), - api_key: "config-token".into(), - api_key_env: "CUSTOM_QR_TEST_AI_KEY".into(), - }) + let api_key = resolve_api_key( + &ProviderConfig { + protocol: AiProtocol::OpenAi, + base_url: "https://example.test/v1".into(), + model: "demo".into(), + api_key: "config-token".into(), + api_key_env: "CUSTOM_QR_TEST_AI_KEY".into(), + }, + "primary", + ) .unwrap(); assert_eq!(api_key, "custom-token"); @@ -462,13 +467,16 @@ mod tests { std::env::set_var("ANTHROPIC_API_KEY", "well-known-token"); } - let api_key = resolve_api_key(&ProviderConfig { - protocol: AiProtocol::Anthropic, - base_url: "https://example.test".into(), - model: "claude-demo".into(), - api_key: "config-token".into(), - api_key_env: "CUSTOM_QR_TEST_ANTHROPIC_KEY".into(), - }) + let api_key = resolve_api_key( + &ProviderConfig { + protocol: AiProtocol::Anthropic, + base_url: "https://example.test".into(), + model: "claude-demo".into(), + api_key: "config-token".into(), + api_key_env: "CUSTOM_QR_TEST_ANTHROPIC_KEY".into(), + }, + "primary", + ) .unwrap(); assert_eq!(api_key, "well-known-token"); @@ -486,13 +494,16 @@ mod tests { std::env::remove_var("OPENAI_API_KEY"); } - let api_key = resolve_api_key(&ProviderConfig { - protocol: AiProtocol::OpenAi, - base_url: "https://example.test/v1".into(), - model: "demo".into(), - api_key: "config-token".into(), - api_key_env: "CUSTOM_QR_TEST_AI_KEY".into(), - }) + let api_key = resolve_api_key( + &ProviderConfig { + protocol: AiProtocol::OpenAi, + base_url: "https://example.test/v1".into(), + model: "demo".into(), + api_key: "config-token".into(), + api_key_env: "CUSTOM_QR_TEST_AI_KEY".into(), + }, + "primary", + ) .unwrap(); assert_eq!(api_key, "config-token"); } diff --git a/src/config.rs b/src/config.rs index 2541483..d879a91 100644 --- a/src/config.rs +++ b/src/config.rs @@ -90,18 +90,22 @@ impl AppConfig { } pub fn load_from_env_with_path(path: PathBuf) -> Result { - let mut config = if path.exists() { - let raw = fs::read_to_string(&path) - .with_context(|| format!("Failed to read config file {}", path.display()))?; - Self::load_from_str(&raw)? - } else { - Self::load_from_str(DEFAULT_CONFIG)? - }; - + let mut config = Self::load_file_without_env(&path)?; apply_env_overrides(&mut config)?; Ok(config) } + /// Parse `config.toml` on disk without applying `QR_*` env overrides. + pub fn load_file_without_env(path: &Path) -> Result { + if path.exists() { + let raw = fs::read_to_string(path) + .with_context(|| format!("Failed to read config file {}", path.display()))?; + Self::load_from_str(&raw) + } else { + Self::load_from_str(DEFAULT_CONFIG) + } + } + pub fn ensure_parent_dirs(&self) -> Result<()> { for path in [self.stats_db_path(), cache_file_path(), config_file_path()] { if let Some(parent) = path.parent() { diff --git a/src/main.rs b/src/main.rs index 2ebca2c..66009d3 100644 --- a/src/main.rs +++ b/src/main.rs @@ -280,7 +280,7 @@ fn record_stats(config: &AppConfig, stats: &CommandStats) -> Result<()> { fn execute_init(args: InitArgs) -> Result<()> { let config_path = config_file_path(); let existing = if config_path.exists() { - Some(AppConfig::load_from_env_with_path(config_path.clone())) + Some(AppConfig::load_file_without_env(&config_path)) } else { None }; @@ -344,9 +344,7 @@ fn execute_init(args: InitArgs) -> Result<()> { AppConfig::load_from_env_with_path(config_path.clone())? } else { println!("config already present at {}", config_path.display()); - existing - .transpose()? - .expect("existing config must be loaded when no recreation is needed") + AppConfig::load_from_env_with_path(config_path.clone())? }; if !args.no_shell_wrapper { @@ -391,7 +389,7 @@ fn prompt_provider_fields(label: &str) -> Result<(AiProtocol, String, String, St fn prompt_ai_config(label: &str) -> Result { let (protocol, base_url, model, api_key, api_key_env) = prompt_provider_fields(label)?; - let api_key = maybe_store_key_in_keychain(protocol, &api_key_env, api_key)?; + let api_key = maybe_store_key_in_keychain("primary", protocol, &api_key_env, api_key)?; let fallback = if prompt_bool("Configure a fallback provider?", false)? { Some(prompt_fallback_ai_config()?) @@ -445,6 +443,7 @@ fn run_cost(config: &AppConfig, refresh: bool) -> Result<()> { /// keychain, otherwise the key itself (config storage is the fallback when the /// user declines or the keychain is unavailable). fn maybe_store_key_in_keychain( + role: &str, protocol: AiProtocol, api_key_env: &str, api_key: String, @@ -456,7 +455,7 @@ fn maybe_store_key_in_keychain( println!("API key will be stored in config.toml"); return Ok(api_key); } - let account = secret::account_for(api_key_env, protocol.default_api_key_env()); + let account = secret::account_for(role, api_key_env, protocol.default_api_key_env()); match secret::set(&account, &api_key) { Ok(()) => { println!("✔ stored API key in the OS keychain (account \"{account}\")"); @@ -474,7 +473,7 @@ fn maybe_store_key_in_keychain( fn prompt_fallback_ai_config() -> Result { let (protocol, base_url, model, api_key, api_key_env) = prompt_provider_fields("fallback")?; - let api_key = maybe_store_key_in_keychain(protocol, &api_key_env, api_key)?; + let api_key = maybe_store_key_in_keychain("fallback", protocol, &api_key_env, api_key)?; Ok(FallbackAiConfig { protocol, base_url, diff --git a/src/secret.rs b/src/secret.rs index 9e5e151..6077a9c 100644 --- a/src/secret.rs +++ b/src/secret.rs @@ -10,10 +10,30 @@ use std::sync::Once; const SERVICE: &str = "quick-runner"; -/// Keychain account name for an API key: the custom env-var name when one is -/// configured, otherwise `default` (the protocol's well-known env var). Keeps -/// `qr init`'s store and the AI client's lookup in agreement. -pub fn account_for(api_key_env: &str, default: &str) -> String { +/// Keychain account name for an API key. `role` distinguishes primary vs fallback +/// when both use the same env-var name (e.g. two openai-compatible endpoints). +pub fn account_for(role: &str, api_key_env: &str, default: &str) -> String { + let base = if api_key_env.trim().is_empty() { + default.to_string() + } else { + api_key_env.to_string() + }; + format!("{role}:{base}") +} + +/// Look up a role-scoped keychain entry, falling back to the legacy bare account +/// for primary keys stored before role-prefixing shipped. +pub fn get_for_role(role: &str, api_key_env: &str, default: &str) -> Option { + get(&account_for(role, api_key_env, default)).or_else(|| { + if role == "primary" { + get(&legacy_account(api_key_env, default)) + } else { + None + } + }) +} + +fn legacy_account(api_key_env: &str, default: &str) -> String { if api_key_env.trim().is_empty() { default.to_string() } else { @@ -49,3 +69,24 @@ fn configure_test_backend() { } }); } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn primary_and_fallback_accounts_differ_when_env_name_matches() { + assert_eq!( + account_for("primary", "", "OPENAI_API_KEY"), + "primary:OPENAI_API_KEY" + ); + assert_eq!( + account_for("fallback", "", "OPENAI_API_KEY"), + "fallback:OPENAI_API_KEY" + ); + assert_ne!( + account_for("primary", "", "OPENAI_API_KEY"), + account_for("fallback", "", "OPENAI_API_KEY") + ); + } +} diff --git a/tests/cli.rs b/tests/cli.rs index b012b4d..ef0eba6 100644 --- a/tests/cli.rs +++ b/tests/cli.rs @@ -23,6 +23,7 @@ fn clear_test_env() { "QR_CONFIG_DIR", "QR_TEST_AI_KEY", "QR_TEST_USE_MOCK_KEYCHAIN", + "QR_SCAN_DEPTH", ] { unsafe { std::env::remove_var(key); @@ -385,19 +386,24 @@ fn init_stores_fallback_key_in_keychain_when_requested() { .arg("--no-prices") .write_stdin(format!( "{root}\n\nopenai-compatible\n\nopenai-model\nprimary-secret\n\n\ -y\ny\nanthropic-compatible\n\nclaude-fallback\nfallback-secret\nFALLBACK_SECRET\ny\n", +y\ny\nopenai-compatible\n\nfallback-model\nfallback-secret\n\ny\n", root = project_root.display() )) .assert() .success() - .stdout(contains("stored API key in the OS keychain")); + .stdout(contains( + "stored API key in the OS keychain (account \"primary:OPENAI_API_KEY\")", + )) + .stdout(contains( + "stored API key in the OS keychain (account \"fallback:OPENAI_API_KEY\")", + )); let raw = fs::read_to_string(cfg_dir.join("config.toml")).unwrap(); assert_eq!(raw.matches("api_key = \"\"").count(), 2); assert!(!raw.contains("primary-secret")); assert!(!raw.contains("fallback-secret")); assert!(raw.contains("api_key_env = \"OPENAI_API_KEY\"")); - assert!(raw.contains("api_key_env = \"FALLBACK_SECRET\"")); + assert!(raw.contains("model = \"fallback-model\"")); unsafe { std::env::remove_var("QR_CONFIG_DIR"); @@ -405,6 +411,61 @@ y\ny\nanthropic-compatible\n\nclaude-fallback\nfallback-secret\nFALLBACK_SECRET\ } } +#[test] +fn init_keeps_valid_config_when_env_override_is_invalid() { + let _guard = env_lock().lock().unwrap(); + clear_test_env(); + let tmp = tempfile::tempdir().unwrap(); + let cfg_dir = tmp.path().join("cfg"); + fs::create_dir_all(&cfg_dir).unwrap(); + fs::write( + cfg_dir.join("config.toml"), + r#"[general] +default_run_mode = "output" + +[projects] +roots = ["~/Development"] +scan_depth = 2 +scan_interval_hours = 1 + +[ai] +protocol = "openai" +base_url = "https://api.openai.com/v1" +model = "gpt-4o" +api_key = "keep-me" +api_key_env = "OPENAI_API_KEY" + +[stats] +enabled = false +db_path = "__default__" +"#, + ) + .unwrap(); + + unsafe { + std::env::set_var("QR_CONFIG_DIR", &cfg_dir); + std::env::set_var("QR_SCAN_DEPTH", "not-a-number"); + } + + Command::cargo_bin("qr") + .unwrap() + .args(["init", "--no-shell-wrapper", "--no-cron", "--no-prices"]) + .write_stdin("\n") + .assert() + .failure() + .stdout(contains("config already present")) + .stdout(contains("recreating").not()) + .stderr(contains("QR_SCAN_DEPTH")); + + let raw = fs::read_to_string(cfg_dir.join("config.toml")).unwrap(); + assert!(raw.contains("api_key = \"keep-me\"")); + + unsafe { + std::env::remove_var("QR_CONFIG_DIR"); + std::env::remove_var("QR_SCAN_DEPTH"); + } +} + #[test] fn init_fails_fast_on_closed_stdin_instead_of_hanging() { let _guard = env_lock().lock().unwrap();