context_forge/
session.rs

1use std::collections::BTreeMap;
2
3use crate::ContextEntry;
4
5/// A grouped set of entries that belong to the same derived session.
6#[derive(Debug, Clone)]
7pub struct SessionGroup {
8    /// Derived or explicit session identifier.
9    pub session_id: String,
10    /// Owned entries for this session group.
11    pub entries: Vec<ContextEntry>,
12}
13
14/// Groups entries by explicit session id, then by fallback timestamp proximity.
15///
16/// # Panics
17///
18/// Panics if `proximity_threshold_secs` is negative.
19#[must_use]
20pub fn group_entries_by_session(
21    entries: &[ContextEntry],
22    proximity_threshold_secs: i64,
23) -> Vec<SessionGroup> {
24    assert!(
25        proximity_threshold_secs >= 0,
26        "proximity_threshold_secs must be non-negative"
27    );
28
29    let mut explicit_groups: BTreeMap<String, Vec<ContextEntry>> = BTreeMap::new();
30    let mut fallback_entries: Vec<ContextEntry> = Vec::new();
31
32    for entry in entries {
33        if let Some(session_id) = &entry.session_id {
34            explicit_groups
35                .entry(session_id.clone())
36                .or_default()
37                .push(entry.clone());
38        } else {
39            fallback_entries.push(entry.clone());
40        }
41    }
42
43    let mut groups: Vec<SessionGroup> = explicit_groups
44        .into_iter()
45        .map(|(session_id, entries)| SessionGroup {
46            session_id,
47            entries,
48        })
49        .collect();
50
51    for group in &mut groups {
52        group.entries.sort_by_key(|e| e.timestamp);
53    }
54
55    groups.sort_by_key(|g| g.entries.first().map_or(i64::MAX, |e| e.timestamp));
56
57    fallback_entries.sort_by_key(|entry| entry.timestamp);
58
59    let mut current_group: Vec<ContextEntry> = Vec::new();
60    let mut current_earliest_timestamp: Option<i64> = None;
61    let mut previous_timestamp: Option<i64> = None;
62
63    for entry in fallback_entries {
64        if let Some(prev_ts) = previous_timestamp {
65            let gap = entry.timestamp.saturating_sub(prev_ts);
66            if gap > proximity_threshold_secs {
67                if let Some(earliest) = current_earliest_timestamp {
68                    groups.push(SessionGroup {
69                        session_id: format!("__fallback__{earliest}"),
70                        entries: std::mem::take(&mut current_group),
71                    });
72                }
73                current_earliest_timestamp = Some(entry.timestamp);
74            }
75        } else {
76            current_earliest_timestamp = Some(entry.timestamp);
77        }
78
79        previous_timestamp = Some(entry.timestamp);
80        current_group.push(entry);
81    }
82
83    if !current_group.is_empty() {
84        let earliest =
85            current_earliest_timestamp.expect("invariant: timestamp set when group non-empty");
86        groups.push(SessionGroup {
87            session_id: format!("__fallback__{earliest}"),
88            entries: current_group,
89        });
90    }
91
92    groups.sort_by_key(|g| g.entries.first().map_or(i64::MAX, |e| e.timestamp));
93
94    groups
95}
96
97#[cfg(test)]
98mod tests {
99    use super::{group_entries_by_session, SessionGroup};
100    use crate::{kind, ContextEntry};
101
102    fn make_entry(id: &str, timestamp: i64, session_id: Option<&str>) -> ContextEntry {
103        ContextEntry {
104            id: id.to_string(),
105            content: format!("content-{id}"),
106            timestamp,
107            kind: kind::MANUAL.to_owned(),
108            scope: None,
109            session_id: session_id.map(ToString::to_string),
110            token_count: None,
111            metadata: None,
112        }
113    }
114
115    fn assert_group_ids(groups: &[SessionGroup], expected_ids: &[&str]) {
116        let actual_ids: Vec<&str> = groups
117            .iter()
118            .map(|group| group.session_id.as_str())
119            .collect();
120        assert_eq!(actual_ids, expected_ids);
121    }
122
123    #[test]
124    fn empty_input_returns_empty_output() {
125        let groups = group_entries_by_session(&[], 60);
126        assert!(groups.is_empty());
127    }
128
129    #[test]
130    fn single_entry_with_session_id_creates_one_explicit_group() {
131        let entries = vec![make_entry("e1", 100, Some("s1"))];
132
133        let groups = group_entries_by_session(&entries, 60);
134
135        assert_eq!(groups.len(), 1);
136        assert_group_ids(&groups, &["s1"]);
137        assert_eq!(groups[0].entries.len(), 1);
138        assert_eq!(groups[0].entries[0].id, "e1");
139    }
140
141    #[test]
142    fn single_entry_without_session_id_creates_one_fallback_group() {
143        let entries = vec![make_entry("e1", 100, None)];
144
145        let groups = group_entries_by_session(&entries, 60);
146
147        assert_eq!(groups.len(), 1);
148        assert_group_ids(&groups, &["__fallback__100"]);
149        assert_eq!(groups[0].entries.len(), 1);
150        assert_eq!(groups[0].entries[0].id, "e1");
151    }
152
153    #[test]
154    fn same_session_id_entries_are_grouped_together() {
155        let entries = vec![
156            make_entry("e1", 100, Some("s1")),
157            make_entry("e2", 200, Some("s1")),
158            make_entry("e3", 300, Some("s1")),
159        ];
160
161        let groups = group_entries_by_session(&entries, 60);
162
163        assert_eq!(groups.len(), 1);
164        assert_group_ids(&groups, &["s1"]);
165        let ids: Vec<&str> = groups[0]
166            .entries
167            .iter()
168            .map(|entry| entry.id.as_str())
169            .collect();
170        assert_eq!(ids, vec!["e1", "e2", "e3"]);
171    }
172
173    #[test]
174    fn different_session_ids_create_separate_groups() {
175        let entries = vec![
176            make_entry("e1", 100, Some("s1")),
177            make_entry("e2", 200, Some("s2")),
178            make_entry("e3", 300, Some("s1")),
179        ];
180
181        let groups = group_entries_by_session(&entries, 60);
182
183        assert_eq!(groups.len(), 2);
184        assert_group_ids(&groups, &["s1", "s2"]);
185        assert_eq!(groups[0].entries.len(), 2);
186        assert_eq!(groups[1].entries.len(), 1);
187    }
188
189    #[test]
190    fn fallback_entries_within_threshold_stay_in_one_group() {
191        let entries = vec![
192            make_entry("e1", 100, None),
193            make_entry("e2", 150, None),
194            make_entry("e3", 200, None),
195        ];
196
197        let groups = group_entries_by_session(&entries, 60);
198
199        assert_eq!(groups.len(), 1);
200        assert_group_ids(&groups, &["__fallback__100"]);
201        assert_eq!(groups[0].entries.len(), 3);
202    }
203
204    #[test]
205    fn fallback_entries_with_gap_above_threshold_split_groups() {
206        let entries = vec![
207            make_entry("e1", 100, None),
208            make_entry("e2", 150, None),
209            make_entry("e3", 300, None),
210        ];
211
212        let groups = group_entries_by_session(&entries, 60);
213
214        assert_eq!(groups.len(), 2);
215        assert_group_ids(&groups, &["__fallback__100", "__fallback__300"]);
216        assert_eq!(groups[0].entries.len(), 2);
217        assert_eq!(groups[1].entries.len(), 1);
218    }
219
220    #[test]
221    fn mixed_explicit_and_fallback_entries_do_not_mix() {
222        let entries = vec![
223            make_entry("e1", 100, Some("s1")),
224            make_entry("e2", 110, None),
225            make_entry("e3", 120, Some("s1")),
226            make_entry("e4", 130, None),
227        ];
228
229        let groups = group_entries_by_session(&entries, 30);
230
231        assert_eq!(groups.len(), 2);
232        assert_group_ids(&groups, &["s1", "__fallback__110"]);
233        assert_eq!(groups[0].entries.len(), 2);
234        assert_eq!(groups[1].entries.len(), 2);
235    }
236
237    #[test]
238    fn unsorted_fallback_timestamps_are_sorted_before_grouping() {
239        let entries = vec![
240            make_entry("e3", 210, None),
241            make_entry("e1", 100, None),
242            make_entry("e2", 150, None),
243        ];
244
245        let groups = group_entries_by_session(&entries, 60);
246
247        assert_eq!(groups.len(), 1);
248        assert_group_ids(&groups, &["__fallback__100"]);
249
250        let ordered_ids: Vec<&str> = groups[0]
251            .entries
252            .iter()
253            .map(|entry| entry.id.as_str())
254            .collect();
255        assert_eq!(ordered_ids, vec!["e1", "e2", "e3"]);
256    }
257
258    #[test]
259    fn entries_at_exact_threshold_boundary_stay_together() {
260        let entries = vec![
261            make_entry("e1", 100, None),
262            make_entry("e2", 160, None),
263            make_entry("e3", 220, None),
264        ];
265
266        let groups = group_entries_by_session(&entries, 60);
267
268        assert_eq!(groups.len(), 1);
269        assert_group_ids(&groups, &["__fallback__100"]);
270        assert_eq!(groups[0].entries.len(), 3);
271    }
272
273    #[test]
274    fn fallback_groups_interleave_chronologically_with_explicit_groups() {
275        let entries = vec![
276            make_entry("f1", 50, None),
277            make_entry("e1", 200, Some("s1")),
278            make_entry("f2", 60, None),
279        ];
280
281        let groups = group_entries_by_session(&entries, 60);
282
283        assert_eq!(groups.len(), 2);
284        assert_eq!(groups[0].session_id, "__fallback__50");
285        assert_eq!(groups[1].session_id, "s1");
286    }
287
288    #[test]
289    fn explicit_group_entries_are_sorted_by_timestamp_regression() {
290        let entries = vec![
291            make_entry("e3", 300, Some("s1")),
292            make_entry("e2", 200, Some("s1")),
293            make_entry("e1", 100, Some("s1")),
294        ];
295
296        let groups = group_entries_by_session(&entries, 60);
297
298        assert_eq!(groups.len(), 1);
299        assert_group_ids(&groups, &["s1"]);
300        let ids: Vec<&str> = groups[0]
301            .entries
302            .iter()
303            .map(|entry| entry.id.as_str())
304            .collect();
305        assert_eq!(ids, vec!["e1", "e2", "e3"]);
306    }
307
308    #[test]
309    fn explicit_groups_are_sorted_by_earliest_timestamp() {
310        let entries = vec![
311            make_entry("a1", 200, Some("z")),
312            make_entry("b1", 100, Some("a")),
313        ];
314
315        let groups = group_entries_by_session(&entries, 60);
316
317        assert_eq!(groups.len(), 2);
318        assert_group_ids(&groups, &["a", "z"]);
319    }
320
321    #[test]
322    fn threshold_zero_splits_distinct_timestamps() {
323        let entries = vec![
324            make_entry("e1", 100, None),
325            make_entry("e2", 101, None),
326            make_entry("e3", 102, None),
327        ];
328
329        let groups = group_entries_by_session(&entries, 0);
330
331        assert_eq!(groups.len(), 3);
332        assert_group_ids(
333            &groups,
334            &["__fallback__100", "__fallback__101", "__fallback__102"],
335        );
336        assert_eq!(groups[0].entries.len(), 1);
337        assert_eq!(groups[1].entries.len(), 1);
338        assert_eq!(groups[2].entries.len(), 1);
339    }
340
341    #[test]
342    fn duplicate_timestamps_with_zero_threshold_stay_in_same_group() {
343        let entries = vec![
344            make_entry("e1", 100, None),
345            make_entry("e2", 100, None),
346            make_entry("e3", 101, None),
347        ];
348
349        let groups = group_entries_by_session(&entries, 0);
350
351        assert_eq!(groups.len(), 2);
352        assert_group_ids(&groups, &["__fallback__100", "__fallback__101"]);
353
354        let first_ids: Vec<&str> = groups[0]
355            .entries
356            .iter()
357            .map(|entry| entry.id.as_str())
358            .collect();
359        assert_eq!(first_ids, vec!["e1", "e2"]);
360        assert_eq!(groups[1].entries.len(), 1);
361    }
362}
context_forge/session.rs

context_forge/
session.rs
