Skip to main content

wisp/conversation/
items.rs

1use super::TurnState;
2use super::plan_tracker::PlanTracker;
3use super::progress_indicator::ProgressIndicator;
4use super::tool_calls::{ToolCall, ToolStatus, raw_input_fragment};
5use crate::view::markdown::{FenceLine, complete_lines_with_fences};
6use acp_utils::notifications::SubAgentProgressParams;
7use agent_client_protocol::schema::v1 as acp;
8use std::collections::HashMap;
9use std::sync::atomic::{AtomicU64, Ordering};
10
11static NEXT_CONVERSATION_ID: AtomicU64 = AtomicU64::new(1);
12
13#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
14pub struct ConversationId(u64);
15
16#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
17pub struct ConversationItemId(u64);
18
19#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
20pub struct Revision(u64);
21
22/// `Sealed` marks an item whose rendering is final: it may enter the
23/// terminal's native scrollback, which can never be rewritten. An `Open` item
24/// may still redraw in place and must stay in the live viewport.
25#[derive(Debug, Clone, Copy, PartialEq, Eq)]
26pub enum ItemState {
27    Open,
28    Sealed,
29}
30
31#[derive(Debug, Clone, PartialEq, Eq)]
32pub struct TextItem {
33    pub text: String,
34}
35
36#[derive(Debug, Clone, PartialEq, Eq)]
37pub struct Notice {
38    pub text: String,
39}
40
41#[derive(Debug, Clone, PartialEq, Eq)]
42pub enum ConversationContent {
43    User(TextItem),
44    Assistant(TextItem),
45    Tool(ToolCall),
46    Notice(Notice),
47}
48
49#[derive(Debug, Clone, PartialEq, Eq)]
50pub struct ConversationItem {
51    id: ConversationItemId,
52    revision: Revision,
53    state: ItemState,
54    content: ConversationContent,
55}
56
57impl ConversationItem {
58    pub fn id(&self) -> ConversationItemId {
59        self.id
60    }
61
62    pub fn revision(&self) -> Revision {
63        self.revision
64    }
65
66    pub fn state(&self) -> ItemState {
67        self.state
68    }
69
70    pub fn content(&self) -> &ConversationContent {
71        &self.content
72    }
73
74    pub fn text(&self) -> Option<&str> {
75        match &self.content {
76            ConversationContent::User(item) | ConversationContent::Assistant(item) => Some(&item.text),
77            ConversationContent::Notice(notice) => Some(&notice.text),
78            ConversationContent::Tool(_) => None,
79        }
80    }
81
82    pub fn is_open(&self) -> bool {
83        self.state == ItemState::Open
84    }
85}
86
87#[derive(Debug)]
88pub struct Conversation {
89    id: ConversationId,
90    items: Vec<ConversationItem>,
91    tool_index: HashMap<String, usize>,
92    next_item_id: u64,
93    turn: TurnState,
94    plan_tracker: PlanTracker,
95    progress_indicator: ProgressIndicator,
96}
97
98impl Default for Conversation {
99    fn default() -> Self {
100        Self::new()
101    }
102}
103
104impl Conversation {
105    pub fn new() -> Self {
106        Self {
107            id: ConversationId(NEXT_CONVERSATION_ID.fetch_add(1, Ordering::Relaxed)),
108            items: Vec::new(),
109            tool_index: HashMap::new(),
110            next_item_id: 0,
111            turn: TurnState::default(),
112            plan_tracker: PlanTracker::default(),
113            progress_indicator: ProgressIndicator::default(),
114        }
115    }
116
117    pub fn id(&self) -> ConversationId {
118        self.id
119    }
120
121    pub fn items(&self) -> &[ConversationItem] {
122        &self.items
123    }
124
125    pub fn append_user_content(&mut self, text: impl Into<String>) -> ConversationItemId {
126        self.seal_open_assistant();
127        self.push(ItemState::Sealed, ConversationContent::User(TextItem { text: text.into() }))
128    }
129
130    pub fn append_notice(&mut self, text: impl Into<String>) -> ConversationItemId {
131        self.seal_open_assistant();
132        self.push(ItemState::Sealed, ConversationContent::Notice(Notice { text: text.into() }))
133    }
134
135    pub fn append_assistant_chunk(&mut self, chunk: &str) {
136        if chunk.is_empty() {
137            return;
138        }
139        let has_open_assistant = self.items.last().is_some_and(|item| {
140            item.state == ItemState::Open && matches!(item.content, ConversationContent::Assistant(_))
141        });
142        if !has_open_assistant {
143            self.push(ItemState::Open, ConversationContent::Assistant(TextItem { text: String::new() }));
144        }
145        if let Some(item) = self.items.last_mut()
146            && let ConversationContent::Assistant(text) = &mut item.content
147        {
148            text.text.push_str(chunk);
149            item.revision.bump();
150        }
151        while let Some(finalized_end) = self.items.last().and_then(|item| match &item.content {
152            ConversationContent::Assistant(text) => complete_lines_with_fences(&text.text)
153                .find(|(_, line)| matches!(line, FenceLine::Blank))
154                .map(|(offset, _)| offset),
155            _ => None,
156        }) {
157            let trailing = match self.items.last_mut() {
158                Some(ConversationItem { content: ConversationContent::Assistant(text), state, revision, .. }) => {
159                    let trailing = text.text.split_off(finalized_end);
160                    *state = ItemState::Sealed;
161                    revision.bump();
162                    trailing
163                }
164                _ => break,
165            };
166            if trailing.is_empty() {
167                break;
168            }
169            self.push(ItemState::Open, ConversationContent::Assistant(TextItem { text: trailing }));
170        }
171    }
172
173    pub fn finish_current_block(&mut self) {
174        self.seal_open_assistant();
175    }
176
177    pub fn on_tool_call(&mut self, tool_call: &acp::ToolCall) {
178        self.seal_open_assistant();
179        let id = tool_call.tool_call_id.0.to_string();
180        if let Some(&index) = self.tool_index.get(&id) {
181            if self.items[index].state == ItemState::Open
182                && let ConversationContent::Tool(current) = &mut self.items[index].content
183            {
184                if !tool_call.title.is_empty() {
185                    current.title.clone_from(&tool_call.title);
186                }
187                current.status = ToolStatus::Running;
188                current.raw_input = tool_call.raw_input.as_ref().map_or_else(String::new, raw_input_fragment);
189                self.items[index].revision.bump();
190            }
191            return;
192        }
193        let index = self.items.len();
194        self.tool_index.insert(id, index);
195        let item_id = self.next_id();
196        self.items.push(ConversationItem {
197            id: item_id,
198            revision: Revision(0),
199            state: ItemState::Open,
200            content: ConversationContent::Tool(ToolCall::from_acp(tool_call)),
201        });
202    }
203
204    pub fn on_tool_call_update(&mut self, update: &acp::ToolCallUpdate) {
205        let Some(&index) = self.tool_index.get(update.tool_call_id.0.as_ref()) else {
206            return;
207        };
208        self.update_open_tool(index, |tool_call| tool_call.apply_update(update));
209    }
210
211    pub fn on_sub_agent_progress(&mut self, notification: &SubAgentProgressParams) {
212        let Some(&index) = self.tool_index.get(&notification.parent_tool_id) else {
213            return;
214        };
215        self.update_open_tool(index, |tool_call| tool_call.apply_sub_agent_progress(notification));
216    }
217
218    pub fn finish_turn(&mut self, terminal_status: &ToolStatus) {
219        for item in &mut self.items {
220            if item.state != ItemState::Open {
221                continue;
222            }
223            if let ConversationContent::Tool(tool_call) = &mut item.content {
224                tool_call.finalize(terminal_status);
225            }
226            item.state = ItemState::Sealed;
227            item.revision.bump();
228        }
229    }
230
231    pub fn clear(&mut self) {
232        self.id = ConversationId(NEXT_CONVERSATION_ID.fetch_add(1, Ordering::Relaxed));
233        self.items.clear();
234        self.tool_index.clear();
235        self.next_item_id = 0;
236    }
237
238    pub fn turn(&self) -> &TurnState {
239        &self.turn
240    }
241
242    pub fn turn_mut(&mut self) -> &mut TurnState {
243        &mut self.turn
244    }
245
246    pub fn plan_tracker(&self) -> &PlanTracker {
247        &self.plan_tracker
248    }
249
250    pub fn plan_tracker_mut(&mut self) -> &mut PlanTracker {
251        &mut self.plan_tracker
252    }
253
254    pub fn progress_indicator(&self) -> &ProgressIndicator {
255        &self.progress_indicator
256    }
257
258    pub fn progress_indicator_mut(&mut self) -> &mut ProgressIndicator {
259        &mut self.progress_indicator
260    }
261
262    pub fn reset_feature_state(&mut self) {
263        self.turn.reset();
264        self.plan_tracker.clear();
265        self.progress_indicator = ProgressIndicator::default();
266    }
267
268    pub fn any_running(&self) -> bool {
269        self.items.iter().any(|item| match &item.content {
270            ConversationContent::Tool(tool_call) => tool_call.is_running(),
271            _ => false,
272        })
273    }
274
275    fn push(&mut self, state: ItemState, content: ConversationContent) -> ConversationItemId {
276        let id = self.next_id();
277        self.items.push(ConversationItem { id, revision: Revision(0), state, content });
278        id
279    }
280
281    fn next_id(&mut self) -> ConversationItemId {
282        let id = ConversationItemId(self.next_item_id);
283        self.next_item_id = self.next_item_id.saturating_add(1);
284        id
285    }
286
287    fn seal_open_assistant(&mut self) {
288        if let Some(item) = self.items.last_mut()
289            && item.state == ItemState::Open
290            && matches!(item.content, ConversationContent::Assistant(_))
291        {
292            item.state = ItemState::Sealed;
293            item.revision.bump();
294        }
295    }
296
297    fn update_open_tool(&mut self, index: usize, apply: impl FnOnce(&mut ToolCall)) {
298        let item = &mut self.items[index];
299        if item.state == ItemState::Open
300            && let ConversationContent::Tool(tool_call) = &mut item.content
301        {
302            apply(tool_call);
303            if tool_call.rendering_final() {
304                item.state = ItemState::Sealed;
305            }
306            item.revision.bump();
307        }
308    }
309}
310
311impl Revision {
312    fn bump(&mut self) {
313        self.0 = self.0.saturating_add(1);
314    }
315}