Skip to main content

acp_utils/conversation/
mod.rs

1mod activity;
2mod items;
3mod tool_calls;
4mod turn;
5
6pub use activity::Activity;
7pub use items::{ConversationContent, ConversationId, ConversationItem, ConversationItemId, ItemState, Revision};
8pub use tool_calls::{SubAgentState, ToolCall, ToolStatus};
9pub use turn::{TurnFinished, TurnPhase};
10
11use crate::client::AcpEvent;
12use crate::notifications::SubAgentProgressParams;
13use agent_client_protocol::schema::{MaybeUndefined, v2 as acp};
14use items::MessageRole;
15use std::collections::{HashMap, HashSet};
16
17/// One session's conversation items and state of the current turn.
18#[derive(Debug, Default)]
19pub struct Conversation {
20    id: ConversationId,
21    revision: Revision,
22    items: Vec<ConversationItem>,
23    tool_index: HashMap<String, usize>,
24    message_index: HashMap<acp::MessageId, usize>,
25    next_item_id: u64,
26    turn: TurnPhase,
27    turn_ended: bool,
28    activity: Activity,
29    compactions: HashSet<acp::CompactionId>,
30    context_usage: Option<acp::UsageUpdate>,
31    plan: Option<acp::PlanItems>,
32}
33
34impl Conversation {
35    pub fn new() -> Self {
36        Self::default()
37    }
38
39    pub fn apply_event(&mut self, event: &AcpEvent) -> Option<TurnFinished> {
40        match event {
41            AcpEvent::SessionUpdate(notification) => return self.apply_update(&notification.update),
42            AcpEvent::SubAgentProgress(progress) => self.apply_sub_agent_progress(progress),
43            AcpEvent::ContextCleared(_) => self.clear(),
44            AcpEvent::ConnectionClosed => self.connection_closed(),
45            AcpEvent::AuthMethodsUpdated(_)
46            | AcpEvent::McpNotification(_)
47            | AcpEvent::GitDiffEvent(_)
48            | AcpEvent::ElicitationRequest { .. } => {}
49        }
50        None
51    }
52
53    pub fn clear(&mut self) {
54        *self = Self { revision: self.revision, ..Self::new() };
55        self.advance();
56    }
57
58    pub fn append_notice(&mut self, text: impl Into<String>) -> ConversationItemId {
59        self.push(ItemState::Sealed, ConversationContent::Notice(text.into()))
60    }
61
62    pub fn id(&self) -> ConversationId {
63        self.id
64    }
65
66    pub fn revision(&self) -> Revision {
67        self.revision
68    }
69
70    pub fn items(&self) -> &[ConversationItem] {
71        &self.items
72    }
73
74    pub fn turn(&self) -> TurnPhase {
75        self.turn
76    }
77
78    pub fn activity(&self) -> Activity {
79        self.activity
80    }
81
82    pub fn plan(&self) -> Option<&acp::PlanItems> {
83        self.plan.as_ref()
84    }
85
86    pub fn context_usage(&self) -> Option<&acp::UsageUpdate> {
87        self.context_usage.as_ref()
88    }
89
90    pub fn is_compacting(&self) -> bool {
91        !self.compactions.is_empty()
92    }
93
94    pub fn any_running(&self) -> bool {
95        self.items.iter().any(|item| match item.content() {
96            ConversationContent::Tool(tool_call) => tool_call.is_running(),
97            _ => false,
98        })
99    }
100
101    fn apply_update(&mut self, update: &acp::SessionUpdate) -> Option<TurnFinished> {
102        if matches!(update, acp::SessionUpdate::StateUpdate(acp::StateUpdate::Running(_))) && self.turn.is_idle() {
103            self.begin_turn();
104        } else if !self.turn.is_idle()
105            && let Some(activity) = Activity::after(update)
106        {
107            self.set_activity(activity);
108        }
109        match update {
110            acp::SessionUpdate::CompactionUpdate(update) => {
111                if !self.turn_ended {
112                    self.apply_compaction(update);
113                }
114            }
115            acp::SessionUpdate::StateUpdate(acp::StateUpdate::Idle(idle)) if !self.turn.is_idle() => {
116                return Some(self.finish_turn(idle.stop_reason.clone()));
117            }
118            acp::SessionUpdate::UserMessage(message) => {
119                self.upsert_message(MessageRole::User, &message.message_id, &message.content);
120            }
121            acp::SessionUpdate::AgentMessage(message) => {
122                self.upsert_message(MessageRole::Assistant, &message.message_id, &message.content);
123            }
124            acp::SessionUpdate::AgentThought(message) => {
125                self.upsert_message(MessageRole::Thought, &message.message_id, &message.content);
126            }
127            acp::SessionUpdate::UserMessageChunk(chunk) => self.append_message_chunk(MessageRole::User, chunk),
128            acp::SessionUpdate::AgentMessageChunk(chunk) => self.append_message_chunk(MessageRole::Assistant, chunk),
129            acp::SessionUpdate::AgentThoughtChunk(chunk) => self.append_message_chunk(MessageRole::Thought, chunk),
130            acp::SessionUpdate::ToolCallUpdate(update) => {
131                let index = self.tool_slot(&update.tool_call_id);
132                self.update_tool(index, |tool_call| tool_call.apply_update(update));
133            }
134            acp::SessionUpdate::ToolCallContentChunk(chunk) => {
135                let index = self.tool_slot(&chunk.tool_call_id);
136                self.update_tool(index, |tool_call| tool_call.append_content(chunk.content.clone()));
137            }
138            acp::SessionUpdate::PlanUpdate(update) => {
139                if let acp::PlanUpdateContent::Items(items) = &update.plan
140                    && self.plan.as_ref() != Some(items)
141                {
142                    self.plan = Some(items.clone());
143                    self.advance();
144                }
145            }
146            acp::SessionUpdate::UsageUpdate(usage) if self.context_usage.as_ref() != Some(usage) => {
147                self.context_usage = Some(usage.clone());
148                self.advance();
149            }
150            _ => {}
151        }
152        None
153    }
154
155    fn apply_sub_agent_progress(&mut self, progress: &SubAgentProgressParams) {
156        if self.turn_ended {
157            return;
158        }
159        let Some(&index) = self.tool_index.get(&progress.parent_tool_id) else {
160            return;
161        };
162        self.update_tool(index, |tool_call| tool_call.apply_sub_agent_progress(progress));
163    }
164
165    fn connection_closed(&mut self) {
166        if !self.turn.is_idle() {
167            self.turn = TurnPhase::Idle;
168            self.advance();
169        }
170        self.turn_ended = true;
171        self.set_activity(Activity::Idle);
172    }
173
174    fn begin_turn(&mut self) {
175        self.turn = TurnPhase::Running;
176        self.turn_ended = false;
177        self.activity = Activity::Thinking;
178        self.advance();
179    }
180
181    fn set_activity(&mut self, activity: Activity) {
182        if self.activity != activity {
183            self.activity = activity;
184            self.advance();
185        }
186    }
187
188    fn finish_turn(&mut self, stop_reason: Option<acp::StopReason>) -> TurnFinished {
189        let status = match stop_reason {
190            Some(acp::StopReason::Cancelled) => ToolStatus::Cancelled,
191            _ => ToolStatus::Success,
192        };
193        self.turn = TurnPhase::Idle;
194        self.turn_ended = true;
195        self.compactions.clear();
196        self.activity = Activity::Idle;
197        let revision = self.advance();
198        for item in self.items.iter_mut().filter(|item| item.is_open()) {
199            if let ConversationContent::Tool(tool_call) = &mut item.content {
200                tool_call.finalize(status);
201            }
202            item.state = ItemState::Sealed;
203            item.touch(revision, false);
204        }
205        TurnFinished { stop_reason }
206    }
207
208    fn apply_compaction(&mut self, update: &acp::CompactionUpdate) {
209        let changed = match update.status {
210            acp::CompactionStatus::InProgress => self.compactions.insert(update.compaction_id.clone()),
211            acp::CompactionStatus::Completed | acp::CompactionStatus::Failed | acp::CompactionStatus::Cancelled => {
212                self.compactions.remove(&update.compaction_id)
213            }
214            _ => false,
215        };
216        if changed {
217            self.advance();
218        }
219    }
220
221    fn upsert_message(
222        &mut self,
223        role: MessageRole,
224        message_id: &acp::MessageId,
225        content: &MaybeUndefined<Vec<acp::ContentBlock>>,
226    ) {
227        let blocks = match content {
228            MaybeUndefined::Value(blocks) => blocks.clone(),
229            MaybeUndefined::Null if self.message_index.contains_key(message_id) => Vec::new(),
230            MaybeUndefined::Null | MaybeUndefined::Undefined => return,
231        };
232        let index = self.message_slot(role, message_id.clone());
233        let item = &mut self.items[index];
234        let content = role.content(blocks);
235        if item.content != content {
236            item.content = content;
237            self.touch(index, true);
238        }
239    }
240
241    fn append_message_chunk(&mut self, role: MessageRole, chunk: &acp::ContentChunk) {
242        let index = self.message_slot(role, chunk.message_id.clone());
243        let item = &mut self.items[index];
244        let rewrites = !item.is_open() || role == MessageRole::User;
245        if role.blocks_mut(&mut item.content).is_some_and(|blocks| items::append_block(blocks, &chunk.content)) {
246            self.touch(index, rewrites);
247        }
248    }
249
250    fn message_slot(&mut self, role: MessageRole, message_id: acp::MessageId) -> usize {
251        if let Some(&index) = self.message_index.get(&message_id) {
252            return index;
253        }
254        let index = self.items.len();
255        self.push(ItemState::Open, role.content(Vec::new()));
256        self.items[index].message_id = Some(message_id.clone());
257        self.message_index.insert(message_id, index);
258        index
259    }
260
261    fn tool_slot(&mut self, id: &acp::ToolCallId) -> usize {
262        if let Some(&index) = self.tool_index.get(id.0.as_ref()) {
263            return index;
264        }
265        let index = self.items.len();
266        self.push(
267            ItemState::Open,
268            ConversationContent::Tool(ToolCall::from_update(&acp::ToolCallUpdate::new(id.clone()))),
269        );
270        self.tool_index.insert(id.to_string(), index);
271        index
272    }
273
274    fn push(&mut self, state: ItemState, content: ConversationContent) -> ConversationItemId {
275        let id = ConversationItemId(self.next_item_id);
276        self.next_item_id = self.next_item_id.saturating_add(1);
277        let revision = self.advance();
278        self.items.push(ConversationItem::new(id, revision, state, content));
279        id
280    }
281
282    fn update_tool(&mut self, index: usize, apply: impl FnOnce(&mut ToolCall)) {
283        let item = &mut self.items[index];
284        let rewrites = !item.is_open();
285        let ConversationContent::Tool(tool_call) = &mut item.content else {
286            return;
287        };
288        let previous = tool_call.clone();
289        apply(tool_call);
290        if *tool_call == previous {
291            return;
292        }
293        item.state = if tool_call.rendering_final() { ItemState::Sealed } else { ItemState::Open };
294        self.touch(index, rewrites);
295    }
296
297    fn touch(&mut self, index: usize, rewrites: bool) {
298        let revision = self.advance();
299        self.items[index].touch(revision, rewrites);
300    }
301
302    fn advance(&mut self) -> Revision {
303        self.revision.advance()
304    }
305}