ragrig_gpui/model/
chat.rs

1//! Chat transcript data and its reducer.
2
3use std::time::Duration;
4
5/// Stable identity of one message in a [`ChatState`].
6///
7/// Identifiers stay valid for the lifetime of the transcript. Loading a new
8/// transcript assigns fresh identifiers; hosts should not persist them.
9#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
10pub struct MessageId(u64);
11
12impl MessageId {
13    /// The numeric value, useful for element IDs and logs.
14    pub fn get(self) -> u64 {
15        self.0
16    }
17}
18
19impl std::fmt::Display for MessageId {
20    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
21        write!(formatter, "{}", self.0)
22    }
23}
24
25/// Who authored a message.
26#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
27pub enum ChatRole {
28    /// The human user.
29    User,
30    /// The assistant or model.
31    Assistant,
32    /// A system notice, rendered without a bubble.
33    System,
34}
35
36impl ChatRole {
37    /// Short lowercase name, e.g. for logs.
38    pub fn label(self) -> &'static str {
39        match self {
40            Self::User => "user",
41            Self::Assistant => "assistant",
42            Self::System => "system",
43        }
44    }
45}
46
47/// Lifecycle of one message.
48#[derive(Clone, Debug, PartialEq, Eq, Default)]
49pub enum MessageState {
50    /// Fully received.
51    #[default]
52    Complete,
53    /// Tokens are still arriving.
54    Streaming,
55    /// The host stopped generation; the partial body is kept.
56    Stopped,
57    /// Generation failed. The partial body, if any, is kept.
58    Error {
59        /// Human-readable failure shown beside the body.
60        message: String,
61    },
62}
63
64impl MessageState {
65    /// Returns whether tokens may still be appended.
66    pub fn is_streaming(&self) -> bool {
67        matches!(self, Self::Streaming)
68    }
69}
70
71/// One retrieved source shown under an assistant message.
72#[derive(Clone, Debug, PartialEq)]
73pub struct SourceRef {
74    /// Stable source identifier, used to deduplicate repeated chunks.
75    pub id: String,
76    /// Short title shown on the chip.
77    pub title: String,
78    /// Optional location such as `p. 4` or `methods section`.
79    pub location: Option<String>,
80    /// Optional retrieval score.
81    pub score: Option<f32>,
82}
83
84impl SourceRef {
85    /// A source with a title and no location or score.
86    pub fn new(id: impl Into<String>, title: impl Into<String>) -> Self {
87        Self {
88            id: id.into(),
89            title: title.into(),
90            location: None,
91            score: None,
92        }
93    }
94
95    /// Set the location label.
96    pub fn location(mut self, location: impl Into<String>) -> Self {
97        self.location = Some(location.into());
98        self
99    }
100
101    /// Set the retrieval score.
102    pub fn score(mut self, score: f32) -> Self {
103        self.score = Some(score);
104        self
105    }
106}
107
108/// Pipeline metadata shown in a message footer.
109#[derive(Clone, Debug, Default, PartialEq)]
110pub struct MessageMeta {
111    /// Model that produced the answer.
112    pub model: Option<String>,
113    /// Number of retrieved chunks that entered the prompt.
114    pub chunks: Option<usize>,
115    /// Distinct sources among the retrieved chunks.
116    pub sources: Vec<SourceRef>,
117    /// Number of streamed tokens, when the host counts them.
118    pub tokens: Option<usize>,
119    /// Wall-clock generation time.
120    pub elapsed: Option<Duration>,
121}
122
123impl MessageMeta {
124    /// One-line summary, e.g. `deepseek-chat · 12 chunks · 3 sources · 421 tokens · 8.2 s`.
125    ///
126    /// Source titles are rendered separately as chips, so only their count is
127    /// included here.
128    pub fn summary(&self) -> String {
129        let mut parts = Vec::new();
130        if let Some(model) = &self.model {
131            parts.push(model.clone());
132        }
133        if let Some(chunks) = self.chunks {
134            parts.push(format!("{chunks} chunks"));
135        }
136        if !self.sources.is_empty() {
137            let count = self.sources.len();
138            parts.push(if count == 1 {
139                "1 source".to_owned()
140            } else {
141                format!("{count} sources")
142            });
143        }
144        if let Some(tokens) = self.tokens {
145            parts.push(format!("{tokens} tokens"));
146        }
147        if let Some(elapsed) = self.elapsed {
148            parts.push(format!("{:.1} s", elapsed.as_secs_f32()));
149        }
150        parts.join(" · ")
151    }
152
153    /// Returns whether the footer has anything to show.
154    pub fn is_empty(&self) -> bool {
155        self.model.is_none()
156            && self.chunks.is_none()
157            && self.sources.is_empty()
158            && self.tokens.is_none()
159            && self.elapsed.is_none()
160    }
161
162    /// Add a source unless one with the same [`SourceRef::id`] is present.
163    ///
164    /// Returns whether the source was added.
165    pub fn add_source(&mut self, source: SourceRef) -> bool {
166        if self.sources.iter().any(|existing| existing.id == source.id) {
167            return false;
168        }
169        self.sources.push(source);
170        true
171    }
172}
173
174/// A message as stored in a [`ChatState`].
175#[derive(Clone, Debug, PartialEq)]
176pub struct ChatMessage {
177    id: MessageId,
178    role: ChatRole,
179    text: String,
180    state: MessageState,
181    meta: Option<MessageMeta>,
182}
183
184impl ChatMessage {
185    /// Stable message identity.
186    pub fn id(&self) -> MessageId {
187        self.id
188    }
189
190    /// Author of the message.
191    pub fn role(&self) -> ChatRole {
192        self.role
193    }
194
195    /// Message body in Markdown.
196    pub fn text(&self) -> &str {
197        &self.text
198    }
199
200    /// Lifecycle state.
201    pub fn state(&self) -> &MessageState {
202        &self.state
203    }
204
205    /// Footer metadata, if any.
206    pub fn meta(&self) -> Option<&MessageMeta> {
207        self.meta.as_ref()
208    }
209}
210
211/// Input for [`ChatState::push`] and [`ChatState::replace`].
212#[derive(Clone, Debug, PartialEq)]
213pub struct NewMessage {
214    /// Author of the message.
215    pub role: ChatRole,
216    /// Message body in Markdown.
217    pub text: String,
218    /// Lifecycle state.
219    pub state: MessageState,
220    /// Optional footer metadata.
221    pub meta: Option<MessageMeta>,
222}
223
224impl NewMessage {
225    /// A complete user message.
226    pub fn user(text: impl Into<String>) -> Self {
227        Self::new(ChatRole::User, text)
228    }
229
230    /// A complete assistant message.
231    pub fn assistant(text: impl Into<String>) -> Self {
232        Self::new(ChatRole::Assistant, text)
233    }
234
235    /// A complete system notice.
236    pub fn system(text: impl Into<String>) -> Self {
237        Self::new(ChatRole::System, text)
238    }
239
240    /// A message with the given role and complete state.
241    pub fn new(role: ChatRole, text: impl Into<String>) -> Self {
242        Self {
243            role,
244            text: text.into(),
245            state: MessageState::Complete,
246            meta: None,
247        }
248    }
249
250    /// Mark the message as streaming.
251    pub fn streaming(mut self) -> Self {
252        self.state = MessageState::Streaming;
253        self
254    }
255
256    /// Override the lifecycle state.
257    pub fn state(mut self, state: MessageState) -> Self {
258        self.state = state;
259        self
260    }
261
262    /// Attach footer metadata.
263    pub fn meta(mut self, meta: MessageMeta) -> Self {
264        self.meta = Some(meta);
265        self
266    }
267}
268
269/// The transcript reducer.
270///
271/// This is plain data: construct it in a host or wrap it in an entity. The
272/// rules are deliberately small:
273///
274/// - only [`MessageState::Streaming`] messages accept [`append`](Self::append);
275/// - [`finish`](Self::finish) only moves a streaming message to complete;
276/// - [`stop`](Self::stop) keeps the partial body;
277/// - [`fail`](Self::fail) keeps the partial body and records the reason.
278#[derive(Debug, Default)]
279pub struct ChatState {
280    next_id: u64,
281    messages: Vec<ChatMessage>,
282}
283
284impl ChatState {
285    /// All messages in display order.
286    pub fn messages(&self) -> &[ChatMessage] {
287        &self.messages
288    }
289
290    /// Number of messages.
291    pub fn len(&self) -> usize {
292        self.messages.len()
293    }
294
295    /// Returns whether the transcript is empty.
296    pub fn is_empty(&self) -> bool {
297        self.messages.is_empty()
298    }
299
300    /// Look up a message by identity.
301    pub fn get(&self, id: MessageId) -> Option<&ChatMessage> {
302        self.messages.iter().find(|message| message.id == id)
303    }
304
305    /// The most recent message.
306    pub fn last(&self) -> Option<&ChatMessage> {
307        self.messages.last()
308    }
309
310    /// The most recent assistant message.
311    pub fn last_assistant(&self) -> Option<&ChatMessage> {
312        self.messages
313            .iter()
314            .rev()
315            .find(|message| message.role == ChatRole::Assistant)
316    }
317
318    /// The message currently streaming, if any.
319    pub fn streaming_id(&self) -> Option<MessageId> {
320        self.messages
321            .iter()
322            .rev()
323            .find(|message| message.state.is_streaming())
324            .map(|message| message.id)
325    }
326
327    /// Returns whether any message is still streaming.
328    pub fn is_streaming(&self) -> bool {
329        self.streaming_id().is_some()
330    }
331
332    /// Append a message and assign it a fresh identity.
333    pub fn push(&mut self, message: NewMessage) -> MessageId {
334        let id = MessageId(self.next_id);
335        self.next_id += 1;
336        self.messages.push(ChatMessage {
337            id,
338            role: message.role,
339            text: message.text,
340            state: message.state,
341            meta: message.meta,
342        });
343        id
344    }
345
346    /// Insert a message immediately before `before`.
347    ///
348    /// Returns `None` when `before` is unknown; the transcript is unchanged in
349    /// that case.
350    pub fn insert_before(&mut self, before: MessageId, message: NewMessage) -> Option<MessageId> {
351        let index = self
352            .messages
353            .iter()
354            .position(|candidate| candidate.id == before)?;
355        let id = MessageId(self.next_id);
356        self.next_id += 1;
357        self.messages.insert(
358            index,
359            ChatMessage {
360                id,
361                role: message.role,
362                text: message.text,
363                state: message.state,
364                meta: message.meta,
365            },
366        );
367        Some(id)
368    }
369
370    /// Append a complete user message.
371    pub fn push_user(&mut self, text: impl Into<String>) -> MessageId {
372        self.push(NewMessage::user(text))
373    }
374
375    /// Append a complete system notice.
376    pub fn push_system(&mut self, text: impl Into<String>) -> MessageId {
377        self.push(NewMessage::system(text))
378    }
379
380    /// Append an empty streaming assistant message.
381    pub fn begin_assistant(&mut self) -> MessageId {
382        self.push(NewMessage::assistant("").streaming())
383    }
384
385    /// Append `delta` to a streaming message.
386    ///
387    /// Returns `false` for unknown messages and for messages that are no
388    /// longer streaming, so late tokens can be dropped safely.
389    pub fn append(&mut self, id: MessageId, delta: &str) -> bool {
390        let Some(message) = self.messages.iter_mut().find(|message| message.id == id) else {
391            return false;
392        };
393        if !message.state.is_streaming() {
394            return false;
395        }
396        message.text.push_str(delta);
397        true
398    }
399
400    /// Mark a streaming message as complete.
401    pub fn finish(&mut self, id: MessageId) -> bool {
402        self.transition(id, &MessageState::Complete)
403    }
404
405    /// Mark a streaming message as stopped by the user.
406    pub fn stop(&mut self, id: MessageId) -> bool {
407        self.transition(id, &MessageState::Stopped)
408    }
409
410    /// Mark a message as failed, keeping any partial body.
411    pub fn fail(&mut self, id: MessageId, error: impl Into<String>) -> bool {
412        let Some(message) = self.messages.iter_mut().find(|message| message.id == id) else {
413            return false;
414        };
415        message.state = MessageState::Error {
416            message: error.into(),
417        };
418        true
419    }
420
421    /// Replace a message's footer metadata.
422    pub fn set_meta(&mut self, id: MessageId, meta: MessageMeta) -> bool {
423        let Some(message) = self.messages.iter_mut().find(|message| message.id == id) else {
424            return false;
425        };
426        message.meta = Some(meta);
427        true
428    }
429
430    /// Replace a message's body, e.g. after lifting a protocol block out of a
431    /// finished answer. Returns whether the message existed.
432    pub fn set_text(&mut self, id: MessageId, text: impl Into<String>) -> bool {
433        let Some(message) = self.messages.iter_mut().find(|message| message.id == id) else {
434            return false;
435        };
436        message.text = text.into();
437        true
438    }
439
440    /// Add one source to a message's footer, deduplicating by source id.
441    ///
442    /// Returns whether the source was added.
443    pub fn add_source(&mut self, id: MessageId, source: SourceRef) -> bool {
444        let Some(message) = self.messages.iter_mut().find(|message| message.id == id) else {
445            return false;
446        };
447        message
448            .meta
449            .get_or_insert_with(MessageMeta::default)
450            .add_source(source)
451    }
452
453    /// Remove every message. Identities continue from the previous run.
454    pub fn clear(&mut self) {
455        self.messages.clear();
456    }
457
458    /// Remove one message, keeping the relative order of the rest.
459    ///
460    /// Returns whether the message existed.
461    pub fn remove(&mut self, id: MessageId) -> bool {
462        let before = self.messages.len();
463        self.messages.retain(|message| message.id != id);
464        self.messages.len() != before
465    }
466
467    /// Replace the transcript, assigning fresh identities.
468    pub fn replace(&mut self, messages: impl IntoIterator<Item = NewMessage>) -> Vec<MessageId> {
469        self.messages.clear();
470        messages
471            .into_iter()
472            .map(|message| self.push(message))
473            .collect()
474    }
475
476    fn transition(&mut self, id: MessageId, next: &MessageState) -> bool {
477        let Some(message) = self.messages.iter_mut().find(|message| message.id == id) else {
478            return false;
479        };
480        if !message.state.is_streaming() {
481            return false;
482        }
483        message.state = next.clone();
484        true
485    }
486}
487
488#[cfg(test)]
489mod tests {
490    use super::*;
491
492    #[test]
493    fn a_streamed_message_is_reduced_in_order() {
494        let mut state = ChatState::default();
495        let user = state.push_user("Explain the method");
496        let assistant = state.begin_assistant();
497
498        assert_eq!(state.len(), 2);
499        assert_eq!(state.messages()[0].id(), user);
500        assert_eq!(state.messages()[0].role(), ChatRole::User);
501        assert!(state.append(assistant, "The "));
502        assert!(state.append(assistant, "method."));
503        assert!(state.finish(assistant));
504
505        let message = state.get(assistant).expect("assistant message");
506        assert_eq!(message.text(), "The method.");
507        assert_eq!(message.state(), &MessageState::Complete);
508        assert!(!state.is_streaming());
509    }
510
511    #[test]
512    fn stopping_keeps_the_partial_body_and_rejects_late_tokens() {
513        let mut state = ChatState::default();
514        let assistant = state.begin_assistant();
515        assert!(state.append(assistant, "Partial"));
516
517        assert!(state.stop(assistant));
518        assert!(!state.append(assistant, " late"));
519        assert!(!state.finish(assistant));
520
521        let message = state.get(assistant).expect("assistant message");
522        assert_eq!(message.text(), "Partial");
523        assert_eq!(message.state(), &MessageState::Stopped);
524    }
525
526    #[test]
527    fn failing_keeps_the_partial_body_and_reason() {
528        let mut state = ChatState::default();
529        let assistant = state.begin_assistant();
530        assert!(state.append(assistant, "Partial"));
531
532        assert!(state.fail(assistant, "connection reset"));
533
534        let message = state.get(assistant).expect("assistant message");
535        assert_eq!(message.text(), "Partial");
536        assert_eq!(
537            message.state(),
538            &MessageState::Error {
539                message: "connection reset".to_owned()
540            }
541        );
542    }
543
544    #[test]
545    fn metadata_and_sources_are_attached_and_deduplicated() {
546        let mut state = ChatState::default();
547        let assistant = state.begin_assistant();
548        assert!(state.set_meta(
549            assistant,
550            MessageMeta {
551                model: Some("deepseek-chat".to_owned()),
552                chunks: Some(12),
553                ..Default::default()
554            }
555        ));
556        assert!(state.add_source(assistant, SourceRef::new("paper", "Paper")));
557        assert!(!state.add_source(assistant, SourceRef::new("paper", "Paper again")));
558
559        let message = state.get(assistant).expect("assistant message");
560        let meta = message.meta().expect("meta");
561        assert_eq!(meta.sources.len(), 1);
562        assert_eq!(meta.summary(), "deepseek-chat · 12 chunks · 1 source");
563    }
564
565    #[test]
566    fn removing_a_message_keeps_the_rest_in_order() {
567        let mut state = ChatState::default();
568        let first = state.push_user("first");
569        let second = state.push_user("second");
570        let third = state.push_user("third");
571
572        assert!(state.remove(second));
573        assert!(!state.remove(second));
574        assert_eq!(
575            state
576                .messages()
577                .iter()
578                .map(ChatMessage::id)
579                .collect::<Vec<_>>(),
580            vec![first, third]
581        );
582    }
583
584    #[test]
585    fn inserting_before_keeps_the_transcript_order() {
586        let mut state = ChatState::default();
587        let user = state.push_user("question");
588        let assistant = state.begin_assistant();
589
590        let note = state
591            .insert_before(assistant, NewMessage::system("Attached: @Kano2025"))
592            .expect("insert before the assistant");
593        assert!(note > assistant);
594        assert_eq!(
595            state
596                .messages()
597                .iter()
598                .map(ChatMessage::role)
599                .collect::<Vec<_>>(),
600            vec![ChatRole::User, ChatRole::System, ChatRole::Assistant]
601        );
602        assert_eq!(state.messages()[0].id(), user);
603        assert_eq!(state.messages()[2].id(), assistant);
604
605        assert!(
606            state
607                .insert_before(MessageId(999), NewMessage::system("late"))
608                .is_none()
609        );
610        assert_eq!(state.len(), 3);
611    }
612
613    #[test]
614    fn replacing_a_transcript_assigns_fresh_identities() {
615        let mut state = ChatState::default();
616        let old = state.push_user("old");
617        let ids = state.replace([NewMessage::user("new"), NewMessage::assistant("answer")]);
618
619        assert_eq!(state.len(), 2);
620        assert_eq!(state.messages()[0].text(), "new");
621        assert!(ids.iter().all(|id| *id > old));
622    }
623
624    #[test]
625    fn append_and_transitions_ignore_unknown_ids() {
626        let mut state = ChatState::default();
627        let id = MessageId(999);
628        assert!(!state.append(id, "x"));
629        assert!(!state.finish(id));
630        assert!(!state.stop(id));
631        assert!(!state.fail(id, "x"));
632        assert!(!state.set_meta(id, MessageMeta::default()));
633        assert!(!state.add_source(id, SourceRef::new("s", "S")));
634    }
635
636    #[test]
637    fn a_finished_message_does_not_accept_more_tokens() {
638        let mut state = ChatState::default();
639        let assistant = state.begin_assistant();
640        assert!(state.finish(assistant));
641        assert!(!state.append(assistant, "late"));
642
643        assert_eq!(state.get(assistant).expect("assistant").text(), "");
644    }
645}
ragrig_gpui/model/chat.rs

ragrig_gpui/model/
chat.rs
