use std::collections::{HashMap, VecDeque};
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex, Once, RwLock};
use std::time::{SystemTime, UNIX_EPOCH};
use serde::{Deserialize, Serialize};
use ts_rs::TS;
use crate::language_model_message::LanguageModelPrompt;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, TS)]
#[ts(export)]
pub enum SessionSource {
Explicit,
Inferred,
}
#[derive(Debug, Clone, Serialize, Deserialize, TS)]
#[ts(export)]
pub struct SessionCall {
pub call_id: String,
pub step: u32,
pub recorded_at: String,
}
#[derive(Debug, Clone, Serialize, Deserialize, TS)]
#[ts(export)]
pub struct SessionView {
pub session_id: String,
pub source: SessionSource,
pub calls: Vec<SessionCall>,
}
pub const DEFAULT_MAX_SESSIONS: usize = 256;
pub const DEFAULT_MAX_CALLS_PER_SESSION: usize = 64;
pub struct SessionStore {
inner: Mutex<Inner>,
max_sessions: usize,
max_calls_per_session: usize,
}
#[derive(Default)]
struct Inner {
sessions: HashMap<String, Entry>,
lru: VecDeque<String>,
}
struct Entry {
source: SessionSource,
calls: VecDeque<SessionCall>,
next_step: u64,
}
impl SessionStore {
#[must_use]
pub fn new() -> Self {
Self::with_capacity(DEFAULT_MAX_SESSIONS, DEFAULT_MAX_CALLS_PER_SESSION)
}
#[must_use]
pub fn with_capacity(max_sessions: usize, max_calls_per_session: usize) -> Self {
assert!(
max_sessions > 0 && max_calls_per_session > 0,
"SessionStore bounds must be positive"
);
Self {
inner: Mutex::new(Inner::default()),
max_sessions,
max_calls_per_session,
}
}
pub fn append(
&self,
session_id: &str,
call_id: Option<&str>,
source: SessionSource,
) -> SessionCall {
let mut inner = self.inner.lock().unwrap();
inner.touch(session_id);
if !inner.sessions.contains_key(session_id) {
if inner.sessions.len() >= self.max_sessions
&& let Some(oldest) = inner.lru.pop_front()
{
inner.sessions.remove(&oldest);
}
inner.sessions.insert(
session_id.to_string(),
Entry {
source,
calls: VecDeque::new(),
next_step: 0,
},
);
}
let entry = inner.sessions.get_mut(session_id).unwrap();
let step = u32::try_from(entry.next_step).unwrap_or(u32::MAX);
entry.next_step += 1;
let call = SessionCall {
call_id: call_id.map(str::to_string).unwrap_or_else(new_call_id),
step,
recorded_at: rfc3339_now(),
};
if entry.calls.len() >= self.max_calls_per_session {
entry.calls.pop_front();
}
entry.calls.push_back(call.clone());
call
}
pub fn session_calls(&self, session_id: &str) -> Vec<SessionCall> {
let mut inner = self.inner.lock().unwrap();
if inner.sessions.contains_key(session_id) {
inner.touch(session_id);
}
inner
.sessions
.get(session_id)
.map(|e| e.calls.iter().cloned().collect())
.unwrap_or_default()
}
pub fn list_sessions(&self) -> Vec<SessionView> {
let inner = self.inner.lock().unwrap();
inner
.sessions
.iter()
.map(|(id, e)| SessionView {
session_id: id.clone(),
source: e.source,
calls: e.calls.iter().cloned().collect(),
})
.collect()
}
pub fn clear(&self) {
let mut inner = self.inner.lock().unwrap();
inner.sessions.clear();
inner.lru.clear();
}
}
impl Default for SessionStore {
fn default() -> Self {
Self::new()
}
}
impl Inner {
fn touch(&mut self, id: &str) {
if let Some(pos) = self.lru.iter().position(|k| k == id) {
self.lru.remove(pos);
}
self.lru.push_back(id.to_string());
}
}
pub const DEFAULT_INFERER_CAPACITY: usize = 64;
pub struct SessionInferer {
enabled: bool,
recent: VecDeque<(LanguageModelPrompt, String)>,
capacity: usize,
}
impl SessionInferer {
#[must_use]
pub fn new(enabled: bool) -> Self {
Self::with_capacity(enabled, DEFAULT_INFERER_CAPACITY)
}
#[must_use]
pub fn with_capacity(enabled: bool, capacity: usize) -> Self {
assert!(capacity > 0, "SessionInferer capacity must be positive");
Self {
enabled,
recent: VecDeque::new(),
capacity,
}
}
pub fn resolve(
&mut self,
explicit: Option<&str>,
prompt: &LanguageModelPrompt,
) -> Option<(String, SessionSource)> {
if let Some(id) = explicit {
self.remember(prompt, id);
return Some((id.to_string(), SessionSource::Explicit));
}
if !self.enabled {
return None;
}
let id = match self.find_by_strong_prefix(prompt) {
Some(id) => id,
None => format!("auto-{}", new_call_id()),
};
self.remember(prompt, &id);
Some((id, SessionSource::Inferred))
}
fn find_by_strong_prefix(&self, prompt: &LanguageModelPrompt) -> Option<String> {
for (old, id) in &self.recent {
if prompt.len() >= old.len() && prompt[..old.len()] == old[..] {
return Some(id.clone());
}
}
None
}
fn remember(&mut self, prompt: &LanguageModelPrompt, session_id: &str) {
if self.recent.len() >= self.capacity {
self.recent.pop_back();
}
self.recent
.push_front((prompt.clone(), session_id.to_string()));
}
}
static SESSION_STORE: RwLock<Option<Arc<SessionStore>>> = RwLock::new(None);
static SESSION_INFERER: RwLock<Option<Arc<Mutex<SessionInferer>>>> = RwLock::new(None);
static INFER_ENV_CHECKED: Once = Once::new();
pub fn init_session_store(store: Arc<SessionStore>) {
*SESSION_STORE.write().unwrap() = Some(store);
}
pub fn session_store() -> Option<Arc<SessionStore>> {
SESSION_STORE.read().unwrap().clone()
}
pub fn init_session_infer(enabled: bool) {
*SESSION_INFERER.write().unwrap() = Some(Arc::new(Mutex::new(SessionInferer::new(enabled))));
}
pub fn session_inferer() -> Option<Arc<Mutex<SessionInferer>>> {
SESSION_INFERER.read().unwrap().clone()
}
pub fn ensure_inferer_from_env() {
if INFER_ENV_CHECKED.is_completed() {
return;
}
INFER_ENV_CHECKED.call_once(|| {
if std::env::var("AIMUX_SESSION_INFER").as_deref() == Ok("1") && session_inferer().is_none()
{
init_session_infer(true);
}
});
}
#[must_use]
pub fn resolve_session_id(
explicit: Option<&str>,
prompt: &LanguageModelPrompt,
) -> Option<(String, SessionSource)> {
ensure_inferer_from_env();
if let Some(inferer) = session_inferer() {
return inferer.lock().unwrap().resolve(explicit, prompt);
}
explicit.map(|id| (id.to_string(), SessionSource::Explicit))
}
#[must_use]
pub fn session_calls(session_id: &str) -> Vec<SessionCall> {
session_store()
.map(|s| s.session_calls(session_id))
.unwrap_or_default()
}
#[must_use]
pub fn list_sessions() -> Vec<SessionView> {
session_store()
.map(|s| s.list_sessions())
.unwrap_or_default()
}
static CALL_SEQ: AtomicU64 = AtomicU64::new(0);
fn new_call_id() -> String {
let ns = SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|d| d.as_nanos())
.unwrap_or(0);
format!("call-{ns}-{}", CALL_SEQ.fetch_add(1, Ordering::Relaxed))
}
fn rfc3339_now() -> String {
let d = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default();
let secs = d.as_secs() as i64;
let millis = d.subsec_millis();
let (year, month, day) = civil_from_days(secs.div_euclid(86_400));
let hms = secs.rem_euclid(86_400);
format!(
"{year:04}-{month:02}-{day:02}T{:02}:{:02}:{:02}.{millis:03}Z",
hms / 3600,
(hms % 3600) / 60,
hms % 60
)
}
fn civil_from_days(z: i64) -> (i64, u32, u32) {
let z = z + 719_468;
let era = (if z >= 0 { z } else { z - 146_096 }) / 146_097;
let doe = (z - era * 146_097) as u64;
let yoe = (doe - doe / 1460 + doe / 36_524 - doe / 146_096) / 365;
let y = yoe as i64 + era * 400;
let doy = doe - (365 * yoe + yoe / 4 - yoe / 100);
let mp = (5 * doy + 2) / 153;
let d = (doy - (153 * mp + 2) / 5 + 1) as u32;
let m = if mp < 10 { mp + 3 } else { mp - 9 } as u32;
(if m <= 2 { y + 1 } else { y }, m, d)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::content::ContentPart;
use crate::language_model_message::LanguageModelPromptMessage;
use crate::message::Role;
fn msg(text: &str) -> LanguageModelPromptMessage {
LanguageModelPromptMessage {
role: Role::User,
content: vec![ContentPart::Text {
text: text.to_string(),
provider_options: None,
}],
provider_options: None,
}
}
fn prompt(messages: Vec<LanguageModelPromptMessage>) -> LanguageModelPrompt {
messages
}
fn is_rfc3339(s: &str) -> bool {
let b = s.as_bytes();
s.len() >= 24
&& b[4] == b'-'
&& b[7] == b'-'
&& b[10] == b'T'
&& b[13] == b':'
&& b[16] == b':'
&& b[19] == b'.'
&& s.ends_with('Z')
}
#[test]
fn append_increments_step_and_keeps_order() {
let store = SessionStore::new();
store.append("s1", None, SessionSource::Explicit);
store.append("s1", None, SessionSource::Explicit);
store.append("s1", None, SessionSource::Explicit);
let calls = store.session_calls("s1");
assert_eq!(calls.len(), 3);
assert_eq!(calls[0].step, 0);
assert_eq!(calls[1].step, 1);
assert_eq!(calls[2].step, 2);
let mut ids: Vec<_> = calls.iter().map(|c| c.call_id.as_str()).collect();
ids.sort_unstable();
ids.dedup();
assert_eq!(ids.len(), 3);
for c in &calls {
assert!(
is_rfc3339(&c.recorded_at),
"bad timestamp: {}",
c.recorded_at
);
}
}
#[test]
fn unknown_session_is_empty_and_sessions_list_is_complete() {
let store = SessionStore::new();
assert!(store.session_calls("nope").is_empty());
assert!(store.list_sessions().is_empty());
store.append("s1", None, SessionSource::Explicit);
store.append("s2", None, SessionSource::Inferred);
let views = store.list_sessions();
assert_eq!(views.len(), 2);
let s1 = views.iter().find(|v| v.session_id == "s1").unwrap();
assert_eq!(s1.source, SessionSource::Explicit);
assert_eq!(s1.calls.len(), 1);
let s2 = views.iter().find(|v| v.session_id == "s2").unwrap();
assert_eq!(s2.source, SessionSource::Inferred);
}
#[test]
fn lru_evicts_oldest_session_when_full() {
let store = SessionStore::with_capacity(2, 8);
store.append("a", None, SessionSource::Explicit);
store.append("b", None, SessionSource::Explicit);
store.session_calls("a");
store.append("c", None, SessionSource::Explicit);
assert!(
store.session_calls("b").is_empty(),
"oldest session evicted"
);
assert_eq!(store.session_calls("a").len(), 1);
assert_eq!(store.session_calls("c").len(), 1);
}
#[test]
fn calls_per_session_are_bounded() {
let store = SessionStore::with_capacity(8, 2);
store.append("s1", None, SessionSource::Explicit);
store.append("s1", None, SessionSource::Explicit);
store.append("s1", None, SessionSource::Explicit);
let calls = store.session_calls("s1");
assert_eq!(calls.len(), 2, "oldest call dropped");
assert_eq!(calls[0].step, 1);
assert_eq!(calls[1].step, 2);
}
#[test]
fn steps_are_monotonic_across_call_eviction() {
let store = SessionStore::with_capacity(8, 2);
for _ in 0..4 {
store.append("s1", None, SessionSource::Explicit);
}
let calls = store.session_calls("s1");
assert_eq!(calls.len(), 2, "only the two newest calls are retained");
assert_eq!(calls[0].step, 2);
assert_eq!(calls[1].step, 3);
}
#[test]
fn querying_unknown_sessions_does_not_break_session_capacity() {
let store = SessionStore::with_capacity(2, 8);
store.append("a", None, SessionSource::Explicit);
store.append("b", None, SessionSource::Explicit);
for i in 0..16 {
assert!(store.session_calls(&format!("unknown-{i}")).is_empty());
}
for id in ["c", "d", "e", "f", "g"] {
store.append(id, None, SessionSource::Explicit);
}
assert!(
store.list_sessions().len() <= 2,
"session count must stay within max_sessions, got {}",
store.list_sessions().len()
);
assert_eq!(store.session_calls("a").len(), 0, "a was evicted");
assert_eq!(store.session_calls("g").len(), 1);
}
#[test]
fn clear_drops_all_sessions() {
let store = SessionStore::new();
store.append("s1", None, SessionSource::Explicit);
store.clear();
assert!(store.list_sessions().is_empty());
}
#[test]
fn inferer_disabled_returns_none() {
let mut inf = SessionInferer::new(false);
assert!(inf.resolve(None, &prompt(vec![msg("hi")])).is_none());
}
#[test]
fn inferer_explicit_wins_and_is_remembered() {
let mut inf = SessionInferer::new(true);
let p1 = prompt(vec![msg("hi")]);
let (id, source) = inf.resolve(Some("sess-1"), &p1).unwrap();
assert_eq!(id, "sess-1");
assert_eq!(source, SessionSource::Explicit);
let p2 = prompt(vec![msg("hi"), msg("hello")]);
let (id2, source2) = inf.resolve(None, &p2).unwrap();
assert_eq!(id2, "sess-1");
assert_eq!(source2, SessionSource::Inferred);
}
#[test]
fn inferer_strong_prefix_continuation_groups_calls() {
let mut inf = SessionInferer::new(true);
let p1 = prompt(vec![msg("u1")]);
let (id1, source1) = inf.resolve(None, &p1).unwrap();
assert!(id1.starts_with("auto-"));
assert_eq!(source1, SessionSource::Inferred);
let p2 = prompt(vec![msg("u1"), msg("a1"), msg("u2")]);
let (id2, _) = inf.resolve(None, &p2).unwrap();
assert_eq!(id2, id1);
let (id3, _) = inf.resolve(None, &p1).unwrap();
assert_eq!(id3, id1);
}
#[test]
fn inferer_distinct_prompts_start_new_sessions() {
let mut inf = SessionInferer::new(true);
let (id1, _) = inf
.resolve(None, &prompt(vec![msg("what is rust")]))
.unwrap();
let (id2, _) = inf
.resolve(None, &prompt(vec![msg("tell me a joke")]))
.unwrap();
assert_ne!(id1, id2);
}
#[test]
fn inferer_prefix_is_directional() {
let mut inf = SessionInferer::new(true);
let (id1, _) = inf.resolve(None, &prompt(vec![msg("abc")])).unwrap();
let (id2, _) = inf
.resolve(None, &prompt(vec![msg("abc"), msg("a1"), msg("u2")]))
.unwrap();
assert_eq!(id2, id1);
let (id3, _) = inf.resolve(None, &prompt(vec![msg("abd")])).unwrap();
assert_ne!(id3, id1, "abd is not a prefix continuation");
}
#[test]
fn inferer_chain_continuation_stays_in_session() {
let mut inf = SessionInferer::new(true);
let p1 = prompt(vec![msg("common")]);
let p2 = prompt(vec![msg("common"), msg("b")]);
let p3 = prompt(vec![msg("common"), msg("b"), msg("c")]);
let (id1, _) = inf.resolve(None, &p1).unwrap();
let (id2, _) = inf.resolve(None, &p2).unwrap();
let (id3, _) = inf.resolve(None, &p3).unwrap();
assert_eq!(id2, id1);
assert_eq!(id3, id1);
}
#[test]
fn inferer_capacity_evicts_oldest_prompt() {
let mut inf = SessionInferer::with_capacity(true, 2);
let p1 = prompt(vec![msg("one")]);
let (id1, _) = inf.resolve(None, &p1).unwrap();
let p2 = prompt(vec![msg("two")]);
let (id2, _) = inf.resolve(None, &p2).unwrap();
let p3 = prompt(vec![msg("three")]);
let (id3, _) = inf.resolve(None, &p3).unwrap();
assert_ne!(id1, id2);
assert_ne!(id2, id3);
let p1_ext = prompt(vec![msg("one"), msg("more")]);
let (id4, _) = inf.resolve(None, &p1_ext).unwrap();
assert_ne!(id4, id1, "evicted prompt no longer matches");
}
#[test]
fn rfc3339_now_is_well_formed() {
let s = rfc3339_now();
assert!(is_rfc3339(&s), "bad timestamp: {s}");
}
}