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

vtcode_core/utils/
ansi.rs
