Skip to main content

wisp/components/
conversation_window.rs

1use std::mem::{Discriminant, discriminant};
2
3use crate::components::input_prompt::prompt_text_start_col;
4use crate::components::thought_message::ThoughtMessage;
5use crate::components::tool_call_statuses::ToolCallStatuses;
6use tui::{Line, Style, ViewContext, render_markdown};
7
8#[derive(Debug, Clone)]
9pub enum SegmentContent {
10    UserMessage(String),
11    Text(String),
12    Thought(String),
13    ToolCall(String),
14}
15
16#[derive(Debug)]
17struct Segment {
18    content: SegmentContent,
19}
20
21#[doc = include_str!("../docs/conversation_window.md")]
22pub struct ConversationBuffer {
23    segments: Vec<Segment>,
24    thought_block_open: bool,
25}
26
27impl Default for ConversationBuffer {
28    fn default() -> Self {
29        Self::new()
30    }
31}
32
33impl ConversationBuffer {
34    pub fn new() -> Self {
35        Self { segments: Vec::new(), thought_block_open: false }
36    }
37
38    #[cfg(test)]
39    pub(crate) fn segments(&self) -> impl ExactSizeIterator<Item = &SegmentContent> {
40        self.segments.iter().map(|s| &s.content)
41    }
42
43    pub fn push_user_message(&mut self, text: &str) {
44        self.close_thought_block();
45        self.segments.push(Segment { content: SegmentContent::UserMessage(text.to_string()) });
46    }
47
48    pub fn append_text_chunk(&mut self, chunk: &str) {
49        if chunk.is_empty() {
50            return;
51        }
52
53        self.close_thought_block();
54
55        if let Some(segment) = self.segments.last_mut()
56            && let SegmentContent::Text(existing) = &mut segment.content
57        {
58            existing.push_str(chunk);
59        } else {
60            self.segments.push(Segment { content: SegmentContent::Text(chunk.to_string()) });
61        }
62    }
63
64    pub fn append_thought_chunk(&mut self, chunk: &str) {
65        if chunk.is_empty() {
66            return;
67        }
68
69        if self.thought_block_open
70            && let Some(segment) = self.segments.last_mut()
71            && let SegmentContent::Thought(existing) = &mut segment.content
72        {
73            existing.push_str(chunk);
74            return;
75        }
76
77        self.segments.push(Segment { content: SegmentContent::Thought(chunk.to_string()) });
78        self.thought_block_open = true;
79    }
80
81    pub(crate) fn close_thought_block(&mut self) {
82        self.thought_block_open = false;
83    }
84
85    pub(crate) fn clear(&mut self) {
86        self.segments.clear();
87        self.thought_block_open = false;
88    }
89
90    pub(crate) fn ensure_tool_segment(&mut self, tool_id: &str) {
91        let has_segment =
92            self.segments.iter().any(|s| matches!(&s.content, SegmentContent::ToolCall(id) if id == tool_id));
93
94        if !has_segment {
95            self.segments.push(Segment { content: SegmentContent::ToolCall(tool_id.to_string()) });
96        }
97    }
98
99    #[cfg(test)]
100    fn drain_segments_except(&mut self, mut keep: impl FnMut(&SegmentContent) -> bool) -> Vec<Segment> {
101        let old = std::mem::take(&mut self.segments);
102        let (kept, removed) = old.into_iter().partition(|s| keep(&s.content));
103        self.segments = kept;
104        removed
105    }
106
107    #[cfg(test)]
108    pub(crate) fn drain_completed(
109        &mut self,
110        tool_call_statuses: &ToolCallStatuses,
111    ) -> (Vec<SegmentContent>, Vec<String>) {
112        let drained = self.drain_segments_except(
113            |seg| matches!(seg, SegmentContent::ToolCall(id) if tool_call_statuses.is_tool_running(id)),
114        );
115
116        let mut content = Vec::new();
117        let mut completed_tool_ids = Vec::new();
118
119        for segment in drained {
120            if let SegmentContent::ToolCall(ref id) = segment.content {
121                completed_tool_ids.push(id.clone());
122            }
123            content.push(segment.content);
124        }
125
126        (content, completed_tool_ids)
127    }
128}
129
130pub struct ConversationWindow<'a> {
131    pub conversation: &'a ConversationBuffer,
132    pub tool_call_statuses: &'a ToolCallStatuses,
133    pub content_padding: usize,
134}
135
136impl ConversationWindow<'_> {
137    pub fn render(&self, context: &ViewContext) -> Vec<Line> {
138        let mut lines = Vec::new();
139        let mut last_segment_kind = None;
140
141        for segment in &self.conversation.segments {
142            let kind = discriminant(&segment.content);
143            let mut rendered = render_stream_segment(&segment.content, self.tool_call_statuses, context);
144            if !matches!(segment.content, SegmentContent::UserMessage(_)) {
145                pad_lines(&mut rendered, self.content_padding);
146            }
147            extend_with_vertical_margin(&mut lines, &mut last_segment_kind, kind, &rendered);
148        }
149
150        lines
151    }
152}
153
154fn render_stream_segment(
155    segment: &SegmentContent,
156    tool_call_statuses: &ToolCallStatuses,
157    context: &ViewContext,
158) -> Vec<Line> {
159    match segment {
160        SegmentContent::UserMessage(text) => render_user_message_block(text, context),
161        SegmentContent::Thought(text) => ThoughtMessage { text }.render(context),
162        SegmentContent::Text(text) => render_markdown(text, context),
163        SegmentContent::ToolCall(id) => tool_call_statuses.render_tool(id, context),
164    }
165}
166
167fn render_user_message_block(text: &str, context: &ViewContext) -> Vec<Line> {
168    if text.is_empty() {
169        return vec![];
170    }
171
172    let block_style = Style::fg(context.theme.text_primary()).bg_color(context.theme.sidebar_bg());
173    let block_width = usize::from(context.size.width).max(1);
174    let left_padding = prompt_text_start_col(block_width).min(block_width.saturating_sub(1));
175    let mut rendered_lines = Vec::new();
176    rendered_lines.push(padded_background_line(block_width, block_style));
177
178    for content in text.lines() {
179        rendered_lines.extend(render_user_message_lines(content, left_padding, block_width, block_style));
180    }
181
182    rendered_lines.push(padded_background_line(block_width, block_style));
183    rendered_lines
184}
185
186fn render_user_message_lines(content: &str, left_padding: usize, block_width: usize, block_style: Style) -> Vec<Line> {
187    if content.is_empty() {
188        return vec![padded_background_line(block_width, block_style)];
189    }
190
191    let content_width = block_width.saturating_sub(left_padding).max(1);
192    Line::with_style(content.to_string(), block_style)
193        .soft_wrap(u16::try_from(content_width).unwrap_or(u16::MAX))
194        .into_iter()
195        .map(|line| pad_user_message_line(&line, left_padding, block_width, block_style))
196        .collect()
197}
198
199fn pad_user_message_line(line: &Line, left_padding: usize, block_width: usize, block_style: Style) -> Line {
200    let mut padded_line = Line::with_style(" ".repeat(left_padding), block_style);
201    padded_line.append_line(line);
202
203    let trailing_padding = block_width.saturating_sub(padded_line.display_width());
204    if trailing_padding > 0 {
205        padded_line.push_with_style(" ".repeat(trailing_padding), block_style);
206    }
207
208    padded_line
209}
210
211fn padded_background_line(width: usize, style: Style) -> Line {
212    Line::with_style(" ".repeat(width.max(1)), style)
213}
214
215pub(crate) fn pad_lines(lines: &mut [Line], padding: usize) {
216    if padding == 0 {
217        return;
218    }
219    let prefix = " ".repeat(padding);
220    for line in lines.iter_mut() {
221        *line = std::mem::take(line).prepend(&prefix);
222    }
223}
224
225fn extend_with_vertical_margin(
226    target: &mut Vec<Line>,
227    last_segment_kind: &mut Option<Discriminant<SegmentContent>>,
228    kind: Discriminant<SegmentContent>,
229    lines: &[Line],
230) {
231    if lines.is_empty() {
232        return;
233    }
234
235    if let Some(prev_kind) = *last_segment_kind
236        && prev_kind != kind
237    {
238        target.push(Line::new(String::new()));
239    }
240
241    target.extend_from_slice(lines);
242    *last_segment_kind = Some(kind);
243}
244
245#[cfg(test)]
246mod tests {
247    use super::*;
248    use crate::settings::DEFAULT_CONTENT_PADDING;
249
250    #[test]
251    fn buffer_closes_thought_block_when_text_arrives() {
252        let mut buffer = ConversationBuffer::new();
253        buffer.append_thought_chunk("thinking");
254        buffer.append_text_chunk("answer");
255        buffer.append_thought_chunk("new thought");
256
257        let segments: Vec<_> = buffer.segments().collect();
258        assert_eq!(segments.len(), 3);
259        assert!(matches!(segments[0], SegmentContent::Thought(_)));
260        assert!(matches!(segments[1], SegmentContent::Text(_)));
261        assert!(matches!(segments[2], SegmentContent::Thought(_)));
262    }
263
264    #[test]
265    fn buffer_coalesces_contiguous_thought_chunks() {
266        let mut buffer = ConversationBuffer::new();
267        buffer.append_thought_chunk("a");
268        buffer.append_thought_chunk("b");
269
270        let segments: Vec<_> = buffer.segments().collect();
271        assert_eq!(segments.len(), 1);
272        match segments[0] {
273            SegmentContent::Thought(text) => assert_eq!(text, "ab"),
274            _ => panic!("expected thought segment"),
275        }
276    }
277
278    #[test]
279    fn clear_removes_segments_and_resets_state() {
280        let mut buffer = ConversationBuffer::new();
281        buffer.append_thought_chunk("thinking");
282        buffer.append_text_chunk("answer");
283        assert_eq!(buffer.segments().len(), 2);
284
285        buffer.clear();
286
287        assert_eq!(buffer.segments().len(), 0);
288        buffer.append_thought_chunk("new");
289        assert_eq!(buffer.segments().len(), 1);
290    }
291
292    #[test]
293    fn user_message_renders_with_top_and_bottom_padding_lines() {
294        let mut buffer = ConversationBuffer::new();
295        buffer.push_user_message("hello");
296
297        let tool_call_statuses = ToolCallStatuses::new();
298        let window = ConversationWindow {
299            conversation: &buffer,
300            tool_call_statuses: &tool_call_statuses,
301            content_padding: DEFAULT_CONTENT_PADDING,
302        };
303        let context = ViewContext::new((80, 24));
304
305        let lines = window.render(&context);
306
307        assert_eq!(lines.len(), 3);
308        let left_padding = " ".repeat(prompt_text_start_col(usize::from(context.size.width)));
309        assert_eq!(lines[1].plain_text().trim_end(), format!("{left_padding}hello"));
310        assert!(lines[0].plain_text().trim().is_empty());
311        assert!(lines[2].plain_text().trim().is_empty());
312        assert_user_message_style(&lines[0], &context);
313        assert_user_message_style(&lines[1], &context);
314        assert_user_message_style(&lines[2], &context);
315        assert!(lines.iter().all(|line| line.display_width() == usize::from(context.size.width)));
316    }
317
318    #[test]
319    fn user_message_block_applies_theme_bg_to_all_lines() {
320        let mut buffer = ConversationBuffer::new();
321        buffer.push_user_message("line one\n\nline three");
322
323        let tool_call_statuses = ToolCallStatuses::new();
324        let window = ConversationWindow {
325            conversation: &buffer,
326            tool_call_statuses: &tool_call_statuses,
327            content_padding: DEFAULT_CONTENT_PADDING,
328        };
329        let context = ViewContext::new((80, 24));
330
331        let lines = window.render(&context);
332
333        assert_eq!(lines.len(), 5);
334        let left_padding = " ".repeat(prompt_text_start_col(usize::from(context.size.width)));
335        assert_eq!(lines[1].plain_text().trim_end(), format!("{left_padding}line one"));
336        assert!(lines[2].plain_text().trim().is_empty());
337        assert_eq!(lines[3].plain_text().trim_end(), format!("{left_padding}line three"));
338
339        for line in &lines {
340            assert_user_message_style(line, &context);
341        }
342
343        let first_width = lines[0].display_width();
344        assert_eq!(first_width, usize::from(context.size.width));
345        assert!(lines.iter().all(|line| line.display_width() == first_width));
346    }
347
348    #[test]
349    fn user_message_wrapped_rows_keep_full_width_background() {
350        let mut buffer = ConversationBuffer::new();
351        buffer.push_user_message("0123456789");
352
353        let tool_call_statuses = ToolCallStatuses::new();
354        let window = ConversationWindow {
355            conversation: &buffer,
356            tool_call_statuses: &tool_call_statuses,
357            content_padding: DEFAULT_CONTENT_PADDING,
358        };
359        let context = ViewContext::new((8, 24));
360
361        let lines = window.render(&context);
362
363        assert_eq!(lines.len(), 5);
364        assert_eq!(lines[1].plain_text().trim_end(), "    0123");
365        assert_eq!(lines[2].plain_text().trim_end(), "    4567");
366        assert_eq!(lines[3].plain_text().trim_end(), "    89");
367        assert!(lines.iter().all(|line| line.display_width() == usize::from(context.size.width)));
368        for line in &lines {
369            assert_user_message_style(line, &context);
370        }
371    }
372
373    #[test]
374    fn drain_completed_returns_content_and_tool_ids() {
375        use agent_client_protocol as acp;
376
377        let mut buffer = ConversationBuffer::new();
378        buffer.append_text_chunk("hello");
379        buffer.ensure_tool_segment("tool-1");
380
381        let mut statuses = ToolCallStatuses::new();
382        let tc = acp::ToolCall::new("tool-1", "Read file");
383        statuses.on_tool_call(&tc);
384        let update =
385            acp::ToolCallUpdate::new("tool-1", acp::ToolCallUpdateFields::new().status(acp::ToolCallStatus::Completed));
386        statuses.on_tool_call_update(&update);
387
388        let (content, tool_ids) = buffer.drain_completed(&statuses);
389
390        assert_eq!(content.len(), 2, "should have text and tool content");
391        assert!(matches!(content[0], SegmentContent::Text(_)));
392        assert!(matches!(content[1], SegmentContent::ToolCall(_)));
393        assert_eq!(tool_ids, vec!["tool-1"]);
394        assert_eq!(buffer.segments().len(), 0, "all segments should be drained");
395    }
396
397    #[test]
398    fn drain_completed_keeps_running_tools() {
399        use agent_client_protocol as acp;
400
401        let mut buffer = ConversationBuffer::new();
402        buffer.append_text_chunk("hello");
403        buffer.ensure_tool_segment("tool-1");
404
405        let mut statuses = ToolCallStatuses::new();
406        let tc = acp::ToolCall::new("tool-1", "Read file");
407        statuses.on_tool_call(&tc);
408        // tool-1 stays Running (no completion update)
409
410        let (content, tool_ids) = buffer.drain_completed(&statuses);
411
412        assert_eq!(content.len(), 1, "text segment should still be drained");
413        assert!(matches!(content[0], SegmentContent::Text(_)));
414        assert!(tool_ids.is_empty(), "running tool should not be drained");
415        let segments: Vec<_> = buffer.segments().collect();
416        assert_eq!(segments.len(), 1, "running tool should remain");
417        assert!(matches!(
418            segments[0],
419            SegmentContent::ToolCall(id) if id == "tool-1"
420        ));
421    }
422
423    fn assert_user_message_style(line: &Line, context: &ViewContext) {
424        assert!(!line.spans().is_empty());
425        assert!(line.spans().iter().all(|span| span.style().bg == Some(context.theme.sidebar_bg())));
426        assert!(line.spans().iter().all(|span| span.style().fg == Some(context.theme.text_primary())));
427    }
428}