1use std::collections::{HashMap, VecDeque};
21use std::sync::atomic::{AtomicU64, Ordering};
22use std::sync::{Arc, Mutex, Once, RwLock};
23use std::time::{SystemTime, UNIX_EPOCH};
24
25use serde::{Deserialize, Serialize};
26use ts_rs::TS;
27
28use crate::language_model_message::LanguageModelPrompt;
29
30#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, TS)]
36#[ts(export)]
37pub enum SessionSource {
38 Explicit,
40 Inferred,
42}
43
44#[derive(Debug, Clone, Serialize, Deserialize, TS)]
47#[ts(export)]
48pub struct SessionCall {
49 pub call_id: String,
53 pub step: u32,
55 pub recorded_at: String,
57}
58
59#[derive(Debug, Clone, Serialize, Deserialize, TS)]
61#[ts(export)]
62pub struct SessionView {
63 pub session_id: String,
64 pub source: SessionSource,
66 pub calls: Vec<SessionCall>,
68}
69
70pub const DEFAULT_MAX_SESSIONS: usize = 256;
76pub const DEFAULT_MAX_CALLS_PER_SESSION: usize = 64;
78
79pub struct SessionStore {
83 inner: Mutex<Inner>,
84 max_sessions: usize,
85 max_calls_per_session: usize,
86}
87
88#[derive(Default)]
89struct Inner {
90 sessions: HashMap<String, Entry>,
91 lru: VecDeque<String>,
93}
94
95struct Entry {
96 source: SessionSource,
97 calls: VecDeque<SessionCall>,
98 next_step: u64,
100}
101
102impl SessionStore {
103 #[must_use]
105 pub fn new() -> Self {
106 Self::with_capacity(DEFAULT_MAX_SESSIONS, DEFAULT_MAX_CALLS_PER_SESSION)
107 }
108
109 #[must_use]
111 pub fn with_capacity(max_sessions: usize, max_calls_per_session: usize) -> Self {
112 assert!(
113 max_sessions > 0 && max_calls_per_session > 0,
114 "SessionStore bounds must be positive"
115 );
116 Self {
117 inner: Mutex::new(Inner::default()),
118 max_sessions,
119 max_calls_per_session,
120 }
121 }
122
123 pub fn append(
130 &self,
131 session_id: &str,
132 call_id: Option<&str>,
133 source: SessionSource,
134 ) -> SessionCall {
135 let mut inner = self.inner.lock().unwrap();
136 inner.touch(session_id);
137 if !inner.sessions.contains_key(session_id) {
138 if inner.sessions.len() >= self.max_sessions
140 && let Some(oldest) = inner.lru.pop_front()
141 {
142 inner.sessions.remove(&oldest);
143 }
144 inner.sessions.insert(
145 session_id.to_string(),
146 Entry {
147 source,
148 calls: VecDeque::new(),
149 next_step: 0,
150 },
151 );
152 }
153 let entry = inner.sessions.get_mut(session_id).unwrap();
154 let step = u32::try_from(entry.next_step).unwrap_or(u32::MAX);
159 entry.next_step += 1;
160 let call = SessionCall {
161 call_id: call_id.map(str::to_string).unwrap_or_else(new_call_id),
162 step,
163 recorded_at: rfc3339_now(),
164 };
165 if entry.calls.len() >= self.max_calls_per_session {
167 entry.calls.pop_front();
168 }
169 entry.calls.push_back(call.clone());
170 call
171 }
172
173 pub fn session_calls(&self, session_id: &str) -> Vec<SessionCall> {
175 let mut inner = self.inner.lock().unwrap();
176 if inner.sessions.contains_key(session_id) {
179 inner.touch(session_id);
180 }
181 inner
182 .sessions
183 .get(session_id)
184 .map(|e| e.calls.iter().cloned().collect())
185 .unwrap_or_default()
186 }
187
188 pub fn list_sessions(&self) -> Vec<SessionView> {
190 let inner = self.inner.lock().unwrap();
191 inner
192 .sessions
193 .iter()
194 .map(|(id, e)| SessionView {
195 session_id: id.clone(),
196 source: e.source,
197 calls: e.calls.iter().cloned().collect(),
198 })
199 .collect()
200 }
201
202 pub fn clear(&self) {
204 let mut inner = self.inner.lock().unwrap();
205 inner.sessions.clear();
206 inner.lru.clear();
207 }
208}
209
210impl Default for SessionStore {
211 fn default() -> Self {
212 Self::new()
213 }
214}
215
216impl Inner {
217 fn touch(&mut self, id: &str) {
219 if let Some(pos) = self.lru.iter().position(|k| k == id) {
220 self.lru.remove(pos);
221 }
222 self.lru.push_back(id.to_string());
223 }
224}
225
226pub const DEFAULT_INFERER_CAPACITY: usize = 64;
232
233pub struct SessionInferer {
241 enabled: bool,
242 recent: VecDeque<(LanguageModelPrompt, String)>,
244 capacity: usize,
245}
246
247impl SessionInferer {
248 #[must_use]
250 pub fn new(enabled: bool) -> Self {
251 Self::with_capacity(enabled, DEFAULT_INFERER_CAPACITY)
252 }
253
254 #[must_use]
257 pub fn with_capacity(enabled: bool, capacity: usize) -> Self {
258 assert!(capacity > 0, "SessionInferer capacity must be positive");
259 Self {
260 enabled,
261 recent: VecDeque::new(),
262 capacity,
263 }
264 }
265
266 pub fn resolve(
274 &mut self,
275 explicit: Option<&str>,
276 prompt: &LanguageModelPrompt,
277 ) -> Option<(String, SessionSource)> {
278 if let Some(id) = explicit {
279 self.remember(prompt, id);
280 return Some((id.to_string(), SessionSource::Explicit));
281 }
282 if !self.enabled {
283 return None;
284 }
285 let id = match self.find_by_strong_prefix(prompt) {
286 Some(id) => id,
287 None => format!("auto-{}", new_call_id()),
288 };
289 self.remember(prompt, &id);
290 Some((id, SessionSource::Inferred))
291 }
292
293 fn find_by_strong_prefix(&self, prompt: &LanguageModelPrompt) -> Option<String> {
297 for (old, id) in &self.recent {
298 if prompt.len() >= old.len() && prompt[..old.len()] == old[..] {
299 return Some(id.clone());
300 }
301 }
302 None
303 }
304
305 fn remember(&mut self, prompt: &LanguageModelPrompt, session_id: &str) {
306 if self.recent.len() >= self.capacity {
307 self.recent.pop_back();
308 }
309 self.recent
310 .push_front((prompt.clone(), session_id.to_string()));
311 }
312}
313
314static SESSION_STORE: RwLock<Option<Arc<SessionStore>>> = RwLock::new(None);
319static SESSION_INFERER: RwLock<Option<Arc<Mutex<SessionInferer>>>> = RwLock::new(None);
320static INFER_ENV_CHECKED: Once = Once::new();
321
322pub fn init_session_store(store: Arc<SessionStore>) {
327 *SESSION_STORE.write().unwrap() = Some(store);
328}
329
330pub fn session_store() -> Option<Arc<SessionStore>> {
332 SESSION_STORE.read().unwrap().clone()
333}
334
335pub fn init_session_infer(enabled: bool) {
340 *SESSION_INFERER.write().unwrap() = Some(Arc::new(Mutex::new(SessionInferer::new(enabled))));
341}
342
343pub fn session_inferer() -> Option<Arc<Mutex<SessionInferer>>> {
345 SESSION_INFERER.read().unwrap().clone()
346}
347
348pub fn ensure_inferer_from_env() {
352 if INFER_ENV_CHECKED.is_completed() {
353 return;
354 }
355 INFER_ENV_CHECKED.call_once(|| {
356 if std::env::var("AIMUX_SESSION_INFER").as_deref() == Ok("1") && session_inferer().is_none()
357 {
358 init_session_infer(true);
359 }
360 });
361}
362
363#[must_use]
368pub fn resolve_session_id(
369 explicit: Option<&str>,
370 prompt: &LanguageModelPrompt,
371) -> Option<(String, SessionSource)> {
372 ensure_inferer_from_env();
373 if let Some(inferer) = session_inferer() {
374 return inferer.lock().unwrap().resolve(explicit, prompt);
375 }
376 explicit.map(|id| (id.to_string(), SessionSource::Explicit))
377}
378
379#[must_use]
382pub fn session_calls(session_id: &str) -> Vec<SessionCall> {
383 session_store()
384 .map(|s| s.session_calls(session_id))
385 .unwrap_or_default()
386}
387
388#[must_use]
390pub fn list_sessions() -> Vec<SessionView> {
391 session_store()
392 .map(|s| s.list_sessions())
393 .unwrap_or_default()
394}
395
396static CALL_SEQ: AtomicU64 = AtomicU64::new(0);
405
406fn new_call_id() -> String {
408 let ns = SystemTime::now()
409 .duration_since(UNIX_EPOCH)
410 .map(|d| d.as_nanos())
411 .unwrap_or(0);
412 format!("call-{ns}-{}", CALL_SEQ.fetch_add(1, Ordering::Relaxed))
413}
414
415fn rfc3339_now() -> String {
418 let d = SystemTime::now()
419 .duration_since(UNIX_EPOCH)
420 .unwrap_or_default();
421 let secs = d.as_secs() as i64;
422 let millis = d.subsec_millis();
423 let (year, month, day) = civil_from_days(secs.div_euclid(86_400));
424 let hms = secs.rem_euclid(86_400);
425 format!(
426 "{year:04}-{month:02}-{day:02}T{:02}:{:02}:{:02}.{millis:03}Z",
427 hms / 3600,
428 (hms % 3600) / 60,
429 hms % 60
430 )
431}
432
433fn civil_from_days(z: i64) -> (i64, u32, u32) {
436 let z = z + 719_468;
437 let era = (if z >= 0 { z } else { z - 146_096 }) / 146_097;
438 let doe = (z - era * 146_097) as u64;
439 let yoe = (doe - doe / 1460 + doe / 36_524 - doe / 146_096) / 365;
440 let y = yoe as i64 + era * 400;
441 let doy = doe - (365 * yoe + yoe / 4 - yoe / 100);
442 let mp = (5 * doy + 2) / 153;
443 let d = (doy - (153 * mp + 2) / 5 + 1) as u32;
444 let m = if mp < 10 { mp + 3 } else { mp - 9 } as u32;
445 (if m <= 2 { y + 1 } else { y }, m, d)
446}
447
448#[cfg(test)]
453mod tests {
454 use super::*;
455 use crate::content::ContentPart;
456 use crate::language_model_message::LanguageModelPromptMessage;
457 use crate::message::Role;
458
459 fn msg(text: &str) -> LanguageModelPromptMessage {
460 LanguageModelPromptMessage {
461 role: Role::User,
462 content: vec![ContentPart::Text {
463 text: text.to_string(),
464 provider_options: None,
465 }],
466 provider_options: None,
467 }
468 }
469
470 fn prompt(messages: Vec<LanguageModelPromptMessage>) -> LanguageModelPrompt {
471 messages
472 }
473
474 fn is_rfc3339(s: &str) -> bool {
476 let b = s.as_bytes();
477 s.len() >= 24
478 && b[4] == b'-'
479 && b[7] == b'-'
480 && b[10] == b'T'
481 && b[13] == b':'
482 && b[16] == b':'
483 && b[19] == b'.'
484 && s.ends_with('Z')
485 }
486
487 #[test]
488 fn append_increments_step_and_keeps_order() {
489 let store = SessionStore::new();
490 store.append("s1", None, SessionSource::Explicit);
491 store.append("s1", None, SessionSource::Explicit);
492 store.append("s1", None, SessionSource::Explicit);
493
494 let calls = store.session_calls("s1");
495 assert_eq!(calls.len(), 3);
496 assert_eq!(calls[0].step, 0);
497 assert_eq!(calls[1].step, 1);
498 assert_eq!(calls[2].step, 2);
499 let mut ids: Vec<_> = calls.iter().map(|c| c.call_id.as_str()).collect();
501 ids.sort_unstable();
502 ids.dedup();
503 assert_eq!(ids.len(), 3);
504 for c in &calls {
506 assert!(
507 is_rfc3339(&c.recorded_at),
508 "bad timestamp: {}",
509 c.recorded_at
510 );
511 }
512 }
513
514 #[test]
515 fn unknown_session_is_empty_and_sessions_list_is_complete() {
516 let store = SessionStore::new();
517 assert!(store.session_calls("nope").is_empty());
518 assert!(store.list_sessions().is_empty());
519
520 store.append("s1", None, SessionSource::Explicit);
521 store.append("s2", None, SessionSource::Inferred);
522
523 let views = store.list_sessions();
524 assert_eq!(views.len(), 2);
525 let s1 = views.iter().find(|v| v.session_id == "s1").unwrap();
526 assert_eq!(s1.source, SessionSource::Explicit);
527 assert_eq!(s1.calls.len(), 1);
528 let s2 = views.iter().find(|v| v.session_id == "s2").unwrap();
529 assert_eq!(s2.source, SessionSource::Inferred);
530 }
531
532 #[test]
533 fn lru_evicts_oldest_session_when_full() {
534 let store = SessionStore::with_capacity(2, 8);
535 store.append("a", None, SessionSource::Explicit);
536 store.append("b", None, SessionSource::Explicit);
537 store.session_calls("a");
539 store.append("c", None, SessionSource::Explicit);
540
541 assert!(
542 store.session_calls("b").is_empty(),
543 "oldest session evicted"
544 );
545 assert_eq!(store.session_calls("a").len(), 1);
546 assert_eq!(store.session_calls("c").len(), 1);
547 }
548
549 #[test]
550 fn calls_per_session_are_bounded() {
551 let store = SessionStore::with_capacity(8, 2);
552 store.append("s1", None, SessionSource::Explicit);
553 store.append("s1", None, SessionSource::Explicit);
554 store.append("s1", None, SessionSource::Explicit);
555
556 let calls = store.session_calls("s1");
557 assert_eq!(calls.len(), 2, "oldest call dropped");
558 assert_eq!(calls[0].step, 1);
559 assert_eq!(calls[1].step, 2);
560 }
561
562 #[test]
563 fn steps_are_monotonic_across_call_eviction() {
564 let store = SessionStore::with_capacity(8, 2);
567 for _ in 0..4 {
568 store.append("s1", None, SessionSource::Explicit);
569 }
570 let calls = store.session_calls("s1");
571 assert_eq!(calls.len(), 2, "only the two newest calls are retained");
572 assert_eq!(calls[0].step, 2);
573 assert_eq!(calls[1].step, 3);
574 }
575
576 #[test]
577 fn querying_unknown_sessions_does_not_break_session_capacity() {
578 let store = SessionStore::with_capacity(2, 8);
581 store.append("a", None, SessionSource::Explicit);
582 store.append("b", None, SessionSource::Explicit);
583
584 for i in 0..16 {
586 assert!(store.session_calls(&format!("unknown-{i}")).is_empty());
587 }
588
589 for id in ["c", "d", "e", "f", "g"] {
590 store.append(id, None, SessionSource::Explicit);
591 }
592 assert!(
593 store.list_sessions().len() <= 2,
594 "session count must stay within max_sessions, got {}",
595 store.list_sessions().len()
596 );
597 assert_eq!(store.session_calls("a").len(), 0, "a was evicted");
598 assert_eq!(store.session_calls("g").len(), 1);
599 }
600
601 #[test]
602 fn clear_drops_all_sessions() {
603 let store = SessionStore::new();
604 store.append("s1", None, SessionSource::Explicit);
605 store.clear();
606 assert!(store.list_sessions().is_empty());
607 }
608
609 #[test]
612 fn inferer_disabled_returns_none() {
613 let mut inf = SessionInferer::new(false);
614 assert!(inf.resolve(None, &prompt(vec![msg("hi")])).is_none());
615 }
616
617 #[test]
618 fn inferer_explicit_wins_and_is_remembered() {
619 let mut inf = SessionInferer::new(true);
620 let p1 = prompt(vec![msg("hi")]);
621 let (id, source) = inf.resolve(Some("sess-1"), &p1).unwrap();
622 assert_eq!(id, "sess-1");
623 assert_eq!(source, SessionSource::Explicit);
624
625 let p2 = prompt(vec![msg("hi"), msg("hello")]);
627 let (id2, source2) = inf.resolve(None, &p2).unwrap();
628 assert_eq!(id2, "sess-1");
629 assert_eq!(source2, SessionSource::Inferred);
630 }
631
632 #[test]
633 fn inferer_strong_prefix_continuation_groups_calls() {
634 let mut inf = SessionInferer::new(true);
635 let p1 = prompt(vec![msg("u1")]);
636 let (id1, source1) = inf.resolve(None, &p1).unwrap();
637 assert!(id1.starts_with("auto-"));
638 assert_eq!(source1, SessionSource::Inferred);
639
640 let p2 = prompt(vec![msg("u1"), msg("a1"), msg("u2")]);
642 let (id2, _) = inf.resolve(None, &p2).unwrap();
643 assert_eq!(id2, id1);
644
645 let (id3, _) = inf.resolve(None, &p1).unwrap();
647 assert_eq!(id3, id1);
648 }
649
650 #[test]
651 fn inferer_distinct_prompts_start_new_sessions() {
652 let mut inf = SessionInferer::new(true);
653 let (id1, _) = inf
654 .resolve(None, &prompt(vec![msg("what is rust")]))
655 .unwrap();
656 let (id2, _) = inf
657 .resolve(None, &prompt(vec![msg("tell me a joke")]))
658 .unwrap();
659 assert_ne!(id1, id2);
660 }
661
662 #[test]
663 fn inferer_prefix_is_directional() {
664 let mut inf = SessionInferer::new(true);
665 let (id1, _) = inf.resolve(None, &prompt(vec![msg("abc")])).unwrap();
666 let (id2, _) = inf
669 .resolve(None, &prompt(vec![msg("abc"), msg("a1"), msg("u2")]))
670 .unwrap();
671 assert_eq!(id2, id1);
672 let (id3, _) = inf.resolve(None, &prompt(vec![msg("abd")])).unwrap();
674 assert_ne!(id3, id1, "abd is not a prefix continuation");
675 }
676
677 #[test]
678 fn inferer_chain_continuation_stays_in_session() {
679 let mut inf = SessionInferer::new(true);
680 let p1 = prompt(vec![msg("common")]);
681 let p2 = prompt(vec![msg("common"), msg("b")]);
682 let p3 = prompt(vec![msg("common"), msg("b"), msg("c")]);
683 let (id1, _) = inf.resolve(None, &p1).unwrap();
684 let (id2, _) = inf.resolve(None, &p2).unwrap();
685 let (id3, _) = inf.resolve(None, &p3).unwrap();
686 assert_eq!(id2, id1);
687 assert_eq!(id3, id1);
688 }
689
690 #[test]
691 fn inferer_capacity_evicts_oldest_prompt() {
692 let mut inf = SessionInferer::with_capacity(true, 2);
693 let p1 = prompt(vec![msg("one")]);
694 let (id1, _) = inf.resolve(None, &p1).unwrap();
695 let p2 = prompt(vec![msg("two")]);
696 let (id2, _) = inf.resolve(None, &p2).unwrap();
697 let p3 = prompt(vec![msg("three")]);
698 let (id3, _) = inf.resolve(None, &p3).unwrap();
699 assert_ne!(id1, id2);
700 assert_ne!(id2, id3);
701
702 let p1_ext = prompt(vec![msg("one"), msg("more")]);
704 let (id4, _) = inf.resolve(None, &p1_ext).unwrap();
705 assert_ne!(id4, id1, "evicted prompt no longer matches");
706 }
707
708 #[test]
711 fn rfc3339_now_is_well_formed() {
712 let s = rfc3339_now();
713 assert!(is_rfc3339(&s), "bad timestamp: {s}");
714 }
715}