steer_tui/tui/widgets/chat_widgets/

message_widget.rs

use ratatui::text::{Line, Span};
use steer_core::app::conversation::{AppCommandType, AssistantContent, MessageData};
use steer_core::app::conversation::{Message, UserContent};

use crate::tui::theme::{Component, Theme};
use crate::tui::widgets::formatters::helpers::style_wrap;
use crate::tui::widgets::{ChatRenderable, HeightCache, ViewMode, markdown};

pub struct MessageWidget {
    message: Message,
    cache: HeightCache,
    rendered_lines: Option<Vec<Line<'static>>>,
}

impl MessageWidget {
    pub fn new(message: Message) -> Self {
        Self {
            message,
            cache: HeightCache::new(),
            rendered_lines: None,
        }
    }

    fn render_as_markdown(text: &str, theme: &Theme, max_width: usize) -> markdown::MarkedText {
        let markdown_styles = markdown::MarkdownStyles::from_theme(theme);
        markdown::from_str_with_width(text, &markdown_styles, theme, Some(max_width as u16))
    }
}

impl ChatRenderable for MessageWidget {
    fn lines(&mut self, width: u16, _mode: ViewMode, theme: &Theme) -> &[Line<'static>] {
        if self.rendered_lines.is_some() && self.cache.last_width == width {
            return self.rendered_lines.as_ref().unwrap();
        }

        let max_width = width.saturating_sub(4) as usize; // Account for gutters
        let mut lines = Vec::new();

        match &self.message.data {
            MessageData::User { content, .. } => {
                for user_content in content {
                    match user_content {
                        UserContent::Text { text } => {
                            let marked_text = Self::render_as_markdown(text, theme, max_width);
                            for marked_line in marked_text.lines {
                                if marked_line.no_wrap {
                                    // Don't wrap code block lines
                                    lines.push(marked_line.line);
                                } else {
                                    // Wrap normal lines
                                    let wrapped = style_wrap(marked_line.line, max_width as u16);
                                    for line in wrapped {
                                        lines.push(line);
                                    }
                                }
                            }
                        }
                        UserContent::CommandExecution {
                            command,
                            stdout,
                            stderr,
                            exit_code,
                        } => {
                            // Format command execution
                            let cmd_style = theme.style(Component::CommandPrompt);
                            lines.push(Line::from(Span::styled(format!("$ {command}"), cmd_style)));

                            if !stdout.is_empty() {
                                let output_style = theme.style(Component::UserMessage);
                                for line in stdout.lines() {
                                    let wrapped = textwrap::wrap(line, max_width);
                                    for wrapped_line in wrapped {
                                        lines.push(Line::from(Span::styled(
                                            wrapped_line.to_string(),
                                            output_style,
                                        )));
                                    }
                                }
                            }

                            if !stderr.is_empty() {
                                let error_style = theme.style(Component::ErrorText);
                                for line in stderr.lines() {
                                    let wrapped = textwrap::wrap(line, max_width);
                                    for wrapped_line in wrapped {
                                        lines.push(Line::from(Span::styled(
                                            wrapped_line.to_string(),
                                            error_style,
                                        )));
                                    }
                                }
                            }

                            if *exit_code != 0 {
                                lines.push(Line::from(Span::styled(
                                    format!("Exit code: {exit_code}"),
                                    theme.style(Component::DimText),
                                )));
                            }
                        }
                        UserContent::AppCommand { command, response } => {
                            // Format app command
                            let cmd_style = theme.style(Component::CommandPrompt);
                            let cmd_text = match command {
                                AppCommandType::Model { target } => {
                                    if let Some(model) = target {
                                        format!("/model {model}")
                                    } else {
                                        "/model".to_string()
                                    }
                                }
                                AppCommandType::Compact => "/compact".to_string(),
                                AppCommandType::Clear => "/clear".to_string(),
                            };
                            lines.push(Line::from(Span::styled(cmd_text, cmd_style)));

                            if let Some(resp) = response {
                                let resp_text = match resp {
                                    steer_core::app::conversation::CommandResponse::Text(text) => text.clone(),
                                    steer_core::app::conversation::CommandResponse::Compact(result) => {
                                        match result {
                                            steer_core::app::conversation::CompactResult::Success(summary) => summary.clone(),
                                            steer_core::app::conversation::CompactResult::Cancelled => "Compact cancelled.".to_string(),
                                            steer_core::app::conversation::CompactResult::InsufficientMessages => "Not enough messages to compact.".to_string(),
                                        }
                                    }
                                };

                                let resp_style = theme.style(Component::CommandText);
                                for line in resp_text.lines() {
                                    let wrapped = textwrap::wrap(line, max_width);
                                    for wrapped_line in wrapped {
                                        lines.push(Line::from(Span::styled(
                                            wrapped_line.to_string(),
                                            resp_style,
                                        )));
                                    }
                                }
                            }
                        }
                    }
                }
            }
            MessageData::Assistant { content, .. } => {
                for block in content {
                    match block {
                        AssistantContent::Text { text } => {
                            if text.trim().is_empty() {
                                continue;
                            }

                            let marked_text = Self::render_as_markdown(text, theme, max_width);
                            for marked_line in marked_text.lines {
                                if marked_line.no_wrap {
                                    // Don't wrap code block lines
                                    lines.push(marked_line.line);
                                } else {
                                    // Wrap normal lines
                                    let wrapped = style_wrap(marked_line.line, max_width as u16);
                                    for line in wrapped {
                                        lines.push(line);
                                    }
                                }
                            }
                        }
                        AssistantContent::ToolCall { .. } => {
                            // Tool calls are rendered separately
                            continue;
                        }
                        AssistantContent::Thought { thought } => {
                            let thought_text = thought.display_text();
                            let thought_style = theme.style(Component::ThoughtText);

                            // Parse markdown for the thought
                            let markdown_styles = markdown::MarkdownStyles::from_theme(theme);
                            let markdown_text = markdown::from_str_with_width(
                                &thought_text,
                                &markdown_styles,
                                theme,
                                Some(max_width as u16),
                            );

                            // Process each line with thought styling
                            for marked_line in markdown_text.lines {
                                let mut styled_spans = Vec::new();

                                // Apply thought style to all spans
                                for span in marked_line.line.spans {
                                    styled_spans.push(Span::styled(
                                        span.content.into_owned(),
                                        thought_style,
                                    ));
                                }

                                let thought_line = Line::from(styled_spans);

                                if marked_line.no_wrap {
                                    lines.push(thought_line);
                                } else {
                                    let wrapped = style_wrap(thought_line, max_width as u16);
                                    for line in wrapped {
                                        lines.push(line);
                                    }
                                }
                            }
                        }
                    }
                }
            }
            MessageData::Tool { .. } => {
                // Tools are rendered as part of ToolInteraction blocks
            }
        }

        self.rendered_lines = Some(lines);
        self.rendered_lines.as_ref().unwrap()
    }
}

#[cfg(test)]
mod tests {
    use super::MessageWidget;
    use crate::tui::theme::Theme;
    use crate::tui::widgets::ChatRenderable;
    use crate::tui::widgets::ViewMode;
    use steer_core::app::conversation::AssistantContent;
    use steer_core::app::conversation::{Message, MessageData, UserContent};

    #[test]
    fn test_message_widget_user_text() {
        let theme = Theme::default();
        let user_msg = Message {
            data: MessageData::User {
                content: vec![UserContent::Text {
                    text: "Hello, world!".to_string(),
                }],
            },
            timestamp: 0,
            id: "test-id".to_string(),
            parent_message_id: None,
        };

        let mut widget = MessageWidget::new(user_msg);

        // Test height calculation
        let height = widget.lines(20, ViewMode::Compact, &theme).len();
        assert_eq!(height, 1); // Single line message

        // Test with wrapping
        let height_wrapped = widget.lines(8, ViewMode::Compact, &theme).len();
        assert!(height_wrapped > 1); // Should wrap
    }

    #[test]
    fn test_message_widget_assistant_text() {
        let theme = Theme::default();
        let assistant_msg = Message {
            data: MessageData::Assistant {
                content: vec![AssistantContent::Text {
                    text: "Hello from assistant".to_string(),
                }],
            },
            timestamp: 0,
            id: "test-id".to_string(),
            parent_message_id: None,
        };

        let mut widget = MessageWidget::new(assistant_msg);

        // Test height calculation
        let height = widget.lines(30, ViewMode::Compact, &theme).len();
        assert_eq!(height, 1); // Single line message
    }

    #[test]
    fn test_message_widget_command_execution() {
        let theme = Theme::default();
        let cmd_msg = Message {
            data: MessageData::User {
                content: vec![UserContent::CommandExecution {
                    command: "ls -la".to_string(),
                    stdout: "file1.txt\nfile2.txt".to_string(),
                    stderr: "".to_string(),
                    exit_code: 0,
                }],
            },
            timestamp: 0,
            id: "test-id".to_string(),
            parent_message_id: None,
        };

        let mut widget = MessageWidget::new(cmd_msg);

        // Test height calculation - should have command line + 2 output lines
        let height = widget.lines(30, ViewMode::Compact, &theme).len();
        assert_eq!(height, 3); // $ ls -la + file1.txt + file2.txt
    }

    #[test]
    fn test_unicode_width_handling() {
        use ratatui::buffer::Buffer;
        use ratatui::layout::Rect;

        let theme = Theme::default();

        // Create a message with various Unicode characters
        let unicode_msg = Message {
            data: MessageData::User {
                content: vec![UserContent::Text {
                    text: "Hello 你好 👋 café".to_string(),
                }],
            },
            timestamp: 0,
            id: "test-unicode".to_string(),
            parent_message_id: None,
        };

        let mut widget = MessageWidget::new(unicode_msg);

        // Create buffers for both render methods
        let area = Rect::new(0, 0, 50, 5);
        let buf_regular = Buffer::empty(area);
        let buf_partial = Buffer::empty(area);

        // Render with regular method
        widget.lines(area.width, ViewMode::Compact, &theme);

        // Render with partial method
        widget.lines(area.width, ViewMode::Compact, &theme);

        // Compare the buffers - they should be identical
        for y in 0..area.height {
            for x in 0..area.width {
                let regular_cell = buf_regular.cell((x, y)).unwrap();
                let partial_cell = buf_partial.cell((x, y)).unwrap();

                assert_eq!(
                    regular_cell.symbol(),
                    partial_cell.symbol(),
                    "Mismatch at ({}, {}): regular='{}' partial='{}'",
                    x,
                    y,
                    regular_cell.symbol(),
                    partial_cell.symbol()
                );
            }
        }
    }

    #[test]
    fn test_wide_character_positioning() {
        use ratatui::buffer::Buffer;
        use ratatui::layout::Rect;

        let theme = Theme::default();

        // Create a message with wide characters (CJK takes 2 columns each)
        let wide_msg = Message {
            data: MessageData::User {
                content: vec![UserContent::Text {
                    text: "A中B".to_string(), // A=1 col, 中=2 cols, B=1 col
                }],
            },
            timestamp: 0,
            id: "test-wide".to_string(),
            parent_message_id: None,
        };

        let mut widget = MessageWidget::new(wide_msg);

        // Render with partial method
        let area = Rect::new(0, 0, 10, 1);
        let buf = Buffer::empty(area);
        widget.lines(area.width, ViewMode::Compact, &theme);

        // With correct Unicode handling:
        // Position 0: 'A' (1 column)
        // Position 1-2: '中' (2 columns)
        // Position 3: 'B' (1 column)

        // With the bug (incrementing by 1):
        // Position 0: 'A'
        // Position 1: '中' (but should occupy 2 columns)
        // Position 2: 'B' (overlapping with second half of '中')

        // This test will fail with current implementation
        // because 'B' will be at position 2 instead of position 3

        // Check that B is not at position 2 (would indicate the bug)
        let cell_at_2 = buf.cell((2, 0)).unwrap();
        assert_ne!(
            cell_at_2.symbol(),
            "B",
            "Character 'B' incorrectly positioned due to Unicode width bug"
        );
    }
}