1use macp_core::session::Session;
2use std::collections::{BinaryHeap, HashMap};
3use std::fs;
4use std::path::{Path, PathBuf};
5use std::sync::Arc;
6use tokio::sync::RwLock;
7
8#[derive(serde::Serialize, serde::Deserialize)]
9pub struct PersistedRoot {
10 pub uri: String,
11 pub name: String,
12}
13
14#[derive(serde::Serialize, serde::Deserialize)]
15pub struct PersistedSession {
16 #[serde(default = "default_schema_version")]
17 pub schema_version: u32,
18 pub session_id: String,
19 pub state: macp_core::session::SessionState,
20 pub ttl_expiry: i64,
21 #[serde(default)]
22 pub ttl_ms: i64,
23 pub started_at_unix_ms: i64,
24 pub resolution: Option<Vec<u8>>,
25 pub mode: String,
26 pub mode_state: Vec<u8>,
27 pub participants: Vec<String>,
28 pub seen_message_ids: Vec<String>,
29 pub intent: String,
30 pub mode_version: String,
31 pub configuration_version: String,
32 pub policy_version: String,
33 #[serde(default)]
34 pub context_id: String,
35 #[serde(default)]
36 pub extensions: HashMap<String, Vec<u8>>,
37 pub roots: Vec<PersistedRoot>,
38 pub initiator_sender: String,
39 #[serde(default)]
40 pub policy_definition: Option<macp_core::policy::PolicyDefinition>,
41 #[serde(default)]
42 pub suspended_at_ms: Option<i64>,
43 #[serde(default)]
44 pub accumulated_suspended_ms: i64,
45 #[serde(default)]
48 pub semantics_rev: u32,
49 #[serde(default)]
52 pub max_suspend_ms: i64,
53}
54
55fn default_schema_version() -> u32 {
56 2
57}
58
59impl From<&Session> for PersistedSession {
60 fn from(session: &Session) -> Self {
61 Self {
62 schema_version: 2,
63 session_id: session.session_id.clone(),
64 state: session.state.clone(),
65 ttl_expiry: session.ttl_expiry,
66 ttl_ms: session.ttl_ms,
67 started_at_unix_ms: session.started_at_unix_ms,
68 resolution: session.resolution.clone(),
69 mode: session.mode.clone(),
70 mode_state: session.mode_state.clone(),
71 participants: session.participants.clone(),
72 seen_message_ids: session.seen_message_ids.iter().cloned().collect(),
73 intent: session.intent.clone(),
74 mode_version: session.mode_version.clone(),
75 configuration_version: session.configuration_version.clone(),
76 policy_version: session.policy_version.clone(),
77 context_id: session.context_id.clone(),
78 extensions: session.extensions.clone(),
79 roots: session
80 .roots
81 .iter()
82 .map(|root| PersistedRoot {
83 uri: root.uri.clone(),
84 name: root.name.clone(),
85 })
86 .collect(),
87 initiator_sender: session.initiator_sender.clone(),
88 policy_definition: session.policy_definition.clone(),
89 suspended_at_ms: session.suspended_at_ms,
90 accumulated_suspended_ms: session.accumulated_suspended_ms,
91 semantics_rev: session.semantics_rev,
92 max_suspend_ms: session.max_suspend_ms,
93 }
94 }
95}
96
97impl From<PersistedSession> for Session {
98 fn from(session: PersistedSession) -> Self {
99 let ttl_ms = if session.ttl_ms > 0 {
100 session.ttl_ms
101 } else {
102 session
104 .ttl_expiry
105 .saturating_sub(session.started_at_unix_ms)
106 };
107 Session::builder(session.session_id, session.mode, session.initiator_sender)
108 .state(session.state)
109 .ttl_expiry(session.ttl_expiry)
110 .ttl_ms(ttl_ms)
111 .started_at_unix_ms(session.started_at_unix_ms)
112 .resolution(session.resolution)
113 .mode_state(session.mode_state)
114 .participants(session.participants)
115 .seen_message_ids(session.seen_message_ids.into_iter().collect())
116 .intent(session.intent)
117 .mode_version(session.mode_version)
118 .configuration_version(session.configuration_version)
119 .policy_version(session.policy_version)
120 .context_id(session.context_id)
121 .extensions(session.extensions)
122 .roots(
123 session
124 .roots
125 .into_iter()
126 .map(|root| macp_pb::pb::Root {
127 uri: root.uri,
128 name: root.name,
129 })
130 .collect(),
131 )
132 .policy_definition(session.policy_definition)
133 .suspended_at_ms(session.suspended_at_ms)
134 .accumulated_suspended_ms(session.accumulated_suspended_ms)
135 .semantics_rev(session.semantics_rev)
136 .max_suspend_ms(session.max_suspend_ms)
137 .build()
138 }
139}
140
141pub type SharedSession = Arc<tokio::sync::Mutex<Session>>;
149
150pub struct SessionRegistry {
151 pub sessions: RwLock<HashMap<String, SharedSession>>,
152 persistence_path: Option<PathBuf>,
153}
154
155impl Default for SessionRegistry {
156 fn default() -> Self {
157 Self::new()
158 }
159}
160
161impl SessionRegistry {
162 pub fn new() -> Self {
163 Self {
164 sessions: RwLock::new(HashMap::new()),
165 persistence_path: None,
166 }
167 }
168
169 pub fn with_persistence<P: AsRef<Path>>(dir: P) -> std::io::Result<Self> {
170 let dir = dir.as_ref().to_path_buf();
171 fs::create_dir_all(&dir)?;
172 let path = dir.join("sessions.json");
173 let sessions = Self::load_sessions(&path)?;
174 Ok(Self {
175 sessions: RwLock::new(sessions),
176 persistence_path: Some(path),
177 })
178 }
179
180 fn load_sessions(path: &Path) -> std::io::Result<HashMap<String, SharedSession>> {
181 if !path.exists() {
182 return Ok(HashMap::new());
183 }
184 let bytes = fs::read(path)?;
185 let persisted: HashMap<String, PersistedSession> = match serde_json::from_slice(&bytes) {
186 Ok(v) => v,
187 Err(e) => {
188 eprintln!("warning: failed to deserialize sessions from {}: {e}; starting with empty state", path.display());
189 HashMap::new()
190 }
191 };
192 Ok(persisted
193 .into_iter()
194 .map(|(id, mut record)| {
195 if record.session_id != id {
205 tracing::warn!(
206 map_key = %id,
207 record_session_id = %record.session_id,
208 path = %path.display(),
209 "persisted session key disagrees with its session_id; \
210 repairing to the map key"
211 );
212 record.session_id.clone_from(&id);
213 }
214 let session: Session = record.into();
215 (id, Arc::new(tokio::sync::Mutex::new(session)))
216 })
217 .collect())
218 }
219
220 fn persist_map(
221 path: &Path,
222 sessions: &HashMap<String, PersistedSession>,
223 ) -> std::io::Result<()> {
224 let bytes = serde_json::to_vec_pretty(sessions)?;
225 let tmp_path = path.with_extension("json.tmp");
226 fs::write(&tmp_path, bytes)?;
227 fs::rename(&tmp_path, path)
228 }
229
230 pub async fn persist_snapshot(&self) -> std::io::Result<()> {
233 let Some(path) = self.persistence_path.clone() else {
234 return Ok(());
235 };
236 let arcs: Vec<(String, SharedSession)> = {
237 let guard = self.sessions.read().await;
238 guard
239 .iter()
240 .map(|(id, arc)| (id.clone(), Arc::clone(arc)))
241 .collect()
242 };
243 let mut persisted = HashMap::with_capacity(arcs.len());
244 for (id, arc) in arcs {
245 let session = arc.lock().await;
246 persisted.insert(id, PersistedSession::from(&*session));
247 }
248 Self::persist_map(&path, &persisted)
249 }
250
251 pub async fn get_shared(&self, session_id: &str) -> Option<SharedSession> {
253 let guard = self.sessions.read().await;
254 guard.get(session_id).cloned()
255 }
256
257 pub async fn get_session(&self, session_id: &str) -> Option<Session> {
258 let arc = self.get_shared(session_id).await?;
259 let session = arc.lock().await;
260 Some(session.clone())
261 }
262
263 pub async fn get_all_sessions(&self) -> Vec<Session> {
264 let arcs: Vec<SharedSession> = {
265 let guard = self.sessions.read().await;
266 guard.values().cloned().collect()
267 };
268 let mut out = Vec::with_capacity(arcs.len());
269 for arc in arcs {
270 out.push(arc.lock().await.clone());
271 }
272 out
273 }
274
275 pub async fn session_ids_after(&self, after: Option<&str>, limit: usize) -> Vec<String> {
295 if limit == 0 {
296 return Vec::new();
297 }
298 {
302 let guard = self.sessions.read().await;
303 let capacity = limit.saturating_add(1).min(guard.len().saturating_add(1));
310 let mut heap: BinaryHeap<&String> = BinaryHeap::with_capacity(capacity);
311 for key in guard.keys() {
312 if after.is_none_or(|a| key.as_str() > a) {
313 heap.push(key);
314 if heap.len() > limit {
315 heap.pop();
316 }
317 }
318 }
319 heap.into_sorted_vec().into_iter().cloned().collect()
320 }
321 }
322
323 pub async fn insert_recovered_session(&self, session_id: String, session: Session) {
324 debug_assert_eq!(
336 session.session_id, session_id,
337 "registry map key must equal Session::session_id — ListSessions paging \
338 orders by the key but emits the field (plan D1)"
339 );
340 {
341 let mut guard = self.sessions.write().await;
342 guard.insert(session_id, Arc::new(tokio::sync::Mutex::new(session)));
343 }
344 let _ = self.persist_snapshot().await;
345 }
346
347 pub async fn count_open_sessions_for_initiator(&self, sender: &str) -> usize {
348 let now = chrono::Utc::now().timestamp_millis();
349 let arcs: Vec<SharedSession> = {
350 let guard = self.sessions.read().await;
351 guard.values().cloned().collect()
352 };
353 let mut count = 0;
354 for arc in arcs {
355 let counts = match arc.try_lock() {
358 Ok(session) => {
359 session.initiator_sender == sender
360 && session.state == macp_core::session::SessionState::Open
361 && now <= session.ttl_expiry
362 }
363 Err(_) => true,
364 };
365 if counts {
366 count += 1;
367 }
368 }
369 count
370 }
371}
372
373#[cfg(test)]
374mod tests {
375 use super::*;
376 use macp_core::session::{Session, SessionState};
377 use std::collections::HashSet;
378 use std::time::{SystemTime, UNIX_EPOCH};
379
380 fn sample_session(id: &str) -> Session {
381 Session::builder(id, "macp.mode.decision.v1", "alice")
382 .ttl_expiry(10)
383 .ttl_ms(9)
384 .started_at_unix_ms(1)
385 .mode_state(vec![1, 2, 3])
386 .participants(vec!["alice".into()])
387 .seen_message_ids(HashSet::from(["m1".into()]))
388 .intent("intent")
389 .mode_version("1.0.0")
390 .configuration_version("cfg")
391 .policy_version("pol")
392 .context_id("test-ctx")
393 .roots(vec![macp_pb::pb::Root {
394 uri: "root://1".into(),
395 name: "r1".into(),
396 }])
397 .build()
398 }
399
400 async fn registry_with(ids: &[String]) -> SessionRegistry {
403 let registry = SessionRegistry::new();
404 for id in ids {
405 registry
406 .insert_recovered_session(id.clone(), sample_session(id))
407 .await;
408 }
409 registry
410 }
411
412 fn sort_then_truncate_reference(
415 ids: &[String],
416 after: Option<&str>,
417 limit: usize,
418 ) -> Vec<String> {
419 let mut sorted: Vec<String> = ids.to_vec();
420 sorted.sort();
421 sorted
422 .into_iter()
423 .filter(|id| after.is_none_or(|a| id.as_str() > a))
424 .take(limit)
425 .collect()
426 }
427
428 fn deterministic_ids(count: usize) -> Vec<String> {
432 let mut state: u64 = 0x2545_F491_4F6C_DD1D;
433 let mut ids = Vec::with_capacity(count);
434 for i in 0..count {
435 state = state
436 .wrapping_mul(6_364_136_223_846_793_005)
437 .wrapping_add(1_442_695_040_888_963_407);
438 ids.push(format!("sess-{:016x}-{i:04}", state >> 16));
440 }
441 ids
442 }
443
444 #[tokio::test]
445 async fn session_ids_after_returns_ascending_ids() {
446 let ids: Vec<String> = ["delta", "alpha", "charlie", "bravo"]
447 .iter()
448 .map(|s| s.to_string())
449 .collect();
450 let registry = registry_with(&ids).await;
451
452 let page = registry.session_ids_after(None, 10).await;
453 assert_eq!(page, vec!["alpha", "bravo", "charlie", "delta"]);
454
455 let page = registry.session_ids_after(None, 2).await;
457 assert_eq!(page, vec!["alpha", "bravo"]);
458 }
459
460 #[tokio::test]
461 async fn session_ids_after_respects_limit() {
462 let ids: Vec<String> = (0..10).map(|i| format!("s{i:02}")).collect();
463 let registry = registry_with(&ids).await;
464
465 assert_eq!(registry.session_ids_after(None, 1).await, vec!["s00"]);
466 assert_eq!(
469 registry.session_ids_after(None, 3).await,
470 vec!["s00", "s01", "s02"]
471 );
472 assert_eq!(registry.session_ids_after(None, 100).await.len(), 10);
474 }
475
476 #[tokio::test]
477 async fn session_ids_after_is_exclusive_of_cursor() {
478 let ids: Vec<String> = ["a", "b", "c", "d"].iter().map(|s| s.to_string()).collect();
479 let registry = registry_with(&ids).await;
480
481 let page = registry.session_ids_after(Some("b"), 10).await;
482 assert_eq!(page, vec!["c", "d"]);
483 assert!(!page.contains(&"b".to_string()));
484 assert!(page.iter().all(|id| id.as_str() > "b"));
485
486 assert!(registry.session_ids_after(Some("d"), 10).await.is_empty());
488 assert!(registry.session_ids_after(Some("zzz"), 10).await.is_empty());
490 }
491
492 #[tokio::test]
493 async fn session_ids_after_tolerates_absent_cursor() {
494 let ids: Vec<String> = ["a", "c", "e"].iter().map(|s| s.to_string()).collect();
495 let registry = registry_with(&ids).await;
496
497 assert_eq!(
500 registry.session_ids_after(Some("b"), 10).await,
501 vec!["c", "e"]
502 );
503 assert_eq!(
505 registry.session_ids_after(Some("b"), 10).await,
506 registry.session_ids_after(Some("a"), 10).await
507 );
508 assert_eq!(
513 registry.session_ids_after(Some(""), 10).await,
514 vec!["a", "c", "e"]
515 );
516 }
517
518 #[tokio::test]
519 async fn session_ids_after_zero_limit_is_empty() {
520 let ids: Vec<String> = ["a", "b", "c"].iter().map(|s| s.to_string()).collect();
521 let registry = registry_with(&ids).await;
522
523 assert!(registry.session_ids_after(None, 0).await.is_empty());
524 assert!(registry.session_ids_after(Some("a"), 0).await.is_empty());
525
526 let empty = SessionRegistry::new();
528 assert!(empty.session_ids_after(None, 0).await.is_empty());
529 assert!(empty.session_ids_after(None, 10).await.is_empty());
530 assert!(empty.session_ids_after(Some("a"), 10).await.is_empty());
531 }
532
533 #[tokio::test]
538 async fn session_ids_after_handles_huge_limits() {
539 let ids: Vec<String> = ["a", "b", "c"].iter().map(|s| s.to_string()).collect();
540 let registry = registry_with(&ids).await;
541
542 for limit in [usize::MAX, usize::MAX - 1, 10_000_000, 1 << 40] {
543 assert_eq!(
544 registry.session_ids_after(None, limit).await,
545 vec!["a", "b", "c"],
546 "limit={limit}"
547 );
548 assert_eq!(
549 registry.session_ids_after(Some("a"), limit).await,
550 vec!["b", "c"],
551 "limit={limit}"
552 );
553 }
554
555 let empty = SessionRegistry::new();
557 assert!(empty.session_ids_after(None, usize::MAX).await.is_empty());
558 }
559
560 #[tokio::test]
561 async fn session_ids_after_matches_sort_then_truncate_reference() {
562 let ids = deterministic_ids(200);
563 let registry = registry_with(&ids).await;
564
565 let mut sorted = ids.clone();
566 sorted.sort();
567
568 let cursors: Vec<Option<String>> = std::iter::once(None)
569 .chain(std::iter::once(Some(String::new())))
570 .chain(std::iter::once(Some("sess-".to_string())))
571 .chain(std::iter::once(Some("zzzz".to_string())))
572 .chain(sorted.iter().step_by(17).cloned().map(Some))
574 .chain(std::iter::once(Some(sorted.last().unwrap().clone())))
575 .chain(sorted.iter().step_by(23).map(|k| Some(format!("{k}~"))))
577 .collect();
578
579 for cursor in &cursors {
580 for limit in [1usize, 2, 7, 50, 199, 200, 201, 1000] {
581 let got = registry.session_ids_after(cursor.as_deref(), limit).await;
582 let want = sort_then_truncate_reference(&ids, cursor.as_deref(), limit);
583 assert_eq!(got, want, "cursor={cursor:?} limit={limit}");
584 }
585 }
586 }
587
588 #[tokio::test]
589 async fn session_ids_after_full_traversal_covers_every_id_once() {
590 let ids = deterministic_ids(200);
591 let registry = registry_with(&ids).await;
592
593 for page_size in [1usize, 3, 7, 64, 199, 200, 500] {
594 let mut collected: Vec<String> = Vec::new();
595 let mut cursor: Option<String> = None;
596 loop {
597 let page = registry
598 .session_ids_after(cursor.as_deref(), page_size)
599 .await;
600 let short = page.len() < page_size;
601 assert!(
605 page.len() <= page_size,
606 "page_size={page_size}: page of {} exceeds the limit",
607 page.len()
608 );
609 if let (Some(last), Some(first)) = (collected.last(), page.first()) {
611 assert!(first > last, "page_size={page_size}: page did not advance");
612 }
613 collected.extend(page.iter().cloned());
614 cursor = page.last().cloned();
615 if short {
616 break;
617 }
618 }
619
620 let unique: HashSet<&String> = collected.iter().collect();
621 assert_eq!(
623 collected.len(),
624 unique.len(),
625 "page_size={page_size}: duplicate IDs across pages"
626 );
627 let expected: HashSet<&String> = ids.iter().collect();
628 assert_eq!(unique, expected, "page_size={page_size}: coverage mismatch");
629 assert_eq!(collected.len(), ids.len(), "page_size={page_size}");
630 }
631 }
632
633 #[tokio::test]
634 async fn expired_sessions_not_counted_against_limit() {
635 let registry = SessionRegistry::new();
636 let now = chrono::Utc::now().timestamp_millis();
637 let mut expired = sample_session("expired-s1");
639 expired.initiator_sender = "agent://alice".into();
640 expired.ttl_expiry = now - 1000; expired.state = SessionState::Open; registry
643 .insert_recovered_session("expired-s1".into(), expired)
644 .await;
645
646 let count = registry
648 .count_open_sessions_for_initiator("agent://alice")
649 .await;
650 assert_eq!(count, 0);
651
652 let mut active = sample_session("active-s1");
654 active.initiator_sender = "agent://alice".into();
655 active.ttl_expiry = now + 60_000; active.state = SessionState::Open;
657 registry
658 .insert_recovered_session("active-s1".into(), active)
659 .await;
660
661 let count = registry
662 .count_open_sessions_for_initiator("agent://alice")
663 .await;
664 assert_eq!(count, 1);
665 }
666
667 #[tokio::test]
673 async fn load_sessions_repairs_key_field_mismatch() {
674 let base = std::env::temp_dir().join(format!(
675 "macp-registry-mismatch-{}",
676 SystemTime::now()
677 .duration_since(UNIX_EPOCH)
678 .unwrap()
679 .as_nanos()
680 ));
681 fs::create_dir_all(&base).unwrap();
682
683 let mut persisted = HashMap::new();
684 persisted.insert(
685 "A".to_string(),
686 PersistedSession::from(&sample_session("B")),
687 );
688 SessionRegistry::persist_map(&base.join("sessions.json"), &persisted).unwrap();
689
690 let reopened = SessionRegistry::with_persistence(&base).unwrap();
691
692 let session = reopened.get_session("A").await.unwrap();
694 assert_eq!(session.session_id, "A");
695 assert!(reopened.get_session("B").await.is_none());
697 assert_eq!(reopened.session_ids_after(None, 10).await, vec!["A"]);
700 }
701
702 #[tokio::test]
703 async fn persistent_registry_round_trip() {
704 let base = std::env::temp_dir().join(format!(
705 "macp-registry-test-{}",
706 SystemTime::now()
707 .duration_since(UNIX_EPOCH)
708 .unwrap()
709 .as_nanos()
710 ));
711
712 let registry = SessionRegistry::with_persistence(&base).unwrap();
713 registry
714 .insert_recovered_session("s1".into(), sample_session("s1"))
715 .await;
716
717 let reopened = SessionRegistry::with_persistence(&base).unwrap();
718 let session = reopened.get_session("s1").await.unwrap();
719 assert_eq!(session.mode, "macp.mode.decision.v1");
720 assert_eq!(session.mode_version, "1.0.0");
721 assert!(session.seen_message_ids.contains("m1"));
722 }
723}