rab/agent/ui/components/
assistant_message.rs

1use std::sync::Arc;
2
3use crate::agent::ui::theme::ThemeKey;
4use crate::agent::ui::theme::current_theme;
5use crate::tui::Component;
6use crate::tui::component::RenderCacheKey;
7use crate::tui::components::markdown::{DefaultTextStyle, Markdown, StyleFn};
8const OSC133_ZONE_START: &str = "\x1b]133;A\x07";
9const OSC133_ZONE_END: &str = "\x1b]133;B\x07";
10const OSC133_ZONE_FINAL: &str = "\x1b]133;C\x07";
11
12/// Assistant message component - matches pi's AssistantMessageComponent.
13/// Renders text content with Markdown, optional thinking blocks.
14pub struct AssistantMessageComponent {
15    text: String,
16    thinking: Vec<ThinkingBlock>,
17    hide_thinking: bool,
18    cached_lines: Option<Vec<String>>,
19    cached_width: usize,
20    /// Persistent Markdown instance for the text content.
21    /// Reused across renders so its internal cache (cached_text + cached_lines)
22    /// avoids re-parsing when text hasn't changed (e.g. spinner ticks between deltas).
23    text_md: Option<Markdown>,
24    /// Persistent Markdown instances for each thinking block.
25    thinking_md: Vec<Markdown>,
26}
27
28pub struct ThinkingBlock {
29    pub text: String,
30    pub level: Option<String>,
31}
32
33impl AssistantMessageComponent {
34    pub fn new(text: impl Into<String>) -> Self {
35        Self {
36            text: text.into(),
37            thinking: Vec::new(),
38            hide_thinking: false,
39            cached_lines: None,
40            cached_width: 0,
41            text_md: None,
42            thinking_md: Vec::new(),
43        }
44    }
45
46    /// Ensure the persistent text Markdown instance matches the current text.
47    fn sync_text_md(&mut self) {
48        if self.text.is_empty() {
49            self.text_md = None;
50            return;
51        }
52        let md_theme = crate::agent::ui::theme::get_markdown_theme();
53        let should_recreate = match self.text_md {
54            Some(ref mut md) => {
55                let needs_update = !md.cached_text_matches(&self.text);
56                if needs_update {
57                    md.set_text(&self.text);
58                }
59                false // reuse existing
60            }
61            None => true, // create new
62        };
63        if should_recreate {
64            self.text_md = Some(Markdown::new(self.text.clone(), 1, 0, md_theme, None));
65        }
66    }
67
68    /// Ensure the persistent thinking Markdown instances match current thinking blocks.
69    fn sync_thinking_md(&mut self) {
70        // Remove excess cached instances
71        while self.thinking_md.len() > self.thinking.len() {
72            self.thinking_md.pop();
73        }
74        for (i, block) in self.thinking.iter().enumerate() {
75            if block.text.trim().is_empty() {
76                continue;
77            }
78            if i >= self.thinking_md.len() {
79                // Create new Markdown for this block
80                let color_fn: StyleFn = Arc::new(|s: &str| -> String {
81                    crate::agent::ui::theme::current_theme().fg_key(ThemeKey::ThinkingText, s)
82                });
83                let default_style = DefaultTextStyle {
84                    color: Some(color_fn),
85                    bold: false,
86                    italic: true,
87                    strikethrough: false,
88                    underline: false,
89                };
90                let mut md = Markdown::new(
91                    block.text.clone(),
92                    1,
93                    0,
94                    crate::agent::ui::theme::get_markdown_theme(),
95                    Some(default_style),
96                );
97                // Mark dirty since text may have already changed by the time
98                // this instance is created (the block was pushed earlier).
99                md.invalidate();
100                self.thinking_md.push(md);
101            } else {
102                let needs_update = !self.thinking_md[i].cached_text_matches(&block.text);
103                if needs_update {
104                    self.thinking_md[i].set_text(&block.text);
105                }
106            }
107        }
108    }
109
110    /// Compute a hash of the current state for cache_key.
111    fn state_hash(&self) -> u64 {
112        use std::collections::hash_map::DefaultHasher;
113        use std::hash::{Hash, Hasher};
114        let mut hasher = DefaultHasher::new();
115        self.text.hash(&mut hasher);
116        self.hide_thinking.hash(&mut hasher);
117        for block in &self.thinking {
118            block.text.hash(&mut hasher);
119            block.level.hash(&mut hasher);
120        }
121        hasher.finish()
122    }
123
124    pub fn add_thinking(&mut self, text: impl Into<String>, level: Option<String>) {
125        let text = text.into();
126        if text.is_empty() {
127            return;
128        }
129        if let Some(last) = self.thinking.last_mut() {
130            // Skip exact duplicates (same content sent again by the provider).
131            if text == last.text {
132                return;
133            }
134            // Some providers (e.g. Ollama) send FULL accumulated content in each
135            // chunk instead of a delta. Detect this: if the new text is longer,
136            // and its trimmed version starts with the trimmed existing text,
137            // it's a full accumulation - replace instead of append.
138            if text.len() > last.text.len() {
139                let t_trimmed = text.trim_start();
140                let l_trimmed = last.text.trim_start();
141                if t_trimmed.starts_with(l_trimmed) {
142                    last.text = text;
143                    self.invalidate();
144                    return;
145                }
146            }
147            // Default: treat as delta and append
148            last.text.push_str(&text);
149        } else {
150            self.thinking.push(ThinkingBlock { text, level });
151        }
152        self.invalidate();
153    }
154
155    pub fn append_text(&mut self, delta: &str) {
156        if delta.is_empty() {
157            return;
158        }
159        // Some providers (e.g. Ollama) send FULL accumulated content in each chunk
160        // instead of a delta. Detect this: if delta is longer than existing text,
161        // and its trimmed version starts with the trimmed existing text,
162        // replace instead of append.
163        if delta.len() > self.text.len() {
164            let d_trimmed = delta.trim_start();
165            let s_trimmed = self.text.trim_start();
166            if delta == self.text {
167                return; // Skip exact duplicate
168            }
169            if d_trimmed.starts_with(s_trimmed) {
170                self.text = delta.to_string();
171                self.invalidate();
172                return;
173            }
174        } else if delta == self.text {
175            return; // Skip exact duplicate
176        }
177        self.text.push_str(delta);
178        self.invalidate();
179    }
180
181    pub fn set_text(&mut self, text: impl Into<String>) {
182        self.text = text.into();
183        self.invalidate();
184    }
185}
186
187impl Component for AssistantMessageComponent {
188    // No-op: expand/collapse is controlled separately via set_hide_thinking.
189    // Pi keeps app.thinking.toggle (Ctrl+T) and app.tools.expand (Ctrl+O)
190    // as independent concerns - tool expansion must not affect thinking visibility.
191
192    fn render(&mut self, width: usize) -> Vec<String> {
193        if self.cached_width == width
194            && let Some(ref lines) = self.cached_lines
195        {
196            return lines.clone();
197        }
198
199        // Sync persistent Markdown instances before rendering so they
200        // can leverage their own text-based caching.
201        self.sync_text_md();
202        if !self.thinking.is_empty() {
203            self.sync_thinking_md();
204        }
205
206        let mut lines: Vec<String> = Vec::new();
207
208        let has_thinking = self.thinking.iter().any(|b| !b.text.trim().is_empty());
209        let has_text = !self.text.trim().is_empty();
210        if !has_thinking && !has_text {
211            self.cached_lines = Some(Vec::new());
212            self.cached_width = width;
213            return Vec::new();
214        }
215
216        // Leading blank line before any content (matching pi's leading Spacer(1))
217        lines.push(String::new());
218
219        let mut think_idx = 0;
220        for (block_idx, block) in self.thinking.iter().enumerate() {
221            let trimmed = block.text.trim();
222            if trimmed.is_empty() {
223                continue;
224            }
225            if self.hide_thinking {
226                let theme = current_theme();
227                let label = theme.italic(&theme.fg_key(ThemeKey::ThinkingText, "Thinking..."));
228                let padded = format!(" {} ", label);
229                lines.push(padded);
230            } else {
231                // Use persistent Markdown instance for this thinking block
232                if let Some(md) = self.thinking_md.get_mut(think_idx) {
233                    lines.extend(md.render(width));
234                }
235                think_idx += 1;
236            }
237
238            // Pi: after each block, add Spacer(1) if there's visible content after
239            let has_content_after = self.thinking.iter().skip(block_idx + 1).any(|b| {
240                if self.hide_thinking {
241                    true
242                } else {
243                    !b.text.trim().is_empty()
244                }
245            }) || has_text;
246            if has_content_after {
247                lines.push(String::new());
248            }
249        }
250
251        if has_text {
252            // Use persistent text Markdown instance
253            if let Some(ref mut md) = self.text_md {
254                lines.extend(md.render(width));
255            }
256        }
257
258        // Add OSC133 zones around the entire component (matching pi)
259        if !lines.is_empty() {
260            lines[0] = format!("{}{}", OSC133_ZONE_START, &lines[0]);
261            if let Some(last) = lines.last_mut() {
262                last.push_str(OSC133_ZONE_END);
263                last.push_str(OSC133_ZONE_FINAL);
264            }
265        }
266
267        let result = lines.clone();
268        self.cached_lines = Some(lines);
269        self.cached_width = width;
270        result
271    }
272
273    fn cache_key(&self, width: usize) -> Option<RenderCacheKey> {
274        Some(RenderCacheKey {
275            width,
276            expanded: false,
277            state_hash: self.state_hash(),
278        })
279    }
280
281    fn set_hide_thinking(&mut self, hide: bool) {
282        self.hide_thinking = hide;
283        self.invalidate();
284    }
285
286    fn invalidate(&mut self) {
287        self.cached_lines = None;
288    }
289}
rab/agent/ui/components/assistant_message.rs

rab/agent/ui/components/
assistant_message.rs
