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.
91#[must_use]
92pub fn no_color_env_active() -> bool {
93    no_color_env_active_from(std::env::var_os("NO_COLOR"))
94}
95
96/// Read the current runtime color policy.
97pub fn current_color_output_policy() -> ColorOutputPolicy {
98    Lazy::force(&INIT_FROM_ENV);
99    ColorOutputPolicy {
100        enabled: POLICY_ENABLED.load(Ordering::Relaxed),
101        source: decode_source(POLICY_SOURCE.load(Ordering::Relaxed)),
102    }
103}
104
105/// Replace the current runtime color policy.
106pub fn set_color_output_policy(policy: ColorOutputPolicy) {
107    POLICY_ENABLED.store(policy.enabled, Ordering::Relaxed);
108    POLICY_SOURCE.store(encode_source(policy.source), Ordering::Relaxed);
109}
110
111/// Reset runtime color policy from environment defaults.
112pub fn reset_color_output_policy_from_env() {
113    set_color_output_policy(detect_policy_from_env());
114}
115
116/// Returns true when runtime color output is enabled.
117#[must_use]
118pub fn color_output_enabled() -> bool {
119    current_color_output_policy().enabled
120}
121
122#[cfg(test)]
123mod tests {
124    use super::no_color_env_active_from;
125    use std::ffi::OsString;
126
127    #[test]
128    fn no_color_requires_non_empty_value() {
129        assert!(!no_color_env_active_from(None));
130        assert!(!no_color_env_active_from(Some(OsString::from(""))));
131        assert!(no_color_env_active_from(Some(OsString::from("1"))));
132    }
133}