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    /// Add one source to a message's footer, deduplicating by source id.
431    ///
432    /// Returns whether the source was added.
433    pub fn add_source(&mut self, id: MessageId, source: SourceRef) -> bool {
434        let Some(message) = self.messages.iter_mut().find(|message| message.id == id) else {
435            return false;
436        };
437        message
438            .meta
439            .get_or_insert_with(MessageMeta::default)
440            .add_source(source)
441    }
442
443    /// Remove every message. Identities continue from the previous run.
444    pub fn clear(&mut self) {
445        self.messages.clear();
446    }
447
448    /// Remove one message, keeping the relative order of the rest.
449    ///
450    /// Returns whether the message existed.
451    pub fn remove(&mut self, id: MessageId) -> bool {
452        let before = self.messages.len();
453        self.messages.retain(|message| message.id != id);
454        self.messages.len() != before
455    }
456
457    /// Replace the transcript, assigning fresh identities.
458    pub fn replace(&mut self, messages: impl IntoIterator<Item = NewMessage>) -> Vec<MessageId> {
459        self.messages.clear();
460        messages
461            .into_iter()
462            .map(|message| self.push(message))
463            .collect()
464    }
465
466    fn transition(&mut self, id: MessageId, next: &MessageState) -> bool {
467        let Some(message) = self.messages.iter_mut().find(|message| message.id == id) else {
468            return false;
469        };
470        if !message.state.is_streaming() {
471            return false;
472        }
473        message.state = next.clone();
474        true
475    }
476}
477
478#[cfg(test)]
479mod tests {
480    use super::*;
481
482    #[test]
483    fn a_streamed_message_is_reduced_in_order() {
484        let mut state = ChatState::default();
485        let user = state.push_user("Explain the method");
486        let assistant = state.begin_assistant();
487
488        assert_eq!(state.len(), 2);
489        assert_eq!(state.messages()[0].id(), user);
490        assert_eq!(state.messages()[0].role(), ChatRole::User);
491        assert!(state.append(assistant, "The "));
492        assert!(state.append(assistant, "method."));
493        assert!(state.finish(assistant));
494
495        let message = state.get(assistant).expect("assistant message");
496        assert_eq!(message.text(), "The method.");
497        assert_eq!(message.state(), &MessageState::Complete);
498        assert!(!state.is_streaming());
499    }
500
501    #[test]
502    fn stopping_keeps_the_partial_body_and_rejects_late_tokens() {
503        let mut state = ChatState::default();
504        let assistant = state.begin_assistant();
505        assert!(state.append(assistant, "Partial"));
506
507        assert!(state.stop(assistant));
508        assert!(!state.append(assistant, " late"));
509        assert!(!state.finish(assistant));
510
511        let message = state.get(assistant).expect("assistant message");
512        assert_eq!(message.text(), "Partial");
513        assert_eq!(message.state(), &MessageState::Stopped);
514    }
515
516    #[test]
517    fn failing_keeps_the_partial_body_and_reason() {
518        let mut state = ChatState::default();
519        let assistant = state.begin_assistant();
520        assert!(state.append(assistant, "Partial"));
521
522        assert!(state.fail(assistant, "connection reset"));
523
524        let message = state.get(assistant).expect("assistant message");
525        assert_eq!(message.text(), "Partial");
526        assert_eq!(
527            message.state(),
528            &MessageState::Error {
529                message: "connection reset".to_owned()
530            }
531        );
532    }
533
534    #[test]
535    fn metadata_and_sources_are_attached_and_deduplicated() {
536        let mut state = ChatState::default();
537        let assistant = state.begin_assistant();
538        assert!(state.set_meta(
539            assistant,
540            MessageMeta {
541                model: Some("deepseek-chat".to_owned()),
542                chunks: Some(12),
543                ..Default::default()
544            }
545        ));
546        assert!(state.add_source(assistant, SourceRef::new("paper", "Paper")));
547        assert!(!state.add_source(assistant, SourceRef::new("paper", "Paper again")));
548
549        let message = state.get(assistant).expect("assistant message");
550        let meta = message.meta().expect("meta");
551        assert_eq!(meta.sources.len(), 1);
552        assert_eq!(meta.summary(), "deepseek-chat · 12 chunks · 1 source");
553    }
554
555    #[test]
556    fn removing_a_message_keeps_the_rest_in_order() {
557        let mut state = ChatState::default();
558        let first = state.push_user("first");
559        let second = state.push_user("second");
560        let third = state.push_user("third");
561
562        assert!(state.remove(second));
563        assert!(!state.remove(second));
564        assert_eq!(
565            state
566                .messages()
567                .iter()
568                .map(ChatMessage::id)
569                .collect::<Vec<_>>(),
570            vec![first, third]
571        );
572    }
573
574    #[test]
575    fn inserting_before_keeps_the_transcript_order() {
576        let mut state = ChatState::default();
577        let user = state.push_user("question");
578        let assistant = state.begin_assistant();
579
580        let note = state
581            .insert_before(assistant, NewMessage::system("Attached: @Kano2025"))
582            .expect("insert before the assistant");
583        assert!(note > assistant);
584        assert_eq!(
585            state
586                .messages()
587                .iter()
588                .map(ChatMessage::role)
589                .collect::<Vec<_>>(),
590            vec![ChatRole::User, ChatRole::System, ChatRole::Assistant]
591        );
592        assert_eq!(state.messages()[0].id(), user);
593        assert_eq!(state.messages()[2].id(), assistant);
594
595        assert!(
596            state
597                .insert_before(MessageId(999), NewMessage::system("late"))
598                .is_none()
599        );
600        assert_eq!(state.len(), 3);
601    }
602
603    #[test]
604    fn replacing_a_transcript_assigns_fresh_identities() {
605        let mut state = ChatState::default();
606        let old = state.push_user("old");
607        let ids = state.replace([NewMessage::user("new"), NewMessage::assistant("answer")]);
608
609        assert_eq!(state.len(), 2);
610        assert_eq!(state.messages()[0].text(), "new");
611        assert!(ids.iter().all(|id| *id > old));
612    }
613
614    #[test]
615    fn append_and_transitions_ignore_unknown_ids() {
616        let mut state = ChatState::default();
617        let id = MessageId(999);
618        assert!(!state.append(id, "x"));
619        assert!(!state.finish(id));
620        assert!(!state.stop(id));
621        assert!(!state.fail(id, "x"));
622        assert!(!state.set_meta(id, MessageMeta::default()));
623        assert!(!state.add_source(id, SourceRef::new("s", "S")));
624    }
625
626    #[test]
627    fn a_finished_message_does_not_accept_more_tokens() {
628        let mut state = ChatState::default();
629        let assistant = state.begin_assistant();
630        assert!(state.finish(assistant));
631        assert!(!state.append(assistant, "late"));
632
633        assert_eq!(state.get(assistant).expect("assistant").text(), "");
634    }
635}
ragrig_gpui/model/chat.rs

ragrig_gpui/model/
chat.rs
