use macp_core::session::Session;
use std::collections::{BinaryHeap, HashMap};
use std::fs;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use tokio::sync::RwLock;
#[derive(serde::Serialize, serde::Deserialize)]
pub struct PersistedRoot {
pub uri: String,
pub name: String,
}
#[derive(serde::Serialize, serde::Deserialize)]
pub struct PersistedSession {
#[serde(default = "default_schema_version")]
pub schema_version: u32,
pub session_id: String,
pub state: macp_core::session::SessionState,
pub ttl_expiry: i64,
#[serde(default)]
pub ttl_ms: i64,
pub started_at_unix_ms: i64,
pub resolution: Option<Vec<u8>>,
pub mode: String,
pub mode_state: Vec<u8>,
pub participants: Vec<String>,
pub seen_message_ids: Vec<String>,
pub intent: String,
pub mode_version: String,
pub configuration_version: String,
pub policy_version: String,
#[serde(default)]
pub context_id: String,
#[serde(default)]
pub extensions: HashMap<String, Vec<u8>>,
pub roots: Vec<PersistedRoot>,
pub initiator_sender: String,
#[serde(default)]
pub policy_definition: Option<macp_core::policy::PolicyDefinition>,
#[serde(default)]
pub suspended_at_ms: Option<i64>,
#[serde(default)]
pub accumulated_suspended_ms: i64,
#[serde(default)]
pub semantics_rev: u32,
#[serde(default)]
pub max_suspend_ms: i64,
}
fn default_schema_version() -> u32 {
2
}
impl From<&Session> for PersistedSession {
fn from(session: &Session) -> Self {
Self {
schema_version: 2,
session_id: session.session_id.clone(),
state: session.state.clone(),
ttl_expiry: session.ttl_expiry,
ttl_ms: session.ttl_ms,
started_at_unix_ms: session.started_at_unix_ms,
resolution: session.resolution.clone(),
mode: session.mode.clone(),
mode_state: session.mode_state.clone(),
participants: session.participants.clone(),
seen_message_ids: session.seen_message_ids.iter().cloned().collect(),
intent: session.intent.clone(),
mode_version: session.mode_version.clone(),
configuration_version: session.configuration_version.clone(),
policy_version: session.policy_version.clone(),
context_id: session.context_id.clone(),
extensions: session.extensions.clone(),
roots: session
.roots
.iter()
.map(|root| PersistedRoot {
uri: root.uri.clone(),
name: root.name.clone(),
})
.collect(),
initiator_sender: session.initiator_sender.clone(),
policy_definition: session.policy_definition.clone(),
suspended_at_ms: session.suspended_at_ms,
accumulated_suspended_ms: session.accumulated_suspended_ms,
semantics_rev: session.semantics_rev,
max_suspend_ms: session.max_suspend_ms,
}
}
}
impl From<PersistedSession> for Session {
fn from(session: PersistedSession) -> Self {
let ttl_ms = if session.ttl_ms > 0 {
session.ttl_ms
} else {
session
.ttl_expiry
.saturating_sub(session.started_at_unix_ms)
};
Session::builder(session.session_id, session.mode, session.initiator_sender)
.state(session.state)
.ttl_expiry(session.ttl_expiry)
.ttl_ms(ttl_ms)
.started_at_unix_ms(session.started_at_unix_ms)
.resolution(session.resolution)
.mode_state(session.mode_state)
.participants(session.participants)
.seen_message_ids(session.seen_message_ids.into_iter().collect())
.intent(session.intent)
.mode_version(session.mode_version)
.configuration_version(session.configuration_version)
.policy_version(session.policy_version)
.context_id(session.context_id)
.extensions(session.extensions)
.roots(
session
.roots
.into_iter()
.map(|root| macp_pb::pb::Root {
uri: root.uri,
name: root.name,
})
.collect(),
)
.policy_definition(session.policy_definition)
.suspended_at_ms(session.suspended_at_ms)
.accumulated_suspended_ms(session.accumulated_suspended_ms)
.semantics_rev(session.semantics_rev)
.max_suspend_ms(session.max_suspend_ms)
.build()
}
}
pub type SharedSession = Arc<tokio::sync::Mutex<Session>>;
pub struct SessionRegistry {
pub sessions: RwLock<HashMap<String, SharedSession>>,
persistence_path: Option<PathBuf>,
}
impl Default for SessionRegistry {
fn default() -> Self {
Self::new()
}
}
impl SessionRegistry {
pub fn new() -> Self {
Self {
sessions: RwLock::new(HashMap::new()),
persistence_path: None,
}
}
pub fn with_persistence<P: AsRef<Path>>(dir: P) -> std::io::Result<Self> {
let dir = dir.as_ref().to_path_buf();
fs::create_dir_all(&dir)?;
let path = dir.join("sessions.json");
let sessions = Self::load_sessions(&path)?;
Ok(Self {
sessions: RwLock::new(sessions),
persistence_path: Some(path),
})
}
fn load_sessions(path: &Path) -> std::io::Result<HashMap<String, SharedSession>> {
if !path.exists() {
return Ok(HashMap::new());
}
let bytes = fs::read(path)?;
let persisted: HashMap<String, PersistedSession> = match serde_json::from_slice(&bytes) {
Ok(v) => v,
Err(e) => {
eprintln!("warning: failed to deserialize sessions from {}: {e}; starting with empty state", path.display());
HashMap::new()
}
};
Ok(persisted
.into_iter()
.map(|(id, mut record)| {
if record.session_id != id {
tracing::warn!(
map_key = %id,
record_session_id = %record.session_id,
path = %path.display(),
"persisted session key disagrees with its session_id; \
repairing to the map key"
);
record.session_id.clone_from(&id);
}
let session: Session = record.into();
(id, Arc::new(tokio::sync::Mutex::new(session)))
})
.collect())
}
fn persist_map(
path: &Path,
sessions: &HashMap<String, PersistedSession>,
) -> std::io::Result<()> {
let bytes = serde_json::to_vec_pretty(sessions)?;
let tmp_path = path.with_extension("json.tmp");
fs::write(&tmp_path, bytes)?;
fs::rename(&tmp_path, path)
}
pub async fn persist_snapshot(&self) -> std::io::Result<()> {
let Some(path) = self.persistence_path.clone() else {
return Ok(());
};
let arcs: Vec<(String, SharedSession)> = {
let guard = self.sessions.read().await;
guard
.iter()
.map(|(id, arc)| (id.clone(), Arc::clone(arc)))
.collect()
};
let mut persisted = HashMap::with_capacity(arcs.len());
for (id, arc) in arcs {
let session = arc.lock().await;
persisted.insert(id, PersistedSession::from(&*session));
}
Self::persist_map(&path, &persisted)
}
pub async fn get_shared(&self, session_id: &str) -> Option<SharedSession> {
let guard = self.sessions.read().await;
guard.get(session_id).cloned()
}
pub async fn get_session(&self, session_id: &str) -> Option<Session> {
let arc = self.get_shared(session_id).await?;
let session = arc.lock().await;
Some(session.clone())
}
pub async fn get_all_sessions(&self) -> Vec<Session> {
let arcs: Vec<SharedSession> = {
let guard = self.sessions.read().await;
guard.values().cloned().collect()
};
let mut out = Vec::with_capacity(arcs.len());
for arc in arcs {
out.push(arc.lock().await.clone());
}
out
}
pub async fn session_ids_after(&self, after: Option<&str>, limit: usize) -> Vec<String> {
if limit == 0 {
return Vec::new();
}
{
let guard = self.sessions.read().await;
let capacity = limit.saturating_add(1).min(guard.len().saturating_add(1));
let mut heap: BinaryHeap<&String> = BinaryHeap::with_capacity(capacity);
for key in guard.keys() {
if after.is_none_or(|a| key.as_str() > a) {
heap.push(key);
if heap.len() > limit {
heap.pop();
}
}
}
heap.into_sorted_vec().into_iter().cloned().collect()
}
}
pub async fn insert_recovered_session(&self, session_id: String, session: Session) {
debug_assert_eq!(
session.session_id, session_id,
"registry map key must equal Session::session_id — ListSessions paging \
orders by the key but emits the field (plan D1)"
);
{
let mut guard = self.sessions.write().await;
guard.insert(session_id, Arc::new(tokio::sync::Mutex::new(session)));
}
let _ = self.persist_snapshot().await;
}
pub async fn count_open_sessions_for_initiator(&self, sender: &str) -> usize {
let now = chrono::Utc::now().timestamp_millis();
let arcs: Vec<SharedSession> = {
let guard = self.sessions.read().await;
guard.values().cloned().collect()
};
let mut count = 0;
for arc in arcs {
let counts = match arc.try_lock() {
Ok(session) => {
session.initiator_sender == sender
&& session.state == macp_core::session::SessionState::Open
&& now <= session.ttl_expiry
}
Err(_) => true,
};
if counts {
count += 1;
}
}
count
}
}
#[cfg(test)]
mod tests {
use super::*;
use macp_core::session::{Session, SessionState};
use std::collections::HashSet;
use std::time::{SystemTime, UNIX_EPOCH};
fn sample_session(id: &str) -> Session {
Session::builder(id, "macp.mode.decision.v1", "alice")
.ttl_expiry(10)
.ttl_ms(9)
.started_at_unix_ms(1)
.mode_state(vec![1, 2, 3])
.participants(vec!["alice".into()])
.seen_message_ids(HashSet::from(["m1".into()]))
.intent("intent")
.mode_version("1.0.0")
.configuration_version("cfg")
.policy_version("pol")
.context_id("test-ctx")
.roots(vec![macp_pb::pb::Root {
uri: "root://1".into(),
name: "r1".into(),
}])
.build()
}
async fn registry_with(ids: &[String]) -> SessionRegistry {
let registry = SessionRegistry::new();
for id in ids {
registry
.insert_recovered_session(id.clone(), sample_session(id))
.await;
}
registry
}
fn sort_then_truncate_reference(
ids: &[String],
after: Option<&str>,
limit: usize,
) -> Vec<String> {
let mut sorted: Vec<String> = ids.to_vec();
sorted.sort();
sorted
.into_iter()
.filter(|id| after.is_none_or(|a| id.as_str() > a))
.take(limit)
.collect()
}
fn deterministic_ids(count: usize) -> Vec<String> {
let mut state: u64 = 0x2545_F491_4F6C_DD1D;
let mut ids = Vec::with_capacity(count);
for i in 0..count {
state = state
.wrapping_mul(6_364_136_223_846_793_005)
.wrapping_add(1_442_695_040_888_963_407);
ids.push(format!("sess-{:016x}-{i:04}", state >> 16));
}
ids
}
#[tokio::test]
async fn session_ids_after_returns_ascending_ids() {
let ids: Vec<String> = ["delta", "alpha", "charlie", "bravo"]
.iter()
.map(|s| s.to_string())
.collect();
let registry = registry_with(&ids).await;
let page = registry.session_ids_after(None, 10).await;
assert_eq!(page, vec!["alpha", "bravo", "charlie", "delta"]);
let page = registry.session_ids_after(None, 2).await;
assert_eq!(page, vec!["alpha", "bravo"]);
}
#[tokio::test]
async fn session_ids_after_respects_limit() {
let ids: Vec<String> = (0..10).map(|i| format!("s{i:02}")).collect();
let registry = registry_with(&ids).await;
assert_eq!(registry.session_ids_after(None, 1).await, vec!["s00"]);
assert_eq!(
registry.session_ids_after(None, 3).await,
vec!["s00", "s01", "s02"]
);
assert_eq!(registry.session_ids_after(None, 100).await.len(), 10);
}
#[tokio::test]
async fn session_ids_after_is_exclusive_of_cursor() {
let ids: Vec<String> = ["a", "b", "c", "d"].iter().map(|s| s.to_string()).collect();
let registry = registry_with(&ids).await;
let page = registry.session_ids_after(Some("b"), 10).await;
assert_eq!(page, vec!["c", "d"]);
assert!(!page.contains(&"b".to_string()));
assert!(page.iter().all(|id| id.as_str() > "b"));
assert!(registry.session_ids_after(Some("d"), 10).await.is_empty());
assert!(registry.session_ids_after(Some("zzz"), 10).await.is_empty());
}
#[tokio::test]
async fn session_ids_after_tolerates_absent_cursor() {
let ids: Vec<String> = ["a", "c", "e"].iter().map(|s| s.to_string()).collect();
let registry = registry_with(&ids).await;
assert_eq!(
registry.session_ids_after(Some("b"), 10).await,
vec!["c", "e"]
);
assert_eq!(
registry.session_ids_after(Some("b"), 10).await,
registry.session_ids_after(Some("a"), 10).await
);
assert_eq!(
registry.session_ids_after(Some(""), 10).await,
vec!["a", "c", "e"]
);
}
#[tokio::test]
async fn session_ids_after_zero_limit_is_empty() {
let ids: Vec<String> = ["a", "b", "c"].iter().map(|s| s.to_string()).collect();
let registry = registry_with(&ids).await;
assert!(registry.session_ids_after(None, 0).await.is_empty());
assert!(registry.session_ids_after(Some("a"), 0).await.is_empty());
let empty = SessionRegistry::new();
assert!(empty.session_ids_after(None, 0).await.is_empty());
assert!(empty.session_ids_after(None, 10).await.is_empty());
assert!(empty.session_ids_after(Some("a"), 10).await.is_empty());
}
#[tokio::test]
async fn session_ids_after_handles_huge_limits() {
let ids: Vec<String> = ["a", "b", "c"].iter().map(|s| s.to_string()).collect();
let registry = registry_with(&ids).await;
for limit in [usize::MAX, usize::MAX - 1, 10_000_000, 1 << 40] {
assert_eq!(
registry.session_ids_after(None, limit).await,
vec!["a", "b", "c"],
"limit={limit}"
);
assert_eq!(
registry.session_ids_after(Some("a"), limit).await,
vec!["b", "c"],
"limit={limit}"
);
}
let empty = SessionRegistry::new();
assert!(empty.session_ids_after(None, usize::MAX).await.is_empty());
}
#[tokio::test]
async fn session_ids_after_matches_sort_then_truncate_reference() {
let ids = deterministic_ids(200);
let registry = registry_with(&ids).await;
let mut sorted = ids.clone();
sorted.sort();
let cursors: Vec<Option<String>> = std::iter::once(None)
.chain(std::iter::once(Some(String::new())))
.chain(std::iter::once(Some("sess-".to_string())))
.chain(std::iter::once(Some("zzzz".to_string())))
.chain(sorted.iter().step_by(17).cloned().map(Some))
.chain(std::iter::once(Some(sorted.last().unwrap().clone())))
.chain(sorted.iter().step_by(23).map(|k| Some(format!("{k}~"))))
.collect();
for cursor in &cursors {
for limit in [1usize, 2, 7, 50, 199, 200, 201, 1000] {
let got = registry.session_ids_after(cursor.as_deref(), limit).await;
let want = sort_then_truncate_reference(&ids, cursor.as_deref(), limit);
assert_eq!(got, want, "cursor={cursor:?} limit={limit}");
}
}
}
#[tokio::test]
async fn session_ids_after_full_traversal_covers_every_id_once() {
let ids = deterministic_ids(200);
let registry = registry_with(&ids).await;
for page_size in [1usize, 3, 7, 64, 199, 200, 500] {
let mut collected: Vec<String> = Vec::new();
let mut cursor: Option<String> = None;
loop {
let page = registry
.session_ids_after(cursor.as_deref(), page_size)
.await;
let short = page.len() < page_size;
assert!(
page.len() <= page_size,
"page_size={page_size}: page of {} exceeds the limit",
page.len()
);
if let (Some(last), Some(first)) = (collected.last(), page.first()) {
assert!(first > last, "page_size={page_size}: page did not advance");
}
collected.extend(page.iter().cloned());
cursor = page.last().cloned();
if short {
break;
}
}
let unique: HashSet<&String> = collected.iter().collect();
assert_eq!(
collected.len(),
unique.len(),
"page_size={page_size}: duplicate IDs across pages"
);
let expected: HashSet<&String> = ids.iter().collect();
assert_eq!(unique, expected, "page_size={page_size}: coverage mismatch");
assert_eq!(collected.len(), ids.len(), "page_size={page_size}");
}
}
#[tokio::test]
async fn expired_sessions_not_counted_against_limit() {
let registry = SessionRegistry::new();
let now = chrono::Utc::now().timestamp_millis();
let mut expired = sample_session("expired-s1");
expired.initiator_sender = "agent://alice".into();
expired.ttl_expiry = now - 1000; expired.state = SessionState::Open; registry
.insert_recovered_session("expired-s1".into(), expired)
.await;
let count = registry
.count_open_sessions_for_initiator("agent://alice")
.await;
assert_eq!(count, 0);
let mut active = sample_session("active-s1");
active.initiator_sender = "agent://alice".into();
active.ttl_expiry = now + 60_000; active.state = SessionState::Open;
registry
.insert_recovered_session("active-s1".into(), active)
.await;
let count = registry
.count_open_sessions_for_initiator("agent://alice")
.await;
assert_eq!(count, 1);
}
#[tokio::test]
async fn load_sessions_repairs_key_field_mismatch() {
let base = std::env::temp_dir().join(format!(
"macp-registry-mismatch-{}",
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_nanos()
));
fs::create_dir_all(&base).unwrap();
let mut persisted = HashMap::new();
persisted.insert(
"A".to_string(),
PersistedSession::from(&sample_session("B")),
);
SessionRegistry::persist_map(&base.join("sessions.json"), &persisted).unwrap();
let reopened = SessionRegistry::with_persistence(&base).unwrap();
let session = reopened.get_session("A").await.unwrap();
assert_eq!(session.session_id, "A");
assert!(reopened.get_session("B").await.is_none());
assert_eq!(reopened.session_ids_after(None, 10).await, vec!["A"]);
}
#[tokio::test]
async fn persistent_registry_round_trip() {
let base = std::env::temp_dir().join(format!(
"macp-registry-test-{}",
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_nanos()
));
let registry = SessionRegistry::with_persistence(&base).unwrap();
registry
.insert_recovered_session("s1".into(), sample_session("s1"))
.await;
let reopened = SessionRegistry::with_persistence(&base).unwrap();
let session = reopened.get_session("s1").await.unwrap();
assert_eq!(session.mode, "macp.mode.decision.v1");
assert_eq!(session.mode_version, "1.0.0");
assert!(session.seen_message_ids.contains("m1"));
}
}