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