1use std::collections::{HashMap, VecDeque};
2use std::fmt;
3use std::sync::Mutex as StdMutex;
4use std::time::{Duration, Instant};
5
6use tokio::sync::broadcast;
7
8use crate::runtime::{LashRuntime, RuntimeSessionState};
9
10const SESSION_CURSOR_PREFIX: &str = "lashsc1:";
11const DEFAULT_LIVE_REPLAY_CAPACITY: usize = 2048;
12const DEFAULT_LIVE_REPLAY_TTL: Duration = Duration::from_secs(120);
13
14#[derive(
15 Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, serde::Serialize, serde::Deserialize,
16)]
17#[serde(transparent)]
18pub struct SessionRevision(pub u64);
19
20impl SessionRevision {
21 pub fn new(revision: u64) -> Self {
22 Self(revision)
23 }
24
25 pub fn as_u64(self) -> u64 {
26 self.0
27 }
28
29 pub(super) fn from_runtime(runtime: &LashRuntime) -> Self {
30 Self::from_state(&runtime.export_persisted_state())
31 }
32
33 pub(super) fn from_state(state: &RuntimeSessionState) -> Self {
34 Self(state.head_revision.unwrap_or(state.turn_index as u64))
35 }
36}
37
38#[derive(Clone, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)]
39#[serde(transparent)]
40pub struct SessionCursor(String);
41
42impl SessionCursor {
43 pub(crate) fn new(
44 session_id: impl AsRef<str>,
45 revision: SessionRevision,
46 live_position: u64,
47 ) -> Self {
48 Self(format!(
49 "{SESSION_CURSOR_PREFIX}{}:{live_position}:{}",
50 revision.0,
51 session_id.as_ref()
52 ))
53 }
54
55 #[cfg(test)]
56 pub(super) fn from_raw_for_testing(raw: impl Into<String>) -> Self {
57 Self(raw.into())
58 }
59
60 pub fn as_str(&self) -> &str {
61 &self.0
62 }
63
64 pub(crate) fn parse_for_session(
65 &self,
66 expected_session_id: &str,
67 ) -> Result<ParsedSessionCursor, SessionCursorError> {
68 let parsed = self.parse()?;
69 if parsed.session_id != expected_session_id {
70 return Err(SessionCursorError::WrongSession {
71 expected_session_id: expected_session_id.to_string(),
72 actual_session_id: parsed.session_id,
73 });
74 }
75 Ok(parsed)
76 }
77
78 fn parse(&self) -> Result<ParsedSessionCursor, SessionCursorError> {
79 let payload = self.0.strip_prefix(SESSION_CURSOR_PREFIX).ok_or_else(|| {
80 SessionCursorError::Malformed {
81 message: "missing cursor prefix".to_string(),
82 }
83 })?;
84 let mut parts = payload.splitn(3, ':');
85 let revision = parts
86 .next()
87 .ok_or_else(|| SessionCursorError::Malformed {
88 message: "missing session revision".to_string(),
89 })?
90 .parse::<u64>()
91 .map_err(|err| SessionCursorError::Malformed {
92 message: format!("invalid session revision: {err}"),
93 })?;
94 let live_position = parts
95 .next()
96 .ok_or_else(|| SessionCursorError::Malformed {
97 message: "missing live replay position".to_string(),
98 })?
99 .parse::<u64>()
100 .map_err(|err| SessionCursorError::Malformed {
101 message: format!("invalid live replay position: {err}"),
102 })?;
103 let session_id = parts
104 .next()
105 .filter(|value| !value.is_empty())
106 .ok_or_else(|| SessionCursorError::Malformed {
107 message: "missing session id".to_string(),
108 })?
109 .to_string();
110 Ok(ParsedSessionCursor {
111 session_id,
112 revision: SessionRevision(revision),
113 live_position,
114 })
115 }
116}
117
118impl fmt::Debug for SessionCursor {
119 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
120 f.write_str("SessionCursor(<opaque>)")
121 }
122}
123
124impl fmt::Display for SessionCursor {
125 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
126 f.write_str(&self.0)
127 }
128}
129
130#[derive(Clone, Debug)]
131pub(crate) struct ParsedSessionCursor {
132 pub session_id: String,
133 pub revision: SessionRevision,
134 pub live_position: u64,
135}
136
137#[derive(Clone, Debug, thiserror::Error)]
138pub enum SessionCursorError {
139 #[error("malformed session cursor: {message}")]
140 Malformed { message: String },
141 #[error("session cursor belongs to `{actual_session_id}`, not `{expected_session_id}`")]
142 WrongSession {
143 expected_session_id: String,
144 actual_session_id: String,
145 },
146}
147
148#[derive(Clone, Debug)]
149pub struct SessionObservation {
150 pub read_view: crate::SessionReadView,
151 pub cursor: SessionCursor,
152}
153
154#[derive(Clone, Debug)]
155pub struct SessionObservationEvent {
156 pub session_id: String,
157 pub revision: SessionRevision,
158 pub cursor: SessionCursor,
159 pub payload: SessionObservationEventPayload,
160}
161
162#[derive(Clone, Copy, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
163#[serde(rename_all = "snake_case")]
164pub enum SessionQueueEventKind {
165 Enqueued,
166 Cancelled,
167}
168
169#[derive(Clone, Copy, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
170#[serde(rename_all = "snake_case")]
171pub enum SessionProcessEventKind {
172 Started,
173 Cancelled,
174}
175
176#[derive(Clone, Debug)]
177#[allow(clippy::large_enum_variant)]
178pub enum SessionObservationEventPayload {
179 TurnActivity(crate::TurnActivity),
180 Committed {
181 read_view: crate::SessionReadView,
182 },
183 AgentFrameSwitched {
184 frame_id: String,
185 },
186 QueueChanged {
187 kind: SessionQueueEventKind,
188 batch_ids: Vec<String>,
189 },
190 ProcessChanged {
191 kind: SessionProcessEventKind,
192 process_ids: Vec<String>,
193 },
194}
195
196#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
197pub struct LiveReplayGap {
198 pub session_id: String,
199 pub requested_cursor: SessionCursor,
200 pub latest_cursor: SessionCursor,
201 pub latest_revision: SessionRevision,
202 pub reason: LiveReplayGapReason,
203}
204
205#[derive(Clone, Copy, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
206#[serde(rename_all = "snake_case")]
207pub enum LiveReplayGapReason {
208 Trimmed,
209 Unavailable,
210}
211
212#[derive(Clone, Debug, thiserror::Error)]
213pub enum LiveReplayStoreError {
214 #[error("{0}")]
215 Cursor(#[from] SessionCursorError),
216 #[error("live replay store error: {0}")]
217 Store(String),
218 #[error("live replay subscriber lagged by {0} events")]
219 SubscriberLagged(u64),
220 #[error("live replay channel closed")]
221 Closed,
222}
223
224#[derive(Clone, Debug)]
225pub enum LiveReplayResult {
226 Replayed(Vec<SessionObservationEvent>),
227 Gap(LiveReplayGapReason),
228}
229
230pub enum LiveReplaySubscribeResult {
231 Subscribed(LiveReplaySubscription),
232 Gap(LiveReplayGapReason),
233}
234
235pub struct LiveReplaySubscription {
236 replay: VecDeque<SessionObservationEvent>,
237 receiver: broadcast::Receiver<SessionObservationEvent>,
238}
239
240impl LiveReplaySubscription {
241 fn new(
242 replay: Vec<SessionObservationEvent>,
243 receiver: broadcast::Receiver<SessionObservationEvent>,
244 ) -> Self {
245 Self {
246 replay: replay.into(),
247 receiver,
248 }
249 }
250
251 pub async fn next_event(&mut self) -> Result<SessionObservationEvent, LiveReplayStoreError> {
252 if let Some(event) = self.replay.pop_front() {
253 return Ok(event);
254 }
255 match self.receiver.recv().await {
256 Ok(event) => Ok(event),
257 Err(broadcast::error::RecvError::Lagged(count)) => {
258 Err(LiveReplayStoreError::SubscriberLagged(count))
259 }
260 Err(broadcast::error::RecvError::Closed) => Err(LiveReplayStoreError::Closed),
261 }
262 }
263}
264
265#[derive(Clone, Debug)]
266pub enum SessionResume {
267 Replayed {
268 events: Vec<SessionObservationEvent>,
269 },
270 Gap {
271 observation: SessionObservation,
272 gap: LiveReplayGap,
273 },
274}
275
276pub enum SessionObservationSubscription {
277 Subscribed(LiveReplaySubscription),
278 Gap {
279 observation: SessionObservation,
280 gap: LiveReplayGap,
281 },
282}
283
284pub trait LiveReplayStore: Send + Sync {
292 fn append(
296 &self,
297 session_id: &str,
298 revision: SessionRevision,
299 payload: SessionObservationEventPayload,
300 ) -> Result<SessionObservationEvent, LiveReplayStoreError>;
301
302 fn replay_after_cursor(
306 &self,
307 cursor: &SessionCursor,
308 ) -> Result<LiveReplayResult, LiveReplayStoreError>;
309
310 fn subscribe_after_cursor(
314 &self,
315 cursor: &SessionCursor,
316 ) -> Result<LiveReplaySubscribeResult, LiveReplayStoreError>;
317
318 fn current_cursor(&self, session_id: &str, revision: SessionRevision) -> SessionCursor;
322
323 fn trim_session(&self, session_id: &str) -> Result<(), LiveReplayStoreError>;
327}
328
329#[derive(Clone, Debug)]
330pub struct InMemoryLiveReplayStoreConfig {
331 pub max_events_per_session: usize,
332 pub max_age: Duration,
333}
334
335impl Default for InMemoryLiveReplayStoreConfig {
336 fn default() -> Self {
337 Self {
338 max_events_per_session: DEFAULT_LIVE_REPLAY_CAPACITY,
339 max_age: DEFAULT_LIVE_REPLAY_TTL,
340 }
341 }
342}
343
344#[derive(Debug)]
345pub struct InMemoryLiveReplayStore {
346 config: InMemoryLiveReplayStoreConfig,
347 sessions: StdMutex<HashMap<String, LiveReplaySessionBuffer>>,
348}
349
350impl InMemoryLiveReplayStore {
351 pub fn new(config: InMemoryLiveReplayStoreConfig) -> Self {
352 Self {
353 config,
354 sessions: StdMutex::new(HashMap::new()),
355 }
356 }
357
358 pub fn with_bounds(max_events_per_session: usize, max_age: Duration) -> Self {
359 Self::new(InMemoryLiveReplayStoreConfig {
360 max_events_per_session,
361 max_age,
362 })
363 }
364}
365
366impl Default for InMemoryLiveReplayStore {
367 fn default() -> Self {
368 Self::new(InMemoryLiveReplayStoreConfig::default())
369 }
370}
371
372#[derive(Debug)]
373struct LiveReplaySessionBuffer {
374 events: VecDeque<StoredObservationEvent>,
375 tail_position: u64,
376 sender: broadcast::Sender<SessionObservationEvent>,
377}
378
379impl LiveReplaySessionBuffer {
380 fn new() -> Self {
381 let (sender, _) = broadcast::channel(DEFAULT_LIVE_REPLAY_CAPACITY);
382 Self {
383 events: VecDeque::new(),
384 tail_position: 0,
385 sender,
386 }
387 }
388}
389
390#[derive(Clone, Debug)]
391struct StoredObservationEvent {
392 position: u64,
393 appended_at: Instant,
394 event: SessionObservationEvent,
395}
396
397impl InMemoryLiveReplayStore {
398 fn trim_locked(config: &InMemoryLiveReplayStoreConfig, buffer: &mut LiveReplaySessionBuffer) {
399 let now = Instant::now();
400 while buffer.events.len() > config.max_events_per_session {
401 buffer.events.pop_front();
402 }
403 while buffer
404 .events
405 .front()
406 .is_some_and(|event| now.duration_since(event.appended_at) > config.max_age)
407 {
408 buffer.events.pop_front();
409 }
410 }
411
412 fn gap_reason_for_cursor(
413 buffer: Option<&LiveReplaySessionBuffer>,
414 cursor_position: u64,
415 ) -> Option<LiveReplayGapReason> {
416 let Some(buffer) = buffer else {
417 return (cursor_position > 0).then_some(LiveReplayGapReason::Unavailable);
418 };
419 if cursor_position > buffer.tail_position {
420 return Some(LiveReplayGapReason::Unavailable);
421 }
422 let Some(first) = buffer.events.front() else {
423 return (cursor_position < buffer.tail_position)
424 .then_some(LiveReplayGapReason::Trimmed);
425 };
426 if cursor_position + 1 < first.position {
427 Some(LiveReplayGapReason::Trimmed)
428 } else {
429 None
430 }
431 }
432}
433
434impl LiveReplayStore for InMemoryLiveReplayStore {
435 fn append(
436 &self,
437 session_id: &str,
438 revision: SessionRevision,
439 payload: SessionObservationEventPayload,
440 ) -> Result<SessionObservationEvent, LiveReplayStoreError> {
441 let mut sessions = self
442 .sessions
443 .lock()
444 .map_err(|_| LiveReplayStoreError::Store("live replay mutex poisoned".to_string()))?;
445 let buffer = sessions
446 .entry(session_id.to_string())
447 .or_insert_with(LiveReplaySessionBuffer::new);
448 buffer.tail_position = buffer.tail_position.saturating_add(1);
449 let cursor = SessionCursor::new(session_id, revision, buffer.tail_position);
450 let event = SessionObservationEvent {
451 session_id: session_id.to_string(),
452 revision,
453 cursor,
454 payload,
455 };
456 buffer.events.push_back(StoredObservationEvent {
457 position: buffer.tail_position,
458 appended_at: Instant::now(),
459 event: event.clone(),
460 });
461 Self::trim_locked(&self.config, buffer);
462 let _ = buffer.sender.send(event.clone());
463 Ok(event)
464 }
465
466 fn replay_after_cursor(
467 &self,
468 cursor: &SessionCursor,
469 ) -> Result<LiveReplayResult, LiveReplayStoreError> {
470 let parsed = cursor.parse()?;
471 let _cursor_revision = parsed.revision;
472 let mut sessions = self
473 .sessions
474 .lock()
475 .map_err(|_| LiveReplayStoreError::Store("live replay mutex poisoned".to_string()))?;
476 if let Some(buffer) = sessions.get_mut(&parsed.session_id) {
477 Self::trim_locked(&self.config, buffer);
478 }
479 let buffer = sessions.get(&parsed.session_id);
480 if let Some(reason) = Self::gap_reason_for_cursor(buffer, parsed.live_position) {
481 return Ok(LiveReplayResult::Gap(reason));
482 }
483 let events = buffer
484 .map(|buffer| {
485 buffer
486 .events
487 .iter()
488 .filter(|event| event.position > parsed.live_position)
489 .map(|event| event.event.clone())
490 .collect()
491 })
492 .unwrap_or_default();
493 Ok(LiveReplayResult::Replayed(events))
494 }
495
496 fn subscribe_after_cursor(
497 &self,
498 cursor: &SessionCursor,
499 ) -> Result<LiveReplaySubscribeResult, LiveReplayStoreError> {
500 let parsed = cursor.parse()?;
501 let _cursor_revision = parsed.revision;
502 let mut sessions = self
503 .sessions
504 .lock()
505 .map_err(|_| LiveReplayStoreError::Store("live replay mutex poisoned".to_string()))?;
506 let buffer = sessions
507 .entry(parsed.session_id.clone())
508 .or_insert_with(LiveReplaySessionBuffer::new);
509 Self::trim_locked(&self.config, buffer);
510 let receiver = buffer.sender.subscribe();
511 if let Some(reason) = Self::gap_reason_for_cursor(Some(buffer), parsed.live_position) {
512 return Ok(LiveReplaySubscribeResult::Gap(reason));
513 }
514 let replay = buffer
515 .events
516 .iter()
517 .filter(|event| event.position > parsed.live_position)
518 .map(|event| event.event.clone())
519 .collect();
520 Ok(LiveReplaySubscribeResult::Subscribed(
521 LiveReplaySubscription::new(replay, receiver),
522 ))
523 }
524
525 fn current_cursor(&self, session_id: &str, revision: SessionRevision) -> SessionCursor {
526 let tail_position = self
527 .sessions
528 .lock()
529 .ok()
530 .and_then(|sessions| sessions.get(session_id).map(|buffer| buffer.tail_position))
531 .unwrap_or(0);
532 SessionCursor::new(session_id, revision, tail_position)
533 }
534
535 fn trim_session(&self, session_id: &str) -> Result<(), LiveReplayStoreError> {
536 let mut sessions = self
537 .sessions
538 .lock()
539 .map_err(|_| LiveReplayStoreError::Store("live replay mutex poisoned".to_string()))?;
540 if let Some(buffer) = sessions.get_mut(session_id) {
541 Self::trim_locked(&self.config, buffer);
542 }
543 Ok(())
544 }
545}
546
547#[cfg(test)]
548mod tests {
549 use super::*;
550
551 fn activity(text: &str) -> SessionObservationEventPayload {
552 SessionObservationEventPayload::TurnActivity(crate::TurnActivity::independent(
553 crate::TurnEvent::AssistantProseDelta {
554 text: text.to_string(),
555 },
556 ))
557 }
558
559 #[test]
560 fn session_cursor_round_trips_and_debug_is_opaque() {
561 let cursor = SessionCursor::new("session:with:colon", SessionRevision(3), 9);
562 let encoded = serde_json::to_string(&cursor).expect("serialize");
563 let decoded: SessionCursor = serde_json::from_str(&encoded).expect("deserialize");
564 assert_eq!(decoded, cursor);
565 assert_eq!(format!("{cursor:?}"), "SessionCursor(<opaque>)");
566 let parsed = cursor
567 .parse_for_session("session:with:colon")
568 .expect("parse");
569 assert_eq!(parsed.revision, SessionRevision(3));
570 assert_eq!(parsed.live_position, 9);
571 }
572
573 #[test]
574 fn session_cursor_rejects_malformed_and_wrong_session() {
575 let malformed = SessionCursor::from_raw_for_testing("bad");
576 assert!(matches!(
577 malformed.parse_for_session("s"),
578 Err(SessionCursorError::Malformed { .. })
579 ));
580 let cursor = SessionCursor::new("actual", SessionRevision(0), 0);
581 assert!(matches!(
582 cursor.parse_for_session("expected"),
583 Err(SessionCursorError::WrongSession { .. })
584 ));
585 }
586
587 #[test]
588 fn in_memory_replay_store_replays_after_cursor_in_order() {
589 let store = InMemoryLiveReplayStore::default();
590 let start = store.current_cursor("s", SessionRevision(0));
591 store
592 .append("s", SessionRevision(0), activity("a"))
593 .expect("append a");
594 store
595 .append("s", SessionRevision(0), activity("b"))
596 .expect("append b");
597 let LiveReplayResult::Replayed(events) = store.replay_after_cursor(&start).expect("replay")
598 else {
599 panic!("expected replay");
600 };
601 assert_eq!(events.len(), 2);
602 match &events[0].payload {
603 SessionObservationEventPayload::TurnActivity(activity) => match &activity.event {
604 crate::TurnEvent::AssistantProseDelta { text } => assert_eq!(text, "a"),
605 _ => panic!("wrong event"),
606 },
607 _ => panic!("wrong payload"),
608 }
609 }
610
611 #[test]
612 fn in_memory_replay_store_reports_gap_after_capacity_trim() {
613 let store = InMemoryLiveReplayStore::with_bounds(1, Duration::from_secs(120));
614 let start = store.current_cursor("s", SessionRevision(0));
615 store
616 .append("s", SessionRevision(0), activity("a"))
617 .expect("append a");
618 store
619 .append("s", SessionRevision(0), activity("b"))
620 .expect("append b");
621 assert!(matches!(
622 store.replay_after_cursor(&start).expect("gap"),
623 LiveReplayResult::Gap(LiveReplayGapReason::Trimmed)
624 ));
625 }
626
627 #[test]
628 fn in_memory_replay_store_reports_gap_after_ttl_trim() {
629 let store = InMemoryLiveReplayStore::with_bounds(16, Duration::from_millis(1));
630 let start = store.current_cursor("s", SessionRevision(0));
631 store
632 .append("s", SessionRevision(0), activity("a"))
633 .expect("append a");
634 std::thread::sleep(Duration::from_millis(5));
635 assert!(matches!(
636 store.replay_after_cursor(&start).expect("gap"),
637 LiveReplayResult::Gap(LiveReplayGapReason::Trimmed)
638 ));
639 }
640
641 #[test]
642 fn in_memory_replay_store_reports_unavailable_for_cursor_ahead_of_tail() {
643 let store = InMemoryLiveReplayStore::default();
644 let ahead = SessionCursor::new("s", SessionRevision(0), 99);
645 assert!(matches!(
646 store.replay_after_cursor(&ahead).expect("gap"),
647 LiveReplayResult::Gap(LiveReplayGapReason::Unavailable)
648 ));
649 }
650
651 #[tokio::test]
652 async fn in_memory_replay_subscription_yields_replay_then_live() {
653 let store = InMemoryLiveReplayStore::default();
654 let start = store.current_cursor("s", SessionRevision(0));
655 store
656 .append("s", SessionRevision(0), activity("a"))
657 .expect("append a");
658 let LiveReplaySubscribeResult::Subscribed(mut subscription) =
659 store.subscribe_after_cursor(&start).expect("subscribe")
660 else {
661 panic!("expected subscription");
662 };
663 let first = subscription.next_event().await.expect("replay");
664 assert_eq!(first.session_id, "s");
665 store
666 .append("s", SessionRevision(0), activity("b"))
667 .expect("append b");
668 let second = subscription.next_event().await.expect("live");
669 match second.payload {
670 SessionObservationEventPayload::TurnActivity(activity) => match activity.event {
671 crate::TurnEvent::AssistantProseDelta { text } => assert_eq!(text, "b"),
672 _ => panic!("wrong event"),
673 },
674 _ => panic!("wrong payload"),
675 }
676 }
677
678 #[test]
679 fn in_memory_replay_subscription_reports_gap_after_capacity_trim() {
680 let store = InMemoryLiveReplayStore::with_bounds(1, Duration::from_secs(120));
681 let start = store.current_cursor("s", SessionRevision(0));
682 store
683 .append("s", SessionRevision(0), activity("a"))
684 .expect("append a");
685 store
686 .append("s", SessionRevision(0), activity("b"))
687 .expect("append b");
688 assert!(matches!(
689 store.subscribe_after_cursor(&start).expect("subscribe"),
690 LiveReplaySubscribeResult::Gap(LiveReplayGapReason::Trimmed)
691 ));
692 }
693
694 #[test]
695 fn in_memory_replay_subscription_reports_gap_after_ttl_trim() {
696 let store = InMemoryLiveReplayStore::with_bounds(16, Duration::from_millis(1));
697 let start = store.current_cursor("s", SessionRevision(0));
698 store
699 .append("s", SessionRevision(0), activity("a"))
700 .expect("append a");
701 std::thread::sleep(Duration::from_millis(5));
702 assert!(matches!(
703 store.subscribe_after_cursor(&start).expect("subscribe"),
704 LiveReplaySubscribeResult::Gap(LiveReplayGapReason::Trimmed)
705 ));
706 }
707}