machi_runtime/
state.rs

1//! Conversation state abstractions.
2
3use machi_compaction::max_messages::compact_max_messages;
4use machi_protocol::{MESSAGE_FRAME_TOKENS, estimate_image_tokens, estimate_text_tokens};
5use machi_types::{ContentPart, Message};
6
7/// Mutable conversation backing a turn or session.
8pub trait ConversationState: Send {
9    /// Immutable view of messages.
10    fn messages(&self) -> &[Message];
11    /// Append a message.
12    fn append(&mut self, message: Message);
13    /// Replace the entire message list (compaction / restore).
14    fn replace(&mut self, messages: Vec<Message>);
15    /// Token estimate for compaction triggers (aligned with preflight, W3.2).
16    ///
17    /// Includes framing, multimodal parts, and tool-call argument JSON.
18    fn token_estimate(&self) -> u64 {
19        u64::from(estimate_messages_tokens(self.messages()))
20    }
21}
22
23/// Shared estimator used by [`ConversationState::token_estimate`] and turn preflight.
24#[must_use]
25pub fn estimate_messages_tokens(messages: &[Message]) -> u32 {
26    messages
27        .iter()
28        .fold(0u32, |acc, m| acc.saturating_add(estimate_one_message(m)))
29}
30
31fn estimate_one_message(m: &Message) -> u32 {
32    let mut n = MESSAGE_FRAME_TOKENS;
33    if m.parts.is_empty() {
34        n = n.saturating_add(estimate_text_tokens(&m.text()));
35    } else {
36        for part in &m.parts {
37            n = n.saturating_add(match part {
38                ContentPart::Text { text } => estimate_text_tokens(text),
39                ContentPart::Image { .. } => estimate_image_tokens(),
40                _ => 0,
41            });
42        }
43    }
44    for call in &m.tool_calls {
45        n = n.saturating_add(estimate_text_tokens(&call.name));
46        n = n.saturating_add(estimate_text_tokens(&call.arguments.to_string()));
47    }
48    n
49}
50
51/// In-memory conversation state.
52#[derive(Debug, Clone, Default)]
53pub struct VecConversationState {
54    messages: Vec<Message>,
55}
56
57impl VecConversationState {
58    /// Empty state.
59    #[must_use]
60    pub fn new() -> Self {
61        Self::default()
62    }
63
64    /// Seed with messages.
65    #[must_use]
66    pub fn from_messages(messages: Vec<Message>) -> Self {
67        Self { messages }
68    }
69
70    /// Drop oldest non-system messages until `max_messages` remains.
71    ///
72    /// Delegates to [`machi_compaction::max_messages::compact_max_messages`].
73    pub fn compact_max_messages(&mut self, max_messages: usize) {
74        self.messages = compact_max_messages(std::mem::take(&mut self.messages), max_messages);
75    }
76}
77
78impl ConversationState for VecConversationState {
79    fn messages(&self) -> &[Message] {
80        &self.messages
81    }
82
83    fn append(&mut self, message: Message) {
84        self.messages.push(message);
85    }
86
87    fn replace(&mut self, messages: Vec<Message>) {
88        self.messages = messages;
89    }
90}
91
92#[cfg(test)]
93mod tests {
94    use machi_types::Message;
95
96    use super::*;
97
98    #[test]
99    fn max_messages_keeps_system_and_tail() {
100        let mut state = VecConversationState::from_messages(vec![
101            Message::system("sys"),
102            Message::user("1"),
103            Message::user("2"),
104            Message::user("3"),
105            Message::user("4"),
106        ]);
107        state.compact_max_messages(3);
108        let msgs = state.messages();
109        assert_eq!(msgs.len(), 3);
110        assert_eq!(msgs.first().map(Message::text).as_deref(), Some("sys"));
111        assert_eq!(msgs.get(1).map(Message::text).as_deref(), Some("3"));
112        assert_eq!(msgs.get(2).map(Message::text).as_deref(), Some("4"));
113    }
114}
machi_runtime/state.rs

machi_runtime/
state.rs
