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.
218pub fn check_github_cli() -> Result<()> {
219    // Check if gh CLI is available
220    let gh_check = std::process::Command::new("gh")
221        .args(["--version"])
222        .output();
223
224    match gh_check {
225        Ok(output) if output.status.success() => {
226            // Test if gh can access the current repo
227            let repo_check = std::process::Command::new("gh")
228                .args(["repo", "view", "--json", "name"])
229                .output();
230
231            match repo_check {
232                Ok(repo_output) if repo_output.status.success() => Ok(()),
233                Ok(repo_output) => {
234                    let error_details = String::from_utf8_lossy(&repo_output.stderr);
235                    if error_details.contains("authentication") || error_details.contains("login") {
236                        bail!(
237                            "GitHub CLI authentication failed.\n\
238                             Please run 'gh auth login' or set GITHUB_TOKEN environment variable."
239                        )
240                    }
241                    bail!(
242                        "GitHub CLI cannot access this repository.\n\
243                         Error: {}",
244                        error_details.trim()
245                    )
246                }
247                Err(e) => bail!("Failed to test GitHub CLI access: {e}"),
248            }
249        }
250        _ => bail!(
251            "GitHub CLI (gh) is not installed or not in PATH.\n\
252             Please install it from https://cli.github.com/"
253        ),
254    }
255}
256
257/// Validates that the current directory is in a valid git repository.
258///
259/// This is a lightweight check that opens the repository without
260/// loading any commit data.
261pub fn check_git_repository() -> Result<()> {
262    crate::git::GitRepository::open().context(
263        "Not in a git repository. Please run this command from within a git repository.",
264    )?;
265    Ok(())
266}
267
268/// Validates that the working directory is clean (no uncommitted changes).
269///
270/// This checks for:
271/// - Staged changes
272/// - Unstaged modifications
273/// - Untracked files (excluding ignored files)
274///
275/// Use this before operations that require a clean working directory,
276/// like amending commits.
277pub fn check_working_directory_clean() -> Result<()> {
278    let repo = crate::git::GitRepository::open().context("Failed to open git repository")?;
279
280    let status = repo
281        .get_working_directory_status()
282        .context("Failed to get working directory status")?;
283
284    if !status.clean {
285        let mut message = String::from("Working directory has uncommitted changes:\n");
286        for change in &status.untracked_changes {
287            message.push_str(&format!("  {} {}\n", change.status, change.file));
288        }
289        message.push_str("\nPlease commit or stash your changes before proceeding.");
290        bail!(message);
291    }
292
293    Ok(())
294}
295
296/// Performs combined preflight check for AI commands.
297///
298/// Validates:
299/// - Git repository access
300/// - AI credentials
301///
302/// Returns information about the AI provider that will be used.
303pub fn check_ai_command_prerequisites(model_override: Option<&str>) -> Result<AiCredentialInfo> {
304    check_git_repository()?;
305    check_ai_credentials(model_override)
306}
307
308/// Performs combined preflight check for PR creation.
309///
310/// Validates:
311/// - Git repository access
312/// - AI credentials
313/// - GitHub CLI availability and authentication
314///
315/// Returns information about the AI provider that will be used.
316pub fn check_pr_command_prerequisites(model_override: Option<&str>) -> Result<AiCredentialInfo> {
317    check_git_repository()?;
318    let ai_info = check_ai_credentials(model_override)?;
319    check_github_cli()?;
320    Ok(ai_info)
321}
322
323#[cfg(test)]
324#[allow(clippy::unwrap_used, clippy::expect_used)]
325mod tests {
326    use super::*;
327
328    use std::env;
329    use std::sync::Mutex;
330    use std::sync::OnceLock;
331
332    /// Global lock to ensure environment variable tests don't interfere with each other.
333    static ENV_TEST_LOCK: OnceLock<Mutex<()>> = OnceLock::new();
334
335    /// Manages environment variables in tests to avoid interference.
336    struct EnvGuard {
337        _lock: std::sync::MutexGuard<'static, ()>,
338        vars: Vec<(String, Option<String>)>,
339    }
340
341    impl EnvGuard {
342        fn new() -> Self {
343            let lock = ENV_TEST_LOCK.get_or_init(|| Mutex::new(())).lock().unwrap();
344            Self {
345                _lock: lock,
346                vars: Vec::new(),
347            }
348        }
349
350        fn set(&mut self, key: &str, value: &str) {
351            let original = env::var(key).ok();
352            self.vars.push((key.to_string(), original));
353            env::set_var(key, value);
354        }
355
356        fn remove(&mut self, key: &str) {
357            let original = env::var(key).ok();
358            self.vars.push((key.to_string(), original));
359            env::remove_var(key);
360        }
361    }
362
363    impl Drop for EnvGuard {
364        fn drop(&mut self) {
365            for (key, original_value) in self.vars.drain(..).rev() {
366                match original_value {
367                    Some(value) => env::set_var(&key, value),
368                    None => env::remove_var(&key),
369                }
370            }
371        }
372    }
373
374    #[test]
375    fn ai_provider_display() {
376        assert_eq!(format!("{}", AiProvider::Claude), "Claude API");
377        assert_eq!(format!("{}", AiProvider::Bedrock), "AWS Bedrock");
378        assert_eq!(format!("{}", AiProvider::OpenAi), "OpenAI API");
379        assert_eq!(format!("{}", AiProvider::Ollama), "Ollama");
380        assert_eq!(format!("{}", AiProvider::ClaudeCli), "Claude Code CLI");
381    }
382
383    #[test]
384    fn ai_provider_equality() {
385        assert_eq!(AiProvider::Claude, AiProvider::Claude);
386        assert_ne!(AiProvider::Claude, AiProvider::OpenAi);
387        assert_ne!(AiProvider::Bedrock, AiProvider::Ollama);
388    }
389
390    #[test]
391    fn ai_provider_clone() {
392        let provider = AiProvider::Bedrock;
393        let cloned = provider;
394        assert_eq!(provider, cloned);
395    }
396
397    #[test]
398    fn ai_provider_debug() {
399        let debug_str = format!("{:?}", AiProvider::Claude);
400        assert_eq!(debug_str, "Claude");
401    }
402
403    #[test]
404    fn ai_credential_info_debug() {
405        let info = AiCredentialInfo {
406            provider: AiProvider::Ollama,
407            model: "llama2".to_string(),
408        };
409        let debug_str = format!("{info:?}");
410        assert!(debug_str.contains("Ollama"));
411        assert!(debug_str.contains("llama2"));
412    }
413
414    #[test]
415    fn claude_default_model_from_registry() {
416        let mut guard = EnvGuard::new();
417        // Enable Claude API path with a dummy key, no model override
418        guard.remove("USE_OPENAI");
419        guard.remove("USE_OLLAMA");
420        guard.remove("CLAUDE_CODE_USE_BEDROCK");
421        guard.remove("ANTHROPIC_MODEL");
422        guard.set("ANTHROPIC_API_KEY", "sk-test-dummy");
423
424        let info = check_ai_credentials(None).unwrap();
425        assert_eq!(info.provider, AiProvider::Claude);
426        assert_eq!(info.model, "claude-sonnet-4-6");
427    }
428
429    #[test]
430    fn openai_default_model_from_registry() {
431        let mut guard = EnvGuard::new();
432        guard.set("USE_OPENAI", "true");
433        guard.remove("USE_OLLAMA");
434        guard.remove("OPENAI_MODEL");
435        guard.set("OPENAI_API_KEY", "sk-test-dummy");
436
437        let info = check_ai_credentials(None).unwrap();
438        assert_eq!(info.provider, AiProvider::OpenAi);
439        assert_eq!(info.model, "gpt-5-mini");
440    }
441
442    #[test]
443    fn bedrock_default_model_from_registry() {
444        let mut guard = EnvGuard::new();
445        guard.remove("USE_OPENAI");
446        guard.remove("USE_OLLAMA");
447        guard.set("CLAUDE_CODE_USE_BEDROCK", "true");
448        guard.remove("ANTHROPIC_MODEL");
449        guard.set("ANTHROPIC_AUTH_TOKEN", "test-token");
450        guard.set("ANTHROPIC_BEDROCK_BASE_URL", "https://bedrock.example.com");
451
452        let info = check_ai_credentials(None).unwrap();
453        assert_eq!(info.provider, AiProvider::Bedrock);
454        assert_eq!(info.model, "claude-sonnet-4-6");
455    }
456
457    #[test]
458    fn model_override_takes_precedence() {
459        let mut guard = EnvGuard::new();
460        guard.remove("USE_OPENAI");
461        guard.remove("USE_OLLAMA");
462        guard.remove("CLAUDE_CODE_USE_BEDROCK");
463        guard.remove("ANTHROPIC_MODEL");
464        guard.set("ANTHROPIC_API_KEY", "sk-test-dummy");
465
466        let info = check_ai_credentials(Some("claude-opus-4-6")).unwrap();
467        assert_eq!(info.model, "claude-opus-4-6");
468    }
469
470    #[cfg(unix)]
471    fn make_version_shim(tmp: &tempfile::TempDir, exit_code: i32) -> std::path::PathBuf {
472        let shim = tmp.path().join("claude-bin-shim");
473        crate::test_support::shim::write_exec_script(
474            &shim,
475            &format!("#!/bin/sh\necho 'fake-claude 0.0.0'\nexit {exit_code}\n"),
476        );
477        shim
478    }
479
480    #[test]
481    #[cfg(unix)]
482    fn claude_cli_backend_uses_version_probe() {
483        let _guard = crate::test_support::shim::shim_lock();
484        let tmp = tempfile::TempDir::new().unwrap();
485        let shim = make_version_shim(&tmp, 0);
486
487        let mut guard = EnvGuard::new();
488        guard.remove("USE_OPENAI");
489        guard.remove("USE_OLLAMA");
490        guard.remove("CLAUDE_CODE_USE_BEDROCK");
491        guard.remove("ANTHROPIC_MODEL");
492        guard.remove("CLAUDE_MODEL");
493        guard.remove("CLAUDE_CODE_MODEL");
494        guard.set("OMNI_DEV_AI_BACKEND", "claude-cli");
495        guard.set("OMNI_DEV_CLAUDE_CLI_BIN", shim.to_str().unwrap());
496
497        let info = check_ai_credentials(None).unwrap();
498        assert_eq!(info.provider, AiProvider::ClaudeCli);
499        assert_eq!(info.model, "claude-sonnet-4-6");
500    }
501
502    #[test]
503    #[cfg(unix)]
504    fn claude_cli_backend_uses_model_from_env() {
505        let _guard = crate::test_support::shim::shim_lock();
506        let tmp = tempfile::TempDir::new().unwrap();
507        let shim = make_version_shim(&tmp, 0);
508
509        let mut guard = EnvGuard::new();
510        guard.remove("USE_OPENAI");
511        guard.remove("USE_OLLAMA");
512        guard.remove("CLAUDE_CODE_USE_BEDROCK");
513        guard.remove("ANTHROPIC_MODEL");
514        guard.remove("CLAUDE_CODE_MODEL");
515        guard.set("OMNI_DEV_AI_BACKEND", "claude-cli");
516        guard.set("OMNI_DEV_CLAUDE_CLI_BIN", shim.to_str().unwrap());
517        guard.set("CLAUDE_MODEL", "haiku");
518
519        let info = check_ai_credentials(None).unwrap();
520        assert_eq!(info.provider, AiProvider::ClaudeCli);
521        assert_eq!(info.model, "haiku");
522    }
523
524    #[test]
525    fn claude_cli_backend_missing_binary_fails_preflight() {
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.set("OMNI_DEV_AI_BACKEND", "claude-cli");
531        guard.set("OMNI_DEV_CLAUDE_CLI_BIN", "/nonexistent/claude-binary-xyz");
532
533        let err = check_ai_credentials(None).expect_err("expected missing-binary error");
534        let chain = format!("{err:#}");
535        assert!(
536            chain.contains("Claude Code CLI not available"),
537            "unexpected error: {chain}"
538        );
539    }
540
541    #[test]
542    fn claude_cli_backend_accepts_underscore_alias() {
543        // The factory/preflight accept both `claude-cli` and `claude_cli`.
544        // Verify the second spelling routes the same way (missing-binary
545        // path exercises the selector cheaply).
546        let mut guard = EnvGuard::new();
547        guard.remove("USE_OPENAI");
548        guard.remove("USE_OLLAMA");
549        guard.remove("CLAUDE_CODE_USE_BEDROCK");
550        guard.set("OMNI_DEV_AI_BACKEND", "claude_cli");
551        guard.set("OMNI_DEV_CLAUDE_CLI_BIN", "/nonexistent/claude-binary-xyz");
552
553        let err = check_ai_credentials(None).expect_err("expected missing-binary error");
554        let chain = format!("{err:#}");
555        assert!(chain.contains("Claude Code CLI not available"));
556    }
557}