Skip to main content

omni_dev/utils/
preflight.rs

1//! Preflight validation checks for early failure detection.
2//!
3//! This module provides functions to validate required services and credentials
4//! before starting expensive operations. Commands should call these checks early
5//! to fail fast with clear error messages.
6
7use anyhow::{bail, Context, Result};
8
9use crate::claude::model_config::get_model_registry;
10
11/// Result of AI credential validation.
12#[derive(Debug)]
13pub struct AiCredentialInfo {
14    /// The AI provider that will be used.
15    pub provider: AiProvider,
16    /// The model that will be used.
17    pub model: String,
18}
19
20/// AI provider types.
21#[derive(Debug, Clone, Copy, PartialEq, Eq)]
22pub enum AiProvider {
23    /// Anthropic Claude API.
24    Claude,
25    /// AWS Bedrock with Claude.
26    Bedrock,
27    /// OpenAI API.
28    OpenAi,
29    /// Local Ollama.
30    Ollama,
31    /// `claude -p` subprocess (Claude Code CLI).
32    ClaudeCli,
33}
34
35impl std::fmt::Display for AiProvider {
36    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
37        match self {
38            Self::Claude => write!(f, "Claude API"),
39            Self::Bedrock => write!(f, "AWS Bedrock"),
40            Self::OpenAi => write!(f, "OpenAI API"),
41            Self::Ollama => write!(f, "Ollama"),
42            Self::ClaudeCli => write!(f, "Claude Code CLI"),
43        }
44    }
45}
46
47/// Validates that AI credentials are available before processing.
48///
49/// This performs a lightweight check of environment variables without
50/// creating a full AI client. Use this at the start of commands that
51/// require AI to fail fast if credentials are missing.
52pub fn check_ai_credentials(model_override: Option<&str>) -> Result<AiCredentialInfo> {
53    use crate::utils::settings::{get_env_var, get_env_vars};
54
55    // The `claude -p` subprocess backend is checked first so it wins over
56    // the existing USE_* flags if multiple are set. Credentials for this
57    // backend live inside the `claude` binary's own auth state, so we just
58    // verify the binary is on PATH.
59    if let Ok(val) = get_env_var("OMNI_DEV_AI_BACKEND") {
60        if matches!(val.as_str(), "claude-cli" | "claude_cli") {
61            let binary =
62                get_env_var("OMNI_DEV_CLAUDE_CLI_BIN").unwrap_or_else(|_| "claude".to_string());
63            let probe = std::process::Command::new(&binary)
64                .arg("--version")
65                .output();
66            match probe {
67                Ok(out) if out.status.success() => {
68                    let registry = get_model_registry();
69                    let model = model_override
70                        .map(String::from)
71                        .or_else(|| get_env_var("CLAUDE_MODEL").ok())
72                        .or_else(|| get_env_var("CLAUDE_CODE_MODEL").ok())
73                        .or_else(|| get_env_var("ANTHROPIC_MODEL").ok())
74                        .unwrap_or_else(|| {
75                            registry
76                                .get_default_model("claude")
77                                .unwrap_or("claude-sonnet-4-6")
78                                .to_string()
79                        });
80                    return Ok(AiCredentialInfo {
81                        provider: AiProvider::ClaudeCli,
82                        model,
83                    });
84                }
85                _ => bail!(
86                    "Claude Code CLI not available at '{binary}'.\n\
87                     Install it from https://github.com/anthropics/claude-code \
88                     or set OMNI_DEV_CLAUDE_CLI_BIN to its path."
89                ),
90            }
91        }
92    }
93
94    // Check provider selection flags
95    let use_openai = get_env_var("USE_OPENAI").is_ok_and(|val| val == "true");
96
97    let use_ollama = get_env_var("USE_OLLAMA").is_ok_and(|val| val == "true");
98
99    let use_bedrock = get_env_var("CLAUDE_CODE_USE_BEDROCK").is_ok_and(|val| val == "true");
100
101    // Check Ollama (no credentials required, just model)
102    if use_ollama {
103        let model = model_override
104            .map(String::from)
105            .or_else(|| get_env_var("OLLAMA_MODEL").ok())
106            .unwrap_or_else(|| "llama2".to_string());
107
108        return Ok(AiCredentialInfo {
109            provider: AiProvider::Ollama,
110            model,
111        });
112    }
113
114    // Check OpenAI
115    if use_openai {
116        let registry = get_model_registry();
117        let model = model_override
118            .map(String::from)
119            .or_else(|| get_env_var("OPENAI_MODEL").ok())
120            .unwrap_or_else(|| {
121                registry
122                    .get_default_model("openai")
123                    .unwrap_or("gpt-5")
124                    .to_string()
125            });
126
127        // Verify API key exists
128        get_env_vars(&["OPENAI_API_KEY", "OPENAI_AUTH_TOKEN"]).map_err(|_| {
129            anyhow::anyhow!(
130                "OpenAI API key not found.\n\
131                 Set one of these environment variables:\n\
132                 - OPENAI_API_KEY\n\
133                 - OPENAI_AUTH_TOKEN"
134            )
135        })?;
136
137        return Ok(AiCredentialInfo {
138            provider: AiProvider::OpenAi,
139            model,
140        });
141    }
142
143    // Check Bedrock
144    if use_bedrock {
145        let registry = get_model_registry();
146        let model = model_override
147            .map(String::from)
148            .or_else(|| get_env_var("ANTHROPIC_MODEL").ok())
149            .unwrap_or_else(|| {
150                registry
151                    .get_default_model("claude")
152                    .unwrap_or("claude-sonnet-4-6")
153                    .to_string()
154            });
155
156        // Verify Bedrock configuration
157        get_env_var("ANTHROPIC_AUTH_TOKEN").map_err(|_| {
158            anyhow::anyhow!(
159                "AWS Bedrock authentication not configured.\n\
160                 Set ANTHROPIC_AUTH_TOKEN environment variable."
161            )
162        })?;
163
164        get_env_var("ANTHROPIC_BEDROCK_BASE_URL").map_err(|_| {
165            anyhow::anyhow!(
166                "AWS Bedrock base URL not configured.\n\
167                 Set ANTHROPIC_BEDROCK_BASE_URL environment variable."
168            )
169        })?;
170
171        return Ok(AiCredentialInfo {
172            provider: AiProvider::Bedrock,
173            model,
174        });
175    }
176
177    // Default: Claude API
178    let registry = get_model_registry();
179    let model = model_override
180        .map(String::from)
181        .or_else(|| get_env_var("ANTHROPIC_MODEL").ok())
182        .unwrap_or_else(|| {
183            registry
184                .get_default_model("claude")
185                .unwrap_or("claude-sonnet-4-6")
186                .to_string()
187        });
188
189    // Verify API key exists
190    get_env_vars(&[
191        "CLAUDE_API_KEY",
192        "ANTHROPIC_API_KEY",
193        "ANTHROPIC_AUTH_TOKEN",
194    ])
195    .map_err(|_| {
196        anyhow::anyhow!(
197            "Claude API key not found.\n\
198                 Set one of these environment variables:\n\
199                 - CLAUDE_API_KEY\n\
200                 - ANTHROPIC_API_KEY\n\
201                 - ANTHROPIC_AUTH_TOKEN"
202        )
203    })?;
204
205    Ok(AiCredentialInfo {
206        provider: AiProvider::Claude,
207        model,
208    })
209}
210
211/// Validates that GitHub CLI is available and authenticated.
212///
213/// This checks:
214/// 1. `gh` CLI is installed and in PATH
215/// 2. User is authenticated (can access the current repo)
216///
217/// Use this at the start of commands that require GitHub API access.
218///
219/// `repo_root` anchors the repository-access probe to the injected repository
220/// rather than the process current working directory.
221pub fn check_github_cli(repo_root: &std::path::Path) -> Result<()> {
222    // Check if gh CLI is available. This probe is a PATH availability check
223    // (CWD-independent), so it is not anchored to `repo_root`.
224    let gh_check = std::process::Command::new("gh")
225        .args(["--version"])
226        .output();
227
228    match gh_check {
229        Ok(output) if output.status.success() => {
230            // Test if gh can access the injected repo
231            let repo_check = std::process::Command::new("gh")
232                .args(["repo", "view", "--json", "name"])
233                .current_dir(repo_root)
234                .output();
235
236            match repo_check {
237                Ok(repo_output) if repo_output.status.success() => Ok(()),
238                Ok(repo_output) => {
239                    let error_details = String::from_utf8_lossy(&repo_output.stderr);
240                    if error_details.contains("authentication") || error_details.contains("login") {
241                        bail!(
242                            "GitHub CLI authentication failed.\n\
243                             Please run 'gh auth login' or set GITHUB_TOKEN environment variable."
244                        )
245                    }
246                    bail!(
247                        "GitHub CLI cannot access this repository.\n\
248                         Error: {}",
249                        error_details.trim()
250                    )
251                }
252                Err(e) => bail!("Failed to test GitHub CLI access: {e}"),
253            }
254        }
255        _ => bail!(
256            "GitHub CLI (gh) is not installed or not in PATH.\n\
257             Please install it from https://cli.github.com/"
258        ),
259    }
260}
261
262/// Validates that `repo_root` is a valid git repository.
263///
264/// A lightweight check that opens the repository without loading commit data.
265pub fn check_git_repository_at(repo_root: &std::path::Path) -> Result<()> {
266    crate::git::GitRepository::open_at(repo_root).context(
267        "Not in a git repository. Please run this command from within a git repository.",
268    )?;
269    Ok(())
270}
271
272/// Validates that the working directory at `repo_root` is clean — no
273/// uncommitted changes (staged, unstaged, or untracked non-ignored files).
274///
275/// Use this before operations that require a clean working directory, like
276/// amending commits.
277pub fn check_working_directory_clean_at(repo_root: &std::path::Path) -> Result<()> {
278    let repo =
279        crate::git::GitRepository::open_at(repo_root).context("Failed to open git repository")?;
280    check_working_directory_clean_for(&repo)
281}
282
283/// Shared clean-worktree check over an already-opened repository.
284fn check_working_directory_clean_for(repo: &crate::git::GitRepository) -> Result<()> {
285    let status = repo
286        .get_working_directory_status()
287        .context("Failed to get working directory status")?;
288
289    if !status.clean {
290        let mut message = String::from("Working directory has uncommitted changes:\n");
291        for change in &status.untracked_changes {
292            message.push_str(&format!("  {} {}\n", change.status, change.file));
293        }
294        message.push_str("\nPlease commit or stash your changes before proceeding.");
295        bail!(message);
296    }
297
298    Ok(())
299}
300
301/// Performs combined preflight check for AI commands.
302///
303/// Validates:
304/// - Git repository access
305/// - AI credentials
306///
307/// Returns information about the AI provider that will be used.
308///
309/// `repo_root` anchors the git-repository check to the injected repository
310/// rather than the process current working directory.
311pub fn check_ai_command_prerequisites(
312    model_override: Option<&str>,
313    repo_root: &std::path::Path,
314) -> Result<AiCredentialInfo> {
315    check_git_repository_at(repo_root)?;
316    check_ai_credentials(model_override)
317}
318
319/// Performs combined preflight check for PR creation.
320///
321/// Validates:
322/// - Git repository access
323/// - AI credentials
324/// - GitHub CLI availability and authentication
325///
326/// Returns information about the AI provider that will be used.
327///
328/// `repo_root` anchors the git-repository and GitHub CLI checks to the injected
329/// repository rather than the process current working directory.
330pub fn check_pr_command_prerequisites(
331    model_override: Option<&str>,
332    repo_root: &std::path::Path,
333) -> Result<AiCredentialInfo> {
334    check_git_repository_at(repo_root)?;
335    let ai_info = check_ai_credentials(model_override)?;
336    check_github_cli(repo_root)?;
337    Ok(ai_info)
338}
339
340#[cfg(test)]
341#[allow(clippy::unwrap_used, clippy::expect_used)]
342mod tests {
343    use super::*;
344
345    use std::env;
346    use std::sync::Mutex;
347    use std::sync::OnceLock;
348
349    /// Global lock to ensure environment variable tests don't interfere with each other.
350    static ENV_TEST_LOCK: OnceLock<Mutex<()>> = OnceLock::new();
351
352    /// Manages environment variables in tests to avoid interference.
353    struct EnvGuard {
354        _lock: std::sync::MutexGuard<'static, ()>,
355        vars: Vec<(String, Option<String>)>,
356    }
357
358    impl EnvGuard {
359        fn new() -> Self {
360            let lock = ENV_TEST_LOCK.get_or_init(|| Mutex::new(())).lock().unwrap();
361            Self {
362                _lock: lock,
363                vars: Vec::new(),
364            }
365        }
366
367        fn set(&mut self, key: &str, value: &str) {
368            let original = env::var(key).ok();
369            self.vars.push((key.to_string(), original));
370            env::set_var(key, value);
371        }
372
373        fn remove(&mut self, key: &str) {
374            let original = env::var(key).ok();
375            self.vars.push((key.to_string(), original));
376            env::remove_var(key);
377        }
378    }
379
380    impl Drop for EnvGuard {
381        fn drop(&mut self) {
382            for (key, original_value) in self.vars.drain(..).rev() {
383                match original_value {
384                    Some(value) => env::set_var(&key, value),
385                    None => env::remove_var(&key),
386                }
387            }
388        }
389    }
390
391    #[test]
392    fn ai_provider_display() {
393        assert_eq!(format!("{}", AiProvider::Claude), "Claude API");
394        assert_eq!(format!("{}", AiProvider::Bedrock), "AWS Bedrock");
395        assert_eq!(format!("{}", AiProvider::OpenAi), "OpenAI API");
396        assert_eq!(format!("{}", AiProvider::Ollama), "Ollama");
397        assert_eq!(format!("{}", AiProvider::ClaudeCli), "Claude Code CLI");
398    }
399
400    #[test]
401    fn ai_provider_equality() {
402        assert_eq!(AiProvider::Claude, AiProvider::Claude);
403        assert_ne!(AiProvider::Claude, AiProvider::OpenAi);
404        assert_ne!(AiProvider::Bedrock, AiProvider::Ollama);
405    }
406
407    #[test]
408    fn ai_provider_clone() {
409        let provider = AiProvider::Bedrock;
410        let cloned = provider;
411        assert_eq!(provider, cloned);
412    }
413
414    #[test]
415    fn ai_provider_debug() {
416        let debug_str = format!("{:?}", AiProvider::Claude);
417        assert_eq!(debug_str, "Claude");
418    }
419
420    #[test]
421    fn ai_credential_info_debug() {
422        let info = AiCredentialInfo {
423            provider: AiProvider::Ollama,
424            model: "llama2".to_string(),
425        };
426        let debug_str = format!("{info:?}");
427        assert!(debug_str.contains("Ollama"));
428        assert!(debug_str.contains("llama2"));
429    }
430
431    #[test]
432    fn claude_default_model_from_registry() {
433        let mut guard = EnvGuard::new();
434        // Enable Claude API path with a dummy key, no model override
435        guard.remove("USE_OPENAI");
436        guard.remove("USE_OLLAMA");
437        guard.remove("CLAUDE_CODE_USE_BEDROCK");
438        guard.remove("ANTHROPIC_MODEL");
439        guard.set("ANTHROPIC_API_KEY", "sk-test-dummy");
440
441        let info = check_ai_credentials(None).unwrap();
442        assert_eq!(info.provider, AiProvider::Claude);
443        assert_eq!(info.model, "claude-sonnet-4-6");
444    }
445
446    #[test]
447    fn openai_default_model_from_registry() {
448        let mut guard = EnvGuard::new();
449        guard.set("USE_OPENAI", "true");
450        guard.remove("USE_OLLAMA");
451        guard.remove("OPENAI_MODEL");
452        guard.set("OPENAI_API_KEY", "sk-test-dummy");
453
454        let info = check_ai_credentials(None).unwrap();
455        assert_eq!(info.provider, AiProvider::OpenAi);
456        assert_eq!(info.model, "gpt-5-mini");
457    }
458
459    #[test]
460    fn bedrock_default_model_from_registry() {
461        let mut guard = EnvGuard::new();
462        guard.remove("USE_OPENAI");
463        guard.remove("USE_OLLAMA");
464        guard.set("CLAUDE_CODE_USE_BEDROCK", "true");
465        guard.remove("ANTHROPIC_MODEL");
466        guard.set("ANTHROPIC_AUTH_TOKEN", "test-token");
467        guard.set("ANTHROPIC_BEDROCK_BASE_URL", "https://bedrock.example.com");
468
469        let info = check_ai_credentials(None).unwrap();
470        assert_eq!(info.provider, AiProvider::Bedrock);
471        assert_eq!(info.model, "claude-sonnet-4-6");
472    }
473
474    #[test]
475    fn model_override_takes_precedence() {
476        let mut guard = EnvGuard::new();
477        guard.remove("USE_OPENAI");
478        guard.remove("USE_OLLAMA");
479        guard.remove("CLAUDE_CODE_USE_BEDROCK");
480        guard.remove("ANTHROPIC_MODEL");
481        guard.set("ANTHROPIC_API_KEY", "sk-test-dummy");
482
483        let info = check_ai_credentials(Some("claude-opus-4-6")).unwrap();
484        assert_eq!(info.model, "claude-opus-4-6");
485    }
486
487    #[cfg(unix)]
488    fn make_version_shim(tmp: &tempfile::TempDir, exit_code: i32) -> std::path::PathBuf {
489        let shim = tmp.path().join("claude-bin-shim");
490        crate::test_support::shim::write_exec_script(
491            &shim,
492            &format!("#!/bin/sh\necho 'fake-claude 0.0.0'\nexit {exit_code}\n"),
493        );
494        shim
495    }
496
497    #[test]
498    #[cfg(unix)]
499    fn claude_cli_backend_uses_version_probe() {
500        let _guard = crate::test_support::shim::shim_lock();
501        let tmp = tempfile::TempDir::new().unwrap();
502        let shim = make_version_shim(&tmp, 0);
503
504        let mut guard = EnvGuard::new();
505        guard.remove("USE_OPENAI");
506        guard.remove("USE_OLLAMA");
507        guard.remove("CLAUDE_CODE_USE_BEDROCK");
508        guard.remove("ANTHROPIC_MODEL");
509        guard.remove("CLAUDE_MODEL");
510        guard.remove("CLAUDE_CODE_MODEL");
511        guard.set("OMNI_DEV_AI_BACKEND", "claude-cli");
512        guard.set("OMNI_DEV_CLAUDE_CLI_BIN", shim.to_str().unwrap());
513
514        let info = check_ai_credentials(None).unwrap();
515        assert_eq!(info.provider, AiProvider::ClaudeCli);
516        assert_eq!(info.model, "claude-sonnet-4-6");
517    }
518
519    #[test]
520    #[cfg(unix)]
521    fn claude_cli_backend_uses_model_from_env() {
522        let _guard = crate::test_support::shim::shim_lock();
523        let tmp = tempfile::TempDir::new().unwrap();
524        let shim = make_version_shim(&tmp, 0);
525
526        let mut guard = EnvGuard::new();
527        guard.remove("USE_OPENAI");
528        guard.remove("USE_OLLAMA");
529        guard.remove("CLAUDE_CODE_USE_BEDROCK");
530        guard.remove("ANTHROPIC_MODEL");
531        guard.remove("CLAUDE_CODE_MODEL");
532        guard.set("OMNI_DEV_AI_BACKEND", "claude-cli");
533        guard.set("OMNI_DEV_CLAUDE_CLI_BIN", shim.to_str().unwrap());
534        guard.set("CLAUDE_MODEL", "haiku");
535
536        let info = check_ai_credentials(None).unwrap();
537        assert_eq!(info.provider, AiProvider::ClaudeCli);
538        assert_eq!(info.model, "haiku");
539    }
540
541    #[test]
542    fn claude_cli_backend_missing_binary_fails_preflight() {
543        let mut guard = EnvGuard::new();
544        guard.remove("USE_OPENAI");
545        guard.remove("USE_OLLAMA");
546        guard.remove("CLAUDE_CODE_USE_BEDROCK");
547        guard.set("OMNI_DEV_AI_BACKEND", "claude-cli");
548        guard.set("OMNI_DEV_CLAUDE_CLI_BIN", "/nonexistent/claude-binary-xyz");
549
550        let err = check_ai_credentials(None).expect_err("expected missing-binary error");
551        let chain = format!("{err:#}");
552        assert!(
553            chain.contains("Claude Code CLI not available"),
554            "unexpected error: {chain}"
555        );
556    }
557
558    #[test]
559    fn claude_cli_backend_accepts_underscore_alias() {
560        // The factory/preflight accept both `claude-cli` and `claude_cli`.
561        // Verify the second spelling routes the same way (missing-binary
562        // path exercises the selector cheaply).
563        let mut guard = EnvGuard::new();
564        guard.remove("USE_OPENAI");
565        guard.remove("USE_OLLAMA");
566        guard.remove("CLAUDE_CODE_USE_BEDROCK");
567        guard.set("OMNI_DEV_AI_BACKEND", "claude_cli");
568        guard.set("OMNI_DEV_CLAUDE_CLI_BIN", "/nonexistent/claude-binary-xyz");
569
570        let err = check_ai_credentials(None).expect_err("expected missing-binary error");
571        let chain = format!("{err:#}");
572        assert!(chain.contains("Claude Code CLI not available"));
573    }
574}