wisp/conversation/
items.rs

1use super::TurnState;
2use super::plan_tracker::PlanTracker;
3use super::progress_indicator::ProgressIndicator;
4use super::tool_calls::{ToolCall, ToolStatus};
5use acp_utils::content::{display_content_blocks, map_content_blocks_to_text};
6use acp_utils::notifications::SubAgentProgressParams;
7use agent_client_protocol::schema::{MaybeUndefined, v2 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
22impl Revision {
23    pub fn value(self) -> u64 {
24        self.0
25    }
26}
27
28#[derive(Debug, Clone, Copy, PartialEq, Eq)]
29pub enum ItemState {
30    Open,
31    Sealed,
32}
33
34#[derive(Debug, Clone, Copy, PartialEq, Eq)]
35pub enum MessageRole {
36    User,
37    Assistant,
38}
39
40#[derive(Debug, Clone, PartialEq, Eq)]
41pub struct TextItem {
42    pub text: String,
43}
44
45#[derive(Debug, Clone, PartialEq, Eq)]
46pub struct Notice {
47    pub text: String,
48}
49
50#[derive(Debug, Clone, PartialEq)]
51pub enum ConversationContent {
52    User(TextItem),
53    Assistant(TextItem),
54    Tool(ToolCall),
55    Notice(Notice),
56}
57
58#[derive(Debug, Clone, PartialEq)]
59pub struct ConversationItem {
60    id: ConversationItemId,
61    message_id: Option<acp::MessageId>,
62    preserve_user_display: bool,
63    revision: Revision,
64    replacement_revision: Revision,
65    state: ItemState,
66    content: ConversationContent,
67}
68
69impl ConversationItem {
70    pub fn id(&self) -> ConversationItemId {
71        self.id
72    }
73
74    pub fn message_id(&self) -> Option<&acp::MessageId> {
75        self.message_id.as_ref()
76    }
77
78    pub fn revision(&self) -> Revision {
79        self.revision
80    }
81
82    /// Advances when a change can invalidate rows already committed to native history.
83    pub fn replacement_revision(&self) -> Revision {
84        self.replacement_revision
85    }
86
87    pub fn state(&self) -> ItemState {
88        self.state
89    }
90
91    pub fn content(&self) -> &ConversationContent {
92        &self.content
93    }
94
95    pub fn text(&self) -> Option<&str> {
96        match &self.content {
97            ConversationContent::User(item) | ConversationContent::Assistant(item) => Some(&item.text),
98            ConversationContent::Notice(notice) => Some(&notice.text),
99            ConversationContent::Tool(_) => None,
100        }
101    }
102
103    pub fn is_open(&self) -> bool {
104        self.state == ItemState::Open
105    }
106}
107
108#[derive(Debug)]
109pub struct Conversation {
110    id: ConversationId,
111    items: Vec<ConversationItem>,
112    tool_index: HashMap<String, usize>,
113    message_index: HashMap<acp::MessageId, usize>,
114    pending_user: Option<usize>,
115    next_item_id: u64,
116    turn: TurnState,
117    plan_tracker: PlanTracker,
118    progress_indicator: ProgressIndicator,
119}
120
121impl Default for Conversation {
122    fn default() -> Self {
123        Self::new()
124    }
125}
126
127impl Conversation {
128    pub fn new() -> Self {
129        Self {
130            id: ConversationId(NEXT_CONVERSATION_ID.fetch_add(1, Ordering::Relaxed)),
131            items: Vec::new(),
132            tool_index: HashMap::new(),
133            message_index: HashMap::new(),
134            pending_user: None,
135            next_item_id: 0,
136            turn: TurnState::default(),
137            plan_tracker: PlanTracker::default(),
138            progress_indicator: ProgressIndicator::default(),
139        }
140    }
141
142    pub fn id(&self) -> ConversationId {
143        self.id
144    }
145
146    pub fn items(&self) -> &[ConversationItem] {
147        &self.items
148    }
149
150    pub fn append_user_content(&mut self, text: impl Into<String>) -> ConversationItemId {
151        self.push(ItemState::Sealed, ConversationContent::User(TextItem { text: text.into() }))
152    }
153
154    /// Echo a submitted prompt before the agent acknowledges it; the next
155    /// `user_message` upsert adopts this item instead of appending another.
156    pub fn append_pending_user_content(&mut self, text: impl Into<String>) -> ConversationItemId {
157        let id = self.push(ItemState::Open, ConversationContent::User(TextItem { text: text.into() }));
158        let index = self.items.len() - 1;
159        self.items[index].preserve_user_display = true;
160        self.pending_user = Some(index);
161        id
162    }
163
164    pub fn upsert_message(
165        &mut self,
166        role: MessageRole,
167        message_id: acp::MessageId,
168        content: &MaybeUndefined<Vec<acp::ContentBlock>>,
169    ) {
170        let index = self.message_slot(role, message_id);
171        if self.items[index].preserve_user_display {
172            return;
173        }
174        let text = match content {
175            MaybeUndefined::Undefined => return,
176            MaybeUndefined::Null => String::new(),
177            MaybeUndefined::Value(blocks) => message_display_text(role, blocks),
178        };
179        let content = message_content(role, text);
180        if self.items[index].content != content {
181            self.items[index].content = content;
182            self.items[index].changed(true);
183        }
184    }
185
186    pub fn append_message_chunk(&mut self, role: MessageRole, chunk: &acp::ContentChunk) {
187        let index = self.message_slot(role, chunk.message_id.clone());
188        let item = &mut self.items[index];
189        if item.preserve_user_display {
190            return;
191        }
192        let text = message_display_text(role, std::slice::from_ref(&chunk.content));
193        match &mut item.content {
194            ConversationContent::User(current) | ConversationContent::Assistant(current) => current.text.push_str(&text),
195            ConversationContent::Tool(_) | ConversationContent::Notice(_) => return,
196        }
197        if !text.is_empty() {
198            item.changed(!item.is_open() || role == MessageRole::User);
199        }
200    }
201
202    pub fn append_notice(&mut self, text: impl Into<String>) -> ConversationItemId {
203        self.push(ItemState::Sealed, ConversationContent::Notice(Notice { text: text.into() }))
204    }
205
206    pub fn on_tool_call_update(&mut self, update: &acp::ToolCallUpdate) {
207        let index = self.tool_slot(&update.tool_call_id);
208        self.update_tool(index, |tool_call| tool_call.apply_update(update));
209    }
210
211    pub fn on_tool_call_content_chunk(&mut self, chunk: &acp::ToolCallContentChunk) {
212        let index = self.tool_slot(&chunk.tool_call_id);
213        self.update_tool(index, |tool_call| tool_call.append_content(chunk.content.clone()));
214    }
215
216    pub fn on_sub_agent_progress(&mut self, notification: &SubAgentProgressParams) {
217        let Some(&index) = self.tool_index.get(&notification.parent_tool_id) else {
218            return;
219        };
220        self.update_tool(index, |tool_call| tool_call.apply_sub_agent_progress(notification));
221    }
222
223    pub fn finish_turn(&mut self, terminal_status: &ToolStatus) {
224        self.pending_user = None;
225        for item in &mut self.items {
226            if item.state != ItemState::Open {
227                continue;
228            }
229            if let ConversationContent::Tool(tool_call) = &mut item.content {
230                tool_call.finalize(terminal_status);
231            }
232            item.state = ItemState::Sealed;
233            item.changed(false);
234        }
235    }
236
237    pub fn clear(&mut self) {
238        self.id = ConversationId(NEXT_CONVERSATION_ID.fetch_add(1, Ordering::Relaxed));
239        self.items.clear();
240        self.tool_index.clear();
241        self.message_index.clear();
242        self.pending_user = None;
243        self.next_item_id = 0;
244    }
245
246    pub fn turn(&self) -> &TurnState {
247        &self.turn
248    }
249
250    pub fn turn_mut(&mut self) -> &mut TurnState {
251        &mut self.turn
252    }
253
254    pub fn plan_tracker(&self) -> &PlanTracker {
255        &self.plan_tracker
256    }
257
258    pub fn plan_tracker_mut(&mut self) -> &mut PlanTracker {
259        &mut self.plan_tracker
260    }
261
262    pub fn progress_indicator(&self) -> &ProgressIndicator {
263        &self.progress_indicator
264    }
265
266    pub fn progress_indicator_mut(&mut self) -> &mut ProgressIndicator {
267        &mut self.progress_indicator
268    }
269
270    pub fn reset_feature_state(&mut self) {
271        self.turn.reset();
272        self.plan_tracker.clear();
273        self.progress_indicator = ProgressIndicator::default();
274    }
275
276    pub fn any_running(&self) -> bool {
277        self.items.iter().any(|item| match &item.content {
278            ConversationContent::Tool(tool_call) => tool_call.is_running(),
279            _ => false,
280        })
281    }
282
283    fn message_slot(&mut self, role: MessageRole, message_id: acp::MessageId) -> usize {
284        if let Some(&index) = self.message_index.get(&message_id) {
285            return index;
286        }
287        let index = if role == MessageRole::User { self.pending_user.take() } else { None }.unwrap_or_else(|| {
288            let index = self.items.len();
289            self.push(ItemState::Open, message_content(role, String::new()));
290            index
291        });
292        self.items[index].message_id = Some(message_id.clone());
293        self.message_index.insert(message_id, index);
294        index
295    }
296
297    fn tool_slot(&mut self, id: &acp::ToolCallId) -> usize {
298        if let Some(&index) = self.tool_index.get(id.0.as_ref()) {
299            return index;
300        }
301        let index = self.items.len();
302        self.push(
303            ItemState::Open,
304            ConversationContent::Tool(ToolCall::from_update(&acp::ToolCallUpdate::new(id.clone()))),
305        );
306        self.tool_index.insert(id.to_string(), index);
307        index
308    }
309
310    fn push(&mut self, state: ItemState, content: ConversationContent) -> ConversationItemId {
311        let id = ConversationItemId(self.next_item_id);
312        self.next_item_id = self.next_item_id.saturating_add(1);
313        self.items.push(ConversationItem {
314            id,
315            message_id: None,
316            preserve_user_display: false,
317            revision: Revision(0),
318            replacement_revision: Revision(0),
319            state,
320            content,
321        });
322        id
323    }
324
325    fn update_tool(&mut self, index: usize, apply: impl FnOnce(&mut ToolCall)) {
326        let item = &mut self.items[index];
327        if let ConversationContent::Tool(tool_call) = &mut item.content {
328            let previous = tool_call.clone();
329            apply(tool_call);
330            if *tool_call == previous {
331                return;
332            }
333            let state = if tool_call.rendering_final() { ItemState::Sealed } else { ItemState::Open };
334            item.changed(!item.is_open());
335            item.state = state;
336        }
337    }
338}
339
340fn message_display_text(role: MessageRole, blocks: &[acp::ContentBlock]) -> String {
341    let blocks = match role {
342        MessageRole::User => display_content_blocks(blocks),
343        MessageRole::Assistant => blocks.to_vec(),
344    };
345    map_content_blocks_to_text(blocks)
346}
347
348fn message_content(role: MessageRole, text: String) -> ConversationContent {
349    match role {
350        MessageRole::User => ConversationContent::User(TextItem { text }),
351        MessageRole::Assistant => ConversationContent::Assistant(TextItem { text }),
352    }
353}
354
355impl ConversationItem {
356    fn changed(&mut self, replaces_history: bool) {
357        if replaces_history {
358            self.replacement_revision.bump();
359        }
360        self.revision.bump();
361    }
362}
363
364impl Revision {
365    fn bump(&mut self) {
366        self.0 = self.0.saturating_add(1);
367    }
368}
wisp/conversation/items.rs

wisp/conversation/
items.rs
