Skip to main content

wisp/components/
conversation_window.rs

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