claude_codex/
session.rs

1use crate::config::AliasProvider;
2use crate::registry::normalize_incoming_model;
3use std::collections::{HashMap, VecDeque};
4use std::sync::{LazyLock, Mutex};
5
6const SESSION_IDLE_TTL_MS: u64 = 30 * 60 * 1000;
7pub const MAX_SESSIONS: usize = 10_000;
8
9#[derive(Debug, Clone)]
10pub struct SessionState {
11    pub seq: u64,
12    pub affinity_provider: Option<AliasProvider>,
13    pub last_seen: u64,
14}
15
16#[derive(Default)]
17struct SessionStore {
18    map: HashMap<String, SessionState>,
19    order: VecDeque<String>,
20}
21
22static SESSIONS: LazyLock<Mutex<SessionStore>> =
23    LazyLock::new(|| Mutex::new(SessionStore::default()));
24
25fn now_millis() -> u64 {
26    use std::time::{SystemTime, UNIX_EPOCH};
27    let dur = SystemTime::now()
28        .duration_since(UNIX_EPOCH)
29        .unwrap_or_default();
30    dur.as_millis() as u64
31}
32
33pub fn existing_session(session_id: Option<&str>, now: u64) -> Option<SessionState> {
34    let id = session_id?;
35    let mut store = SESSIONS.lock().expect("session lock");
36    let state = store.map.get(id).cloned()?;
37    if now.saturating_sub(state.last_seen) > SESSION_IDLE_TTL_MS {
38        store.map.remove(id);
39        store.order.retain(|item| item != id);
40        return None;
41    }
42    Some(state)
43}
44
45pub fn existing_session_now(session_id: Option<&str>) -> Option<SessionState> {
46    existing_session(session_id, now_millis())
47}
48
49pub fn record_session_request(
50    session_id: Option<&str>,
51    prior: Option<&SessionState>,
52    provider_name: &str,
53    model: &str,
54    now: u64,
55) -> Option<SessionState> {
56    let id = session_id?;
57    let mut store = SESSIONS.lock().expect("session lock");
58    let mut next = prior.cloned().unwrap_or(SessionState {
59        seq: 0,
60        affinity_provider: None,
61        last_seen: now,
62    });
63    next.seq += 1;
64    next.last_seen = now;
65    if is_alias_routable_provider(provider_name)
66        && !crate::registry::is_anthropic_alias(normalize_incoming_model(model).as_str())
67    {
68        next.affinity_provider = Some(match provider_name {
69            "codex" => AliasProvider::Codex,
70            "kimi" => AliasProvider::Kimi,
71            _ => next.affinity_provider.unwrap_or(AliasProvider::Codex),
72        });
73    }
74
75    if !store.map.contains_key(id) {
76        store.order.push_back(id.to_string());
77    }
78    store.map.insert(id.to_string(), next.clone());
79
80    while store.order.len() > MAX_SESSIONS {
81        if let Some(evict) = store.order.pop_front() {
82            store.map.remove(&evict);
83        } else {
84            break;
85        }
86    }
87
88    Some(next)
89}
90
91fn is_alias_routable_provider(name: &str) -> bool {
92    matches!(name, "codex" | "kimi")
93}
94
95#[cfg(test)]
96pub fn reset_sessions_for_test() {
97    let mut store = SESSIONS.lock().expect("session lock");
98    store.map.clear();
99    store.order.clear();
100}
101
102pub fn affinity_provider_from_session(session: &SessionState) -> Option<AliasProvider> {
103    session.affinity_provider
104}
claude_codex/session.rs

claude_codex/
session.rs
