vtcode_core/utils/
ansi.rs

1use crate::ui::theme;
2use crate::utils::transcript;
3use anstream::{AutoStream, ColorChoice};
4use anstyle::{Reset, Style};
5use anstyle_query::{clicolor, clicolor_force, no_color, term_supports_color};
6use anyhow::Result;
7use std::io::{self, Write};
8
9/// Styles available for rendering messages
10#[derive(Clone, Copy)]
11pub enum MessageStyle {
12    Info,
13    Error,
14    Output,
15    Response,
16    Tool,
17    User,
18}
19
20impl MessageStyle {
21    fn style(self) -> Style {
22        let styles = theme::active_styles();
23        match self {
24            Self::Info => styles.info,
25            Self::Error => styles.error,
26            Self::Output => styles.output,
27            Self::Response => styles.response,
28            Self::Tool => styles.tool,
29            Self::User => styles.user,
30        }
31    }
32
33    fn indent(self) -> &'static str {
34        match self {
35            Self::Response | Self::Tool => "  ",
36            _ => "",
37        }
38    }
39}
40
41/// Renderer with deferred output buffering
42pub struct AnsiRenderer {
43    writer: AutoStream<io::Stdout>,
44    buffer: String,
45    color: bool,
46}
47
48impl AnsiRenderer {
49    /// Create a new renderer for stdout
50    pub fn stdout() -> Self {
51        let color =
52            clicolor_force() || (!no_color() && clicolor().unwrap_or_else(term_supports_color));
53        let choice = if color {
54            ColorChoice::Auto
55        } else {
56            ColorChoice::Never
57        };
58        Self {
59            writer: AutoStream::new(std::io::stdout(), choice),
60            buffer: String::new(),
61            color,
62        }
63    }
64
65    /// Push text into the buffer
66    pub fn push(&mut self, text: &str) {
67        self.buffer.push_str(text);
68    }
69
70    /// Flush the buffer with the given style
71    pub fn flush(&mut self, style: MessageStyle) -> Result<()> {
72        let style = style.style();
73        if self.color {
74            writeln!(self.writer, "{style}{}{Reset}", self.buffer)?;
75        } else {
76            writeln!(self.writer, "{}", self.buffer)?;
77        }
78        self.writer.flush()?;
79        transcript::append(&self.buffer);
80        self.buffer.clear();
81        Ok(())
82    }
83
84    /// Convenience for writing a single line
85    pub fn line(&mut self, style: MessageStyle, text: &str) -> Result<()> {
86        let indent = style.indent();
87
88        if text.contains('\n') {
89            let trailing_newline = text.ends_with('\n');
90            for line in text.lines() {
91                self.buffer.clear();
92                if !indent.is_empty() && !line.is_empty() {
93                    self.buffer.push_str(indent);
94                }
95                self.buffer.push_str(line);
96                self.flush(style)?;
97            }
98            if trailing_newline {
99                self.buffer.clear();
100                if !indent.is_empty() {
101                    self.buffer.push_str(indent);
102                }
103                self.flush(style)?;
104            }
105            Ok(())
106        } else {
107            self.buffer.clear();
108            if !indent.is_empty() && !text.is_empty() {
109                self.buffer.push_str(indent);
110            }
111            self.buffer.push_str(text);
112            self.flush(style)
113        }
114    }
115
116    /// Write styled text without a trailing newline
117    pub fn inline_with_style(&mut self, style: Style, text: &str) -> Result<()> {
118        if self.color {
119            write!(self.writer, "{style}{}{Reset}", text)?;
120        } else {
121            write!(self.writer, "{}", text)?;
122        }
123        self.writer.flush()?;
124        Ok(())
125    }
126
127    /// Write a line with an explicit style
128    pub fn line_with_style(&mut self, style: Style, text: &str) -> Result<()> {
129        if self.color {
130            writeln!(self.writer, "{style}{}{Reset}", text)?;
131        } else {
132            writeln!(self.writer, "{}", text)?;
133        }
134        self.writer.flush()?;
135        transcript::append(text);
136        Ok(())
137    }
138
139    /// Write a raw line without styling
140    pub fn raw_line(&mut self, text: &str) -> Result<()> {
141        writeln!(self.writer, "{}", text)?;
142        self.writer.flush()?;
143        transcript::append(text);
144        Ok(())
145    }
146}
147
148#[cfg(test)]
149mod tests {
150    use super::*;
151
152    #[test]
153    fn test_styles_construct() {
154        let info = MessageStyle::Info.style();
155        assert_eq!(info, MessageStyle::Info.style());
156        let resp = MessageStyle::Response.style();
157        assert_eq!(resp, MessageStyle::Response.style());
158        let tool = MessageStyle::Tool.style();
159        assert_eq!(tool, MessageStyle::Tool.style());
160    }
161
162    #[test]
163    fn test_renderer_buffer() {
164        let mut r = AnsiRenderer::stdout();
165        r.push("hello");
166        assert_eq!(r.buffer, "hello");
167    }
168}