Skip to main content

swink_agent/
display.rs

1//! Display-ready message types for frontend consumption.
2//!
3//! Provides a core display representation that any frontend (TUI, GUI, web)
4//! can wrap with UI-specific fields (collapse state, scroll position, etc.).
5
6use crate::types::{ContentBlock, LlmMessage, StopReason};
7
8/// Role of a message for display styling.
9#[derive(Debug, Clone, Copy, PartialEq, Eq)]
10pub enum DisplayRole {
11    User,
12    Assistant,
13    ToolResult,
14    Error,
15    System,
16}
17
18/// A message converted to a frontend-friendly format.
19///
20/// Contains the essential display data extracted from [`LlmMessage`].
21/// Frontend implementations can wrap this with additional UI-specific
22/// fields (collapse state, scroll position, etc.).
23#[derive(Debug, Clone)]
24pub struct CoreDisplayMessage {
25    pub role: DisplayRole,
26    pub content: String,
27    pub thinking: Option<String>,
28}
29
30/// Convert message types into display-ready representations.
31pub trait IntoDisplayMessages {
32    fn to_display_messages(&self) -> Vec<CoreDisplayMessage>;
33}
34
35impl IntoDisplayMessages for LlmMessage {
36    fn to_display_messages(&self) -> Vec<CoreDisplayMessage> {
37        match self {
38            Self::User(user) => {
39                vec![CoreDisplayMessage {
40                    role: DisplayRole::User,
41                    content: ContentBlock::extract_text(&user.content),
42                    thinking: None,
43                }]
44            }
45            Self::Assistant(assistant) => {
46                let mut text_parts = Vec::new();
47                let mut thinking_parts = Vec::new();
48                for block in &assistant.content {
49                    match block {
50                        ContentBlock::Text { text } => text_parts.push(text.as_str()),
51                        ContentBlock::Thinking { thinking, .. } => {
52                            thinking_parts.push(thinking.as_str());
53                        }
54                        _ => {}
55                    }
56                }
57
58                let content = if !text_parts.is_empty() {
59                    text_parts.join("")
60                } else if assistant.stop_reason == StopReason::Error {
61                    assistant.error_message.clone().unwrap_or_default()
62                } else {
63                    String::new()
64                };
65
66                let thinking = if thinking_parts.is_empty() {
67                    None
68                } else {
69                    Some(thinking_parts.join(""))
70                };
71
72                let role = if assistant.stop_reason == StopReason::Error {
73                    DisplayRole::Error
74                } else {
75                    DisplayRole::Assistant
76                };
77
78                vec![CoreDisplayMessage {
79                    role,
80                    content,
81                    thinking,
82                }]
83            }
84            Self::ToolResult(tool_result) => {
85                let content = ContentBlock::extract_text(&tool_result.content);
86                if content.is_empty() {
87                    return vec![];
88                }
89                let role = if tool_result.is_error {
90                    DisplayRole::Error
91                } else {
92                    DisplayRole::ToolResult
93                };
94                vec![CoreDisplayMessage {
95                    role,
96                    content,
97                    thinking: None,
98                }]
99            }
100        }
101    }
102}
103
104impl IntoDisplayMessages for [LlmMessage] {
105    fn to_display_messages(&self) -> Vec<CoreDisplayMessage> {
106        self.iter()
107            .flat_map(IntoDisplayMessages::to_display_messages)
108            .collect()
109    }
110}
111
112impl IntoDisplayMessages for Vec<LlmMessage> {
113    fn to_display_messages(&self) -> Vec<CoreDisplayMessage> {
114        self.as_slice().to_display_messages()
115    }
116}
117
118#[cfg(test)]
119mod tests {
120    use super::*;
121    use crate::types::{AssistantMessage, Cost, Usage, UserMessage};
122
123    #[test]
124    fn user_message_to_display() {
125        let msg = LlmMessage::User(UserMessage {
126            content: vec![ContentBlock::Text {
127                text: "hello".to_string(),
128            }],
129            timestamp: 0,
130            cache_hint: None,
131        });
132        let display = msg.to_display_messages();
133        assert_eq!(display.len(), 1);
134        assert_eq!(display[0].role, DisplayRole::User);
135        assert_eq!(display[0].content, "hello");
136        assert!(display[0].thinking.is_none());
137    }
138
139    #[test]
140    fn assistant_message_with_thinking() {
141        let msg = LlmMessage::Assistant(AssistantMessage {
142            content: vec![
143                ContentBlock::Thinking {
144                    thinking: "reasoning".to_string(),
145                    signature: None,
146                },
147                ContentBlock::Text {
148                    text: "answer".to_string(),
149                },
150            ],
151            provider: String::new(),
152            model_id: String::new(),
153            usage: Usage::default(),
154            cost: Cost::default(),
155            stop_reason: StopReason::Stop,
156            error_message: None,
157            error_kind: None,
158            timestamp: 0,
159            cache_hint: None,
160        });
161        let display = msg.to_display_messages();
162        assert_eq!(display.len(), 1);
163        assert_eq!(display[0].role, DisplayRole::Assistant);
164        assert_eq!(display[0].content, "answer");
165        assert_eq!(display[0].thinking.as_deref(), Some("reasoning"));
166    }
167
168    #[test]
169    fn assistant_error_message() {
170        let msg = LlmMessage::Assistant(AssistantMessage {
171            content: vec![],
172            provider: String::new(),
173            model_id: String::new(),
174            usage: Usage::default(),
175            cost: Cost::default(),
176            stop_reason: StopReason::Error,
177            error_message: Some("something broke".to_string()),
178            error_kind: None,
179            timestamp: 0,
180            cache_hint: None,
181        });
182        let display = msg.to_display_messages();
183        assert_eq!(display.len(), 1);
184        assert_eq!(display[0].role, DisplayRole::Error);
185        assert_eq!(display[0].content, "something broke");
186    }
187
188    #[test]
189    fn empty_tool_result_produces_no_messages() {
190        let msg = LlmMessage::ToolResult(crate::types::ToolResultMessage {
191            tool_call_id: "tc1".to_string(),
192            content: vec![],
193            is_error: false,
194            timestamp: 0,
195            details: serde_json::Value::Null,
196            cache_hint: None,
197        });
198        let display = msg.to_display_messages();
199        assert!(display.is_empty());
200    }
201}