Skip to main content

vtcode_commons/
color_policy.rs

1//! Runtime color output policy helpers.
2//!
3//! This module centralizes color enable/disable decisions for CLI and
4//! transcript-style output paths. By default it follows the NO_COLOR
5//! environment variable with strict "present and non-empty" semantics.
6
7use once_cell::sync::Lazy;
8use std::ffi::OsString;
9use std::sync::atomic::{AtomicBool, AtomicU8, Ordering};
10
11/// Source that determined the active runtime color policy.
12#[derive(Clone, Copy, Debug, PartialEq, Eq)]
13pub enum ColorOutputPolicySource {
14    /// Default runtime behavior (auto detect + env hints).
15    DefaultAuto,
16    /// Disabled due to NO_COLOR environment variable.
17    NoColorEnv,
18    /// Disabled due to explicit `--no-color`.
19    CliNoColor,
20    /// Disabled due to explicit `--color never`.
21    CliColorNever,
22    /// Enabled due to explicit `--color always`.
23    CliColorAlways,
24    /// Enabled or disabled by explicit config override.
25    ConfigOverride,
26}
27
28/// Runtime color output policy.
29#[derive(Clone, Copy, Debug, PartialEq, Eq)]
30pub struct ColorOutputPolicy {
31    pub enabled: bool,
32    pub source: ColorOutputPolicySource,
33}
34
35const SOURCE_DEFAULT_AUTO: u8 = 0;
36const SOURCE_NO_COLOR_ENV: u8 = 1;
37const SOURCE_CLI_NO_COLOR: u8 = 2;
38const SOURCE_CLI_COLOR_NEVER: u8 = 3;
39const SOURCE_CLI_COLOR_ALWAYS: u8 = 4;
40const SOURCE_CONFIG_OVERRIDE: u8 = 5;
41
42static POLICY_ENABLED: AtomicBool = AtomicBool::new(true);
43static POLICY_SOURCE: AtomicU8 = AtomicU8::new(SOURCE_DEFAULT_AUTO);
44
45static INIT_FROM_ENV: Lazy<()> = Lazy::new(|| {
46    let default_policy = detect_policy_from_env();
47    set_color_output_policy(default_policy);
48});
49
50fn detect_policy_from_env() -> ColorOutputPolicy {
51    if no_color_env_active() {
52        ColorOutputPolicy {
53            enabled: false,
54            source: ColorOutputPolicySource::NoColorEnv,
55        }
56    } else {
57        ColorOutputPolicy {
58            enabled: true,
59            source: ColorOutputPolicySource::DefaultAuto,
60        }
61    }
62}
63
64fn encode_source(source: ColorOutputPolicySource) -> u8 {
65    match source {
66        ColorOutputPolicySource::DefaultAuto => SOURCE_DEFAULT_AUTO,
67        ColorOutputPolicySource::NoColorEnv => SOURCE_NO_COLOR_ENV,
68        ColorOutputPolicySource::CliNoColor => SOURCE_CLI_NO_COLOR,
69        ColorOutputPolicySource::CliColorNever => SOURCE_CLI_COLOR_NEVER,
70        ColorOutputPolicySource::CliColorAlways => SOURCE_CLI_COLOR_ALWAYS,
71        ColorOutputPolicySource::ConfigOverride => SOURCE_CONFIG_OVERRIDE,
72    }
73}
74
75fn decode_source(value: u8) -> ColorOutputPolicySource {
76    match value {
77        SOURCE_NO_COLOR_ENV => ColorOutputPolicySource::NoColorEnv,
78        SOURCE_CLI_NO_COLOR => ColorOutputPolicySource::CliNoColor,
79        SOURCE_CLI_COLOR_NEVER => ColorOutputPolicySource::CliColorNever,
80        SOURCE_CLI_COLOR_ALWAYS => ColorOutputPolicySource::CliColorAlways,
81        SOURCE_CONFIG_OVERRIDE => ColorOutputPolicySource::ConfigOverride,
82        _ => ColorOutputPolicySource::DefaultAuto,
83    }
84}
85
86fn no_color_env_active_from(value: Option<OsString>) -> bool {
87    value.map(|v| !v.is_empty()).unwrap_or(false)
88}
89
90/// Returns true when NO_COLOR is present and non-empty.
91pub fn no_color_env_active() -> bool {
92    no_color_env_active_from(std::env::var_os("NO_COLOR"))
93}
94
95/// Read the current runtime color policy.
96pub fn current_color_output_policy() -> ColorOutputPolicy {
97    Lazy::force(&INIT_FROM_ENV);
98    ColorOutputPolicy {
99        enabled: POLICY_ENABLED.load(Ordering::Relaxed),
100        source: decode_source(POLICY_SOURCE.load(Ordering::Relaxed)),
101    }
102}
103
104/// Replace the current runtime color policy.
105pub fn set_color_output_policy(policy: ColorOutputPolicy) {
106    POLICY_ENABLED.store(policy.enabled, Ordering::Relaxed);
107    POLICY_SOURCE.store(encode_source(policy.source), Ordering::Relaxed);
108}
109
110/// Reset runtime color policy from environment defaults.
111pub fn reset_color_output_policy_from_env() {
112    set_color_output_policy(detect_policy_from_env());
113}
114
115/// Returns true when runtime color output is enabled.
116pub fn color_output_enabled() -> bool {
117    current_color_output_policy().enabled
118}
119
120#[cfg(test)]
121mod tests {
122    use super::no_color_env_active_from;
123    use std::ffi::OsString;
124
125    #[test]
126    fn no_color_requires_non_empty_value() {
127        assert!(!no_color_env_active_from(None));
128        assert!(!no_color_env_active_from(Some(OsString::from(""))));
129        assert!(no_color_env_active_from(Some(OsString::from("1"))));
130    }
131}