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

lash_core/runtime/observation/
replay.rs
