agent_code_lib/services/
history.rs

1//! Conversation history utilities.
2//!
3//! Functions for manipulating, searching, and transforming the
4//! message history. Used by compaction, export, and the query engine.
5
6use crate::llm::message::{ContentBlock, Message};
7
8/// Count messages by type.
9pub fn message_counts(messages: &[Message]) -> (usize, usize, usize) {
10    let mut user = 0;
11    let mut assistant = 0;
12    let mut system = 0;
13
14    for msg in messages {
15        match msg {
16            Message::User(_) => user += 1,
17            Message::Assistant(_) => assistant += 1,
18            Message::System(_) => system += 1,
19        }
20    }
21
22    (user, assistant, system)
23}
24
25/// Extract all text content from messages (for search/export).
26pub fn extract_text(messages: &[Message]) -> String {
27    let mut text = String::new();
28    for msg in messages {
29        let blocks = match msg {
30            Message::User(u) => &u.content,
31            Message::Assistant(a) => &a.content,
32            Message::System(s) => {
33                text.push_str(&s.content);
34                text.push('\n');
35                continue;
36            }
37        };
38        for block in blocks {
39            if let ContentBlock::Text { text: t } = block {
40                text.push_str(t);
41                text.push('\n');
42            }
43        }
44    }
45    text
46}
47
48/// Find the index of the last user message (non-meta).
49pub fn last_user_message_index(messages: &[Message]) -> Option<usize> {
50    messages
51        .iter()
52        .rposition(|m| matches!(m, Message::User(u) if !u.is_meta))
53}
54
55/// Find the index of the last assistant message.
56pub fn last_assistant_index(messages: &[Message]) -> Option<usize> {
57    messages
58        .iter()
59        .rposition(|m| matches!(m, Message::Assistant(_)))
60}
61
62/// Count tool use blocks in the conversation.
63pub fn tool_use_count(messages: &[Message]) -> usize {
64    messages
65        .iter()
66        .filter_map(|m| match m {
67            Message::Assistant(a) => Some(&a.content),
68            _ => None,
69        })
70        .flat_map(|blocks| blocks.iter())
71        .filter(|b| matches!(b, ContentBlock::ToolUse { .. }))
72        .count()
73}
74
75/// Get a list of unique tools used in the conversation.
76pub fn tools_used(messages: &[Message]) -> Vec<String> {
77    let mut tools: Vec<String> = messages
78        .iter()
79        .filter_map(|m| match m {
80            Message::Assistant(a) => Some(&a.content),
81            _ => None,
82        })
83        .flat_map(|blocks| blocks.iter())
84        .filter_map(|b| match b {
85            ContentBlock::ToolUse { name, .. } => Some(name.clone()),
86            _ => None,
87        })
88        .collect();
89
90    tools.sort();
91    tools.dedup();
92    tools
93}
94
95/// Truncate messages to fit within a token budget.
96///
97/// Removes oldest messages (preserving the first system/summary message)
98/// until the estimated token count is within budget.
99pub fn truncate_to_budget(messages: &mut Vec<Message>, max_tokens: u64) {
100    while crate::services::tokens::estimate_context_tokens(messages) > max_tokens
101        && messages.len() > 2
102    {
103        // Remove the second message (preserve index 0 which may be a summary).
104        messages.remove(1);
105    }
106}
107
108#[cfg(test)]
109mod tests {
110    use super::*;
111    use crate::llm::message::{AssistantMessage, ContentBlock, user_message};
112    use uuid::Uuid;
113
114    fn assistant_msg(text: &str) -> Message {
115        Message::Assistant(AssistantMessage {
116            uuid: Uuid::new_v4(),
117            timestamp: String::new(),
118            content: vec![ContentBlock::Text {
119                text: text.to_string(),
120            }],
121            model: None,
122            usage: None,
123            stop_reason: None,
124            request_id: None,
125        })
126    }
127
128    #[test]
129    fn test_message_counts() {
130        let msgs = vec![
131            user_message("hello"),
132            assistant_msg("hi"),
133            user_message("bye"),
134        ];
135        assert_eq!(message_counts(&msgs), (2, 1, 0));
136    }
137
138    #[test]
139    fn test_tool_use_count() {
140        let msgs = vec![Message::Assistant(AssistantMessage {
141            uuid: Uuid::new_v4(),
142            timestamp: String::new(),
143            content: vec![
144                ContentBlock::ToolUse {
145                    id: "1".into(),
146                    name: "Bash".into(),
147                    input: serde_json::json!({}),
148                },
149                ContentBlock::Text {
150                    text: "done".into(),
151                },
152            ],
153            model: None,
154            usage: None,
155            stop_reason: None,
156            request_id: None,
157        })];
158        assert_eq!(tool_use_count(&msgs), 1);
159    }
160
161    #[test]
162    fn test_extract_text() {
163        let msgs = vec![user_message("hello world"), assistant_msg("response here")];
164        let text = extract_text(&msgs);
165        assert!(text.contains("hello world"));
166        assert!(text.contains("response here"));
167    }
168
169    #[test]
170    fn test_last_user_message_index() {
171        let msgs = vec![
172            user_message("first"),
173            assistant_msg("reply"),
174            user_message("second"),
175        ];
176        assert_eq!(last_user_message_index(&msgs), Some(2));
177    }
178
179    #[test]
180    fn test_last_assistant_index() {
181        let msgs = vec![
182            user_message("first"),
183            assistant_msg("reply"),
184            user_message("second"),
185        ];
186        assert_eq!(last_assistant_index(&msgs), Some(1));
187    }
188
189    #[test]
190    fn test_tools_used() {
191        let msgs = vec![Message::Assistant(AssistantMessage {
192            uuid: Uuid::new_v4(),
193            timestamp: String::new(),
194            content: vec![
195                ContentBlock::ToolUse {
196                    id: "1".into(),
197                    name: "Bash".into(),
198                    input: serde_json::json!({}),
199                },
200                ContentBlock::ToolUse {
201                    id: "2".into(),
202                    name: "FileRead".into(),
203                    input: serde_json::json!({}),
204                },
205                ContentBlock::ToolUse {
206                    id: "3".into(),
207                    name: "Bash".into(),
208                    input: serde_json::json!({}),
209                },
210            ],
211            model: None,
212            usage: None,
213            stop_reason: None,
214            request_id: None,
215        })];
216        let used = tools_used(&msgs);
217        assert!(used.contains(&"Bash".to_string()));
218        assert!(used.contains(&"FileRead".to_string()));
219        assert_eq!(used.len(), 2); // Deduplicated.
220    }
221
222    #[test]
223    fn test_empty_messages() {
224        assert_eq!(message_counts(&[]), (0, 0, 0));
225        assert_eq!(tool_use_count(&[]), 0);
226        assert!(extract_text(&[]).is_empty());
227        assert_eq!(last_user_message_index(&[]), None);
228        assert_eq!(last_assistant_index(&[]), None);
229    }
230}
agent_code_lib/services/history.rs

agent_code_lib/services/
history.rs
