lash_core/runtime/observation/
replay.rs

1use std::collections::{HashMap, VecDeque};
2use std::fmt;
3use std::pin::Pin;
4use std::sync::{Arc, Mutex as StdMutex};
5use std::task::{Context, Poll, ready};
6use std::time::{Duration, Instant};
7
8use futures_util::Stream;
9use tokio::sync::broadcast;
10use tokio_util::sync::ReusableBoxFuture;
11
12use crate::runtime::LashRuntime;
13#[cfg(test)]
14use crate::runtime::RuntimeSessionState;
15
16const SESSION_CURSOR_PREFIX: &str = "lashsc1:";
17const DEFAULT_LIVE_REPLAY_CAPACITY: usize = 2048;
18const DEFAULT_LIVE_REPLAY_TTL: Duration = Duration::from_secs(120);
19
20#[derive(
21    Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, serde::Serialize, serde::Deserialize,
22)]
23#[serde(transparent)]
24pub struct SessionRevision(pub u64);
25
26impl SessionRevision {
27    pub fn new(revision: u64) -> Self {
28        Self(revision)
29    }
30
31    pub fn as_u64(self) -> u64 {
32        self.0
33    }
34
35    pub(super) fn from_runtime(runtime: &LashRuntime) -> Self {
36        Self(
37            runtime
38                .state
39                .head_revision
40                .unwrap_or(runtime.state.turn_index as u64),
41        )
42    }
43
44    #[cfg(test)]
45    pub(super) fn from_state(state: &RuntimeSessionState) -> Self {
46        Self(state.head_revision.unwrap_or(state.turn_index as u64))
47    }
48}
49
50#[derive(Clone, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)]
51#[serde(transparent)]
52pub struct SessionCursor(String);
53
54impl SessionCursor {
55    pub(crate) fn new(
56        session_id: impl AsRef<str>,
57        revision: SessionRevision,
58        live_position: u64,
59    ) -> Self {
60        Self(format!(
61            "{SESSION_CURSOR_PREFIX}{}:{live_position}:{}",
62            revision.0,
63            session_id.as_ref()
64        ))
65    }
66
67    #[cfg(test)]
68    pub(super) fn from_raw_for_testing(raw: impl Into<String>) -> Self {
69        Self(raw.into())
70    }
71
72    pub fn as_str(&self) -> &str {
73        &self.0
74    }
75
76    pub(crate) fn parse_for_session(
77        &self,
78        expected_session_id: &str,
79    ) -> Result<ParsedSessionCursor, SessionCursorError> {
80        let parsed = self.parse()?;
81        if parsed.session_id != expected_session_id {
82            return Err(SessionCursorError::WrongSession {
83                expected_session_id: expected_session_id.to_string(),
84                actual_session_id: parsed.session_id,
85            });
86        }
87        Ok(parsed)
88    }
89
90    fn parse(&self) -> Result<ParsedSessionCursor, SessionCursorError> {
91        let payload = self.0.strip_prefix(SESSION_CURSOR_PREFIX).ok_or_else(|| {
92            SessionCursorError::Malformed {
93                message: "missing cursor prefix".to_string(),
94            }
95        })?;
96        let mut parts = payload.splitn(3, ':');
97        let revision = parts
98            .next()
99            .ok_or_else(|| SessionCursorError::Malformed {
100                message: "missing session revision".to_string(),
101            })?
102            .parse::<u64>()
103            .map_err(|err| SessionCursorError::Malformed {
104                message: format!("invalid session revision: {err}"),
105            })?;
106        let live_position = parts
107            .next()
108            .ok_or_else(|| SessionCursorError::Malformed {
109                message: "missing live replay position".to_string(),
110            })?
111            .parse::<u64>()
112            .map_err(|err| SessionCursorError::Malformed {
113                message: format!("invalid live replay position: {err}"),
114            })?;
115        let session_id = parts
116            .next()
117            .filter(|value| !value.is_empty())
118            .ok_or_else(|| SessionCursorError::Malformed {
119                message: "missing session id".to_string(),
120            })?
121            .to_string();
122        Ok(ParsedSessionCursor {
123            session_id,
124            revision: SessionRevision(revision),
125            live_position,
126        })
127    }
128}
129
130impl fmt::Debug for SessionCursor {
131    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
132        f.write_str("SessionCursor(<opaque>)")
133    }
134}
135
136impl fmt::Display for SessionCursor {
137    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
138        f.write_str(&self.0)
139    }
140}
141
142#[derive(Clone, Debug)]
143pub(crate) struct ParsedSessionCursor {
144    pub session_id: String,
145    pub revision: SessionRevision,
146    pub live_position: u64,
147}
148
149#[derive(Clone, Debug, thiserror::Error)]
150pub enum SessionCursorError {
151    #[error("malformed session cursor: {message}")]
152    Malformed { message: String },
153    #[error("session cursor belongs to `{actual_session_id}`, not `{expected_session_id}`")]
154    WrongSession {
155        expected_session_id: String,
156        actual_session_id: String,
157    },
158}
159
160#[derive(Clone, Debug)]
161pub struct SessionObservation {
162    pub read_view: crate::SessionReadView,
163    pub cursor: SessionCursor,
164}
165
166#[derive(Clone, Debug)]
167pub struct SessionObservationEvent {
168    pub session_id: String,
169    pub revision: SessionRevision,
170    pub cursor: SessionCursor,
171    pub payload: SessionObservationEventPayload,
172}
173
174#[derive(Clone, Copy, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
175#[serde(rename_all = "snake_case")]
176pub enum SessionQueueEventKind {
177    Enqueued,
178    Cancelled,
179}
180
181#[derive(Clone, Copy, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
182#[serde(rename_all = "snake_case")]
183pub enum SessionProcessEventKind {
184    Started,
185    Cancelled,
186}
187
188#[derive(Clone, Debug)]
189// justification: the enclosing replay event is already Arc-owned, so another allocation would not bound retained event storage.
190#[allow(clippy::large_enum_variant)]
191pub enum SessionObservationEventPayload {
192    TurnActivity(crate::TurnActivity),
193    Committed {
194        read_view: crate::SessionReadView,
195    },
196    AgentFrameSwitched {
197        frame_id: String,
198    },
199    QueueChanged {
200        kind: SessionQueueEventKind,
201        batch_ids: Vec<String>,
202    },
203    ProcessChanged {
204        kind: SessionProcessEventKind,
205        process_ids: Vec<String>,
206    },
207}
208
209#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
210pub struct LiveReplayGap {
211    pub session_id: String,
212    pub requested_cursor: SessionCursor,
213    pub latest_cursor: SessionCursor,
214    pub latest_revision: SessionRevision,
215    pub reason: LiveReplayGapReason,
216}
217
218#[derive(Clone, Copy, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
219#[serde(rename_all = "snake_case")]
220pub enum LiveReplayGapReason {
221    Trimmed,
222    Unavailable,
223}
224
225#[derive(Clone, Debug, thiserror::Error)]
226pub enum LiveReplayStoreError {
227    #[error("{0}")]
228    Cursor(#[from] SessionCursorError),
229    #[error("live replay store error: {0}")]
230    Store(String),
231    #[error("live replay subscriber lagged by {0} events")]
232    SubscriberLagged(u64),
233    #[error("live replay channel closed")]
234    Closed,
235}
236
237#[derive(Clone, Debug)]
238pub enum LiveReplayResult {
239    Replayed(Vec<Arc<SessionObservationEvent>>),
240    Gap(LiveReplayGapReason),
241}
242
243pub enum LiveReplaySubscribeResult {
244    Subscribed(LiveReplaySubscription),
245    Gap(LiveReplayGapReason),
246}
247
248type LiveReplayRecvResult = (
249    Result<Arc<SessionObservationEvent>, broadcast::error::RecvError>,
250    broadcast::Receiver<Arc<SessionObservationEvent>>,
251);
252
253#[cfg(test)]
254static LIVE_REPLAY_EVENT_CLONES: std::sync::atomic::AtomicUsize =
255    std::sync::atomic::AtomicUsize::new(0);
256
257#[inline]
258fn clone_event(event: &Arc<SessionObservationEvent>) -> Arc<SessionObservationEvent> {
259    #[cfg(test)]
260    LIVE_REPLAY_EVENT_CLONES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
261    Arc::clone(event)
262}
263
264pub struct LiveReplaySubscription {
265    replay: VecDeque<Arc<SessionObservationEvent>>,
266    receiver: ReusableBoxFuture<'static, LiveReplayRecvResult>,
267    closed: bool,
268}
269
270impl LiveReplaySubscription {
271    fn new(
272        replay: Vec<Arc<SessionObservationEvent>>,
273        receiver: broadcast::Receiver<Arc<SessionObservationEvent>>,
274    ) -> Self {
275        Self {
276            replay: replay.into(),
277            receiver: ReusableBoxFuture::new(live_replay_recv(receiver)),
278            closed: false,
279        }
280    }
281
282    pub async fn next_event(
283        &mut self,
284    ) -> Result<Arc<SessionObservationEvent>, LiveReplayStoreError> {
285        futures_util::StreamExt::next(self)
286            .await
287            .unwrap_or(Err(LiveReplayStoreError::Closed))
288    }
289}
290
291async fn live_replay_recv(
292    mut receiver: broadcast::Receiver<Arc<SessionObservationEvent>>,
293) -> LiveReplayRecvResult {
294    let result = receiver.recv().await;
295    #[cfg(test)]
296    if result.is_ok() {
297        LIVE_REPLAY_EVENT_CLONES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
298    }
299    (result, receiver)
300}
301
302impl Stream for LiveReplaySubscription {
303    type Item = Result<Arc<SessionObservationEvent>, LiveReplayStoreError>;
304
305    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
306        if let Some(event) = self.replay.pop_front() {
307            return Poll::Ready(Some(Ok(event)));
308        }
309        if self.closed {
310            return Poll::Ready(None);
311        }
312        let (result, receiver) = ready!(self.receiver.poll(cx));
313        self.receiver.set(live_replay_recv(receiver));
314        match result {
315            Ok(event) => Poll::Ready(Some(Ok(event))),
316            Err(broadcast::error::RecvError::Lagged(count)) => {
317                Poll::Ready(Some(Err(LiveReplayStoreError::SubscriberLagged(count))))
318            }
319            Err(broadcast::error::RecvError::Closed) => {
320                self.closed = true;
321                Poll::Ready(Some(Err(LiveReplayStoreError::Closed)))
322            }
323        }
324    }
325}
326
327#[derive(Clone, Debug)]
328pub enum SessionResume {
329    Replayed {
330        events: Vec<Arc<SessionObservationEvent>>,
331    },
332    Gap {
333        observation: SessionObservation,
334        gap: LiveReplayGap,
335    },
336}
337
338pub enum SessionObservationSubscription {
339    Subscribed(LiveReplaySubscription),
340    Gap {
341        observation: SessionObservation,
342        gap: LiveReplayGap,
343    },
344}
345
346/// Bounded, best-effort live replay for host reconnects.
347///
348/// Runtime turn execution calls this trait from synchronous boundary code. All
349/// methods must therefore be fast and nonblocking from the runtime's point of
350/// view. A custom external store should expose local or buffered behavior here,
351/// or offload blocking transport and durability work internally. Runtime turn
352/// execution must not wait for slow network or storage durability in this path.
353pub trait LiveReplayStore: Send + Sync {
354    /// Append one observation event and return its assigned cursor.
355    ///
356    /// This must be fast and nonblocking from the runtime's point of view.
357    fn append(
358        &self,
359        session_id: &str,
360        revision: SessionRevision,
361        payload: SessionObservationEventPayload,
362    ) -> Result<Arc<SessionObservationEvent>, LiveReplayStoreError>;
363
364    /// Return buffered events after `cursor`, or report a recoverable gap.
365    ///
366    /// This must be fast and nonblocking from the runtime's point of view.
367    fn replay_after_cursor(
368        &self,
369        cursor: &SessionCursor,
370    ) -> Result<LiveReplayResult, LiveReplayStoreError>;
371
372    /// Subscribe after `cursor`, replaying buffered events before live events.
373    ///
374    /// This must be fast and nonblocking from the runtime's point of view.
375    fn subscribe_after_cursor(
376        &self,
377        cursor: &SessionCursor,
378    ) -> Result<LiveReplaySubscribeResult, LiveReplayStoreError>;
379
380    /// Return the latest cursor known locally for a session.
381    ///
382    /// This must be fast and nonblocking from the runtime's point of view.
383    fn current_cursor(&self, session_id: &str, revision: SessionRevision) -> SessionCursor;
384
385    /// Apply best-effort retention trimming for a session.
386    ///
387    /// This must be fast and nonblocking from the runtime's point of view.
388    fn trim_session(&self, session_id: &str) -> Result<(), LiveReplayStoreError>;
389}
390
391#[derive(Clone, Debug)]
392pub struct InMemoryLiveReplayStoreConfig {
393    pub max_events_per_session: usize,
394    pub max_age: Duration,
395}
396
397impl Default for InMemoryLiveReplayStoreConfig {
398    fn default() -> Self {
399        Self {
400            max_events_per_session: DEFAULT_LIVE_REPLAY_CAPACITY,
401            max_age: DEFAULT_LIVE_REPLAY_TTL,
402        }
403    }
404}
405
406#[derive(Debug)]
407pub struct InMemoryLiveReplayStore {
408    config: InMemoryLiveReplayStoreConfig,
409    clock: Arc<dyn crate::Clock>,
410    sessions: StdMutex<HashMap<String, LiveReplaySessionBuffer>>,
411}
412
413impl InMemoryLiveReplayStore {
414    pub fn new(config: InMemoryLiveReplayStoreConfig) -> Self {
415        Self::with_clock(config, Arc::new(crate::SystemClock))
416    }
417
418    pub fn with_clock(config: InMemoryLiveReplayStoreConfig, clock: Arc<dyn crate::Clock>) -> Self {
419        Self {
420            config,
421            clock,
422            sessions: StdMutex::new(HashMap::new()),
423        }
424    }
425
426    pub fn with_bounds(max_events_per_session: usize, max_age: Duration) -> Self {
427        Self::new(InMemoryLiveReplayStoreConfig {
428            max_events_per_session,
429            max_age,
430        })
431    }
432}
433
434impl Default for InMemoryLiveReplayStore {
435    fn default() -> Self {
436        Self::new(InMemoryLiveReplayStoreConfig::default())
437    }
438}
439
440#[derive(Debug)]
441struct LiveReplaySessionBuffer {
442    events: VecDeque<StoredObservationEvent>,
443    tail_position: u64,
444    sender: Option<broadcast::Sender<Arc<SessionObservationEvent>>>,
445}
446
447impl LiveReplaySessionBuffer {
448    fn new() -> Self {
449        Self {
450            events: VecDeque::new(),
451            tail_position: 0,
452            sender: None,
453        }
454    }
455
456    fn subscribe(
457        &mut self,
458        channel_capacity: usize,
459    ) -> broadcast::Receiver<Arc<SessionObservationEvent>> {
460        match self.sender.as_ref() {
461            Some(sender) => sender.subscribe(),
462            None => {
463                let (sender, receiver) = broadcast::channel(channel_capacity.max(1));
464                self.sender = Some(sender);
465                receiver
466            }
467        }
468    }
469
470    fn publish(&mut self, event: Arc<SessionObservationEvent>) {
471        let Some(sender) = self.sender.as_ref() else {
472            return;
473        };
474        if sender.send(event).is_err() {
475            self.sender = None;
476        }
477    }
478}
479
480#[derive(Clone, Debug)]
481struct StoredObservationEvent {
482    position: u64,
483    appended_at: Instant,
484    event: Arc<SessionObservationEvent>,
485}
486
487impl InMemoryLiveReplayStore {
488    fn trim_locked(
489        config: &InMemoryLiveReplayStoreConfig,
490        buffer: &mut LiveReplaySessionBuffer,
491        now: Instant,
492    ) {
493        while buffer.events.len() > config.max_events_per_session {
494            buffer.events.pop_front();
495        }
496        while buffer
497            .events
498            .front()
499            .is_some_and(|event| now.duration_since(event.appended_at) > config.max_age)
500        {
501            buffer.events.pop_front();
502        }
503    }
504
505    fn gap_reason_for_cursor(
506        buffer: Option<&LiveReplaySessionBuffer>,
507        cursor_position: u64,
508    ) -> Option<LiveReplayGapReason> {
509        let Some(buffer) = buffer else {
510            return (cursor_position > 0).then_some(LiveReplayGapReason::Unavailable);
511        };
512        if cursor_position > buffer.tail_position {
513            return Some(LiveReplayGapReason::Unavailable);
514        }
515        let Some(first) = buffer.events.front() else {
516            return (cursor_position < buffer.tail_position)
517                .then_some(LiveReplayGapReason::Trimmed);
518        };
519        if cursor_position + 1 < first.position {
520            Some(LiveReplayGapReason::Trimmed)
521        } else {
522            None
523        }
524    }
525}
526
527impl LiveReplayStore for InMemoryLiveReplayStore {
528    fn append(
529        &self,
530        session_id: &str,
531        revision: SessionRevision,
532        payload: SessionObservationEventPayload,
533    ) -> Result<Arc<SessionObservationEvent>, LiveReplayStoreError> {
534        let now = self.clock.now();
535        let mut sessions = self
536            .sessions
537            .lock()
538            .map_err(|_| LiveReplayStoreError::Store("live replay mutex poisoned".to_string()))?;
539        let buffer = sessions
540            .entry(session_id.to_string())
541            .or_insert_with(LiveReplaySessionBuffer::new);
542        buffer.tail_position = buffer.tail_position.saturating_add(1);
543        let cursor = SessionCursor::new(session_id, revision, buffer.tail_position);
544        let event = Arc::new(SessionObservationEvent {
545            session_id: session_id.to_string(),
546            revision,
547            cursor,
548            payload,
549        });
550        buffer.events.push_back(StoredObservationEvent {
551            position: buffer.tail_position,
552            appended_at: now,
553            event: clone_event(&event),
554        });
555        Self::trim_locked(&self.config, buffer, now);
556        buffer.publish(clone_event(&event));
557        Ok(event)
558    }
559
560    fn replay_after_cursor(
561        &self,
562        cursor: &SessionCursor,
563    ) -> Result<LiveReplayResult, LiveReplayStoreError> {
564        let parsed = cursor.parse()?;
565        let _cursor_revision = parsed.revision;
566        let now = self.clock.now();
567        let mut sessions = self
568            .sessions
569            .lock()
570            .map_err(|_| LiveReplayStoreError::Store("live replay mutex poisoned".to_string()))?;
571        if let Some(buffer) = sessions.get_mut(&parsed.session_id) {
572            Self::trim_locked(&self.config, buffer, now);
573        }
574        let buffer = sessions.get(&parsed.session_id);
575        if let Some(reason) = Self::gap_reason_for_cursor(buffer, parsed.live_position) {
576            return Ok(LiveReplayResult::Gap(reason));
577        }
578        let events = buffer
579            .map(|buffer| {
580                buffer
581                    .events
582                    .iter()
583                    .filter(|event| event.position > parsed.live_position)
584                    .map(|event| clone_event(&event.event))
585                    .collect()
586            })
587            .unwrap_or_default();
588        Ok(LiveReplayResult::Replayed(events))
589    }
590
591    fn subscribe_after_cursor(
592        &self,
593        cursor: &SessionCursor,
594    ) -> Result<LiveReplaySubscribeResult, LiveReplayStoreError> {
595        let parsed = cursor.parse()?;
596        let _cursor_revision = parsed.revision;
597        let now = self.clock.now();
598        let mut sessions = self
599            .sessions
600            .lock()
601            .map_err(|_| LiveReplayStoreError::Store("live replay mutex poisoned".to_string()))?;
602        let buffer = sessions
603            .entry(parsed.session_id.clone())
604            .or_insert_with(LiveReplaySessionBuffer::new);
605        Self::trim_locked(&self.config, buffer, now);
606        if let Some(reason) = Self::gap_reason_for_cursor(Some(buffer), parsed.live_position) {
607            return Ok(LiveReplaySubscribeResult::Gap(reason));
608        }
609        let replay = buffer
610            .events
611            .iter()
612            .filter(|event| event.position > parsed.live_position)
613            .map(|event| clone_event(&event.event))
614            .collect();
615        let receiver = buffer.subscribe(self.config.max_events_per_session);
616        Ok(LiveReplaySubscribeResult::Subscribed(
617            LiveReplaySubscription::new(replay, receiver),
618        ))
619    }
620
621    fn current_cursor(&self, session_id: &str, revision: SessionRevision) -> SessionCursor {
622        let tail_position = self
623            .sessions
624            .lock()
625            .ok()
626            .and_then(|sessions| sessions.get(session_id).map(|buffer| buffer.tail_position))
627            .unwrap_or(0);
628        SessionCursor::new(session_id, revision, tail_position)
629    }
630
631    fn trim_session(&self, session_id: &str) -> Result<(), LiveReplayStoreError> {
632        let now = self.clock.now();
633        let mut sessions = self
634            .sessions
635            .lock()
636            .map_err(|_| LiveReplayStoreError::Store("live replay mutex poisoned".to_string()))?;
637        if let Some(buffer) = sessions.get_mut(session_id) {
638            Self::trim_locked(&self.config, buffer, now);
639        }
640        Ok(())
641    }
642}
643
644#[cfg(test)]
645mod tests {
646    use super::*;
647    use std::alloc::{GlobalAlloc, Layout, System};
648    use std::sync::atomic::{AtomicUsize, Ordering};
649
650    struct CountingAllocator;
651
652    static ALLOCATION_COUNT: AtomicUsize = AtomicUsize::new(0);
653    static ALLOCATED_BYTES: AtomicUsize = AtomicUsize::new(0);
654
655    unsafe impl GlobalAlloc for CountingAllocator {
656        unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
657            ALLOCATION_COUNT.fetch_add(1, Ordering::Relaxed);
658            ALLOCATED_BYTES.fetch_add(layout.size(), Ordering::Relaxed);
659            // SAFETY: forwarding the allocator contract unchanged to System.
660            unsafe { System.alloc(layout) }
661        }
662
663        unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
664            // SAFETY: `ptr` and `layout` came from the forwarded System allocation.
665            unsafe { System.dealloc(ptr, layout) }
666        }
667    }
668
669    #[global_allocator]
670    static TEST_ALLOCATOR: CountingAllocator = CountingAllocator;
671
672    fn activity(text: &str) -> SessionObservationEventPayload {
673        SessionObservationEventPayload::TurnActivity(crate::TurnActivity::independent(
674            crate::TurnEvent::AssistantProseDelta { text: text.into() },
675        ))
676    }
677
678    #[test]
679    fn session_cursor_round_trips_and_debug_is_opaque() {
680        let cursor = SessionCursor::new("session:with:colon", SessionRevision(3), 9);
681        let encoded = serde_json::to_string(&cursor).expect("serialize");
682        let decoded: SessionCursor = serde_json::from_str(&encoded).expect("deserialize");
683        assert_eq!(decoded, cursor);
684        assert_eq!(format!("{cursor:?}"), "SessionCursor(<opaque>)");
685        let parsed = cursor
686            .parse_for_session("session:with:colon")
687            .expect("parse");
688        assert_eq!(parsed.revision, SessionRevision(3));
689        assert_eq!(parsed.live_position, 9);
690    }
691
692    #[test]
693    fn session_cursor_rejects_malformed_and_wrong_session() {
694        let malformed = SessionCursor::from_raw_for_testing("bad");
695        assert!(matches!(
696            malformed.parse_for_session("s"),
697            Err(SessionCursorError::Malformed { .. })
698        ));
699        let cursor = SessionCursor::new("actual", SessionRevision(0), 0);
700        assert!(matches!(
701            cursor.parse_for_session("expected"),
702            Err(SessionCursorError::WrongSession { .. })
703        ));
704    }
705
706    #[test]
707    fn in_memory_replay_store_replays_after_cursor_in_order() {
708        let store = InMemoryLiveReplayStore::default();
709        let start = store.current_cursor("s", SessionRevision(0));
710        store
711            .append("s", SessionRevision(0), activity("a"))
712            .expect("append a");
713        store
714            .append("s", SessionRevision(0), activity("b"))
715            .expect("append b");
716        let LiveReplayResult::Replayed(events) = store.replay_after_cursor(&start).expect("replay")
717        else {
718            panic!("expected replay");
719        };
720        assert_eq!(events.len(), 2);
721        match &events[0].payload {
722            SessionObservationEventPayload::TurnActivity(activity) => match &activity.event {
723                crate::TurnEvent::AssistantProseDelta { text } => assert_eq!(text.as_ref(), "a"),
724                _ => panic!("wrong event"),
725            },
726            _ => panic!("wrong payload"),
727        }
728    }
729
730    #[test]
731    fn in_memory_replay_store_reports_gap_after_capacity_trim() {
732        let store = InMemoryLiveReplayStore::with_bounds(1, Duration::from_secs(120));
733        let start = store.current_cursor("s", SessionRevision(0));
734        store
735            .append("s", SessionRevision(0), activity("a"))
736            .expect("append a");
737        store
738            .append("s", SessionRevision(0), activity("b"))
739            .expect("append b");
740        assert!(matches!(
741            store.replay_after_cursor(&start).expect("gap"),
742            LiveReplayResult::Gap(LiveReplayGapReason::Trimmed)
743        ));
744    }
745
746    #[test]
747    fn in_memory_replay_store_reports_gap_after_ttl_trim() {
748        let store = InMemoryLiveReplayStore::with_bounds(16, Duration::from_millis(1));
749        let start = store.current_cursor("s", SessionRevision(0));
750        store
751            .append("s", SessionRevision(0), activity("a"))
752            .expect("append a");
753        std::thread::sleep(Duration::from_millis(5));
754        assert!(matches!(
755            store.replay_after_cursor(&start).expect("gap"),
756            LiveReplayResult::Gap(LiveReplayGapReason::Trimmed)
757        ));
758    }
759
760    #[test]
761    fn in_memory_replay_store_reports_unavailable_for_cursor_ahead_of_tail() {
762        let store = InMemoryLiveReplayStore::default();
763        let ahead = SessionCursor::new("s", SessionRevision(0), 99);
764        assert!(matches!(
765            store.replay_after_cursor(&ahead).expect("gap"),
766            LiveReplayResult::Gap(LiveReplayGapReason::Unavailable)
767        ));
768    }
769
770    #[tokio::test]
771    async fn in_memory_replay_subscription_yields_replay_then_live() {
772        let store = InMemoryLiveReplayStore::default();
773        let start = store.current_cursor("s", SessionRevision(0));
774        store
775            .append("s", SessionRevision(0), activity("a"))
776            .expect("append a");
777        let LiveReplaySubscribeResult::Subscribed(mut subscription) =
778            store.subscribe_after_cursor(&start).expect("subscribe")
779        else {
780            panic!("expected subscription");
781        };
782        let first = subscription.next_event().await.expect("replay");
783        assert_eq!(first.session_id, "s");
784        store
785            .append("s", SessionRevision(0), activity("b"))
786            .expect("append b");
787        let second = subscription.next_event().await.expect("live");
788        match &second.payload {
789            SessionObservationEventPayload::TurnActivity(activity) => match &activity.event {
790                crate::TurnEvent::AssistantProseDelta { text } => assert_eq!(text.as_ref(), "b"),
791                _ => panic!("wrong event"),
792            },
793            _ => panic!("wrong payload"),
794        }
795    }
796
797    #[tokio::test]
798    #[ignore = "manual lane-O allocation measurement"]
799    async fn measure_streamed_token_allocations() {
800        const TOKENS: usize = 1_000;
801        let store = InMemoryLiveReplayStore::with_bounds(TOKENS + 1, Duration::from_secs(120));
802        let mut cursor = store.current_cursor("perf-session", SessionRevision(7));
803        let LiveReplaySubscribeResult::Subscribed(mut subscription) = store
804            .subscribe_after_cursor(&cursor)
805            .expect("subscribe for allocation measurement")
806        else {
807            panic!("expected subscription");
808        };
809
810        ALLOCATION_COUNT.store(0, Ordering::SeqCst);
811        ALLOCATED_BYTES.store(0, Ordering::SeqCst);
812        LIVE_REPLAY_EVENT_CLONES.store(0, Ordering::SeqCst);
813        for ordinal in 0..TOKENS {
814            let event = store
815                .append(
816                    "perf-session",
817                    SessionRevision(7),
818                    activity(&format!("token-{ordinal}")),
819                )
820                .expect("append token event");
821            let live = subscription.next_event().await.expect("receive live event");
822            assert_eq!(live.cursor, event.cursor);
823            let LiveReplayResult::Replayed(replayed) = store
824                .replay_after_cursor(&cursor)
825                .expect("replay token event")
826            else {
827                panic!("expected replay");
828            };
829            assert_eq!(replayed.len(), 1);
830            cursor = event.cursor.clone();
831        }
832        let allocations = ALLOCATION_COUNT.load(Ordering::SeqCst);
833        let bytes = ALLOCATED_BYTES.load(Ordering::SeqCst);
834        let event_clones = LIVE_REPLAY_EVENT_CLONES.load(Ordering::SeqCst);
835        eprintln!(
836            "streamed-token allocations: total={allocations} per_token={:.3} bytes_total={bytes} bytes_per_token={:.3} deep_event_clones_per_token=0 arc_handle_clones_per_token={:.3}",
837            allocations as f64 / TOKENS as f64,
838            bytes as f64 / TOKENS as f64,
839            event_clones as f64 / TOKENS as f64,
840        );
841    }
842
843    #[test]
844    fn in_memory_replay_store_allocates_live_channel_lazily() {
845        let store = InMemoryLiveReplayStore::default();
846        let start = store.current_cursor("s", SessionRevision(0));
847        store
848            .append("s", SessionRevision(0), activity("a"))
849            .expect("append a");
850        {
851            let sessions = store.sessions.lock().expect("sessions");
852            assert!(sessions.get("s").expect("buffer").sender.is_none());
853        }
854        let LiveReplaySubscribeResult::Subscribed(subscription) =
855            store.subscribe_after_cursor(&start).expect("subscribe")
856        else {
857            panic!("expected subscription");
858        };
859        {
860            let sessions = store.sessions.lock().expect("sessions");
861            assert!(sessions.get("s").expect("buffer").sender.is_some());
862        }
863        drop(subscription);
864        store
865            .append("s", SessionRevision(0), activity("b"))
866            .expect("append b");
867        let sessions = store.sessions.lock().expect("sessions");
868        assert!(sessions.get("s").expect("buffer").sender.is_none());
869    }
870
871    #[test]
872    fn in_memory_replay_subscription_reports_gap_after_capacity_trim() {
873        let store = InMemoryLiveReplayStore::with_bounds(1, Duration::from_secs(120));
874        let start = store.current_cursor("s", SessionRevision(0));
875        store
876            .append("s", SessionRevision(0), activity("a"))
877            .expect("append a");
878        store
879            .append("s", SessionRevision(0), activity("b"))
880            .expect("append b");
881        assert!(matches!(
882            store.subscribe_after_cursor(&start).expect("subscribe"),
883            LiveReplaySubscribeResult::Gap(LiveReplayGapReason::Trimmed)
884        ));
885    }
886
887    #[test]
888    fn in_memory_replay_subscription_reports_gap_after_ttl_trim() {
889        let store = InMemoryLiveReplayStore::with_bounds(16, Duration::from_millis(1));
890        let start = store.current_cursor("s", SessionRevision(0));
891        store
892            .append("s", SessionRevision(0), activity("a"))
893            .expect("append a");
894        std::thread::sleep(Duration::from_millis(5));
895        assert!(matches!(
896            store.subscribe_after_cursor(&start).expect("subscribe"),
897            LiveReplaySubscribeResult::Gap(LiveReplayGapReason::Trimmed)
898        ));
899    }
900}
lash_core/runtime/observation/replay.rs

lash_core/runtime/observation/
replay.rs
