Skip to main content

lash_core/runtime/
observation.rs

1mod replay;
2
3use arc_swap::ArcSwap;
4use std::sync::Arc;
5use tokio::sync::Mutex;
6
7use super::{LashRuntime, ProcessHandleGrantEntry, ProcessHandleSummary, ProcessRegistry};
8
9pub use replay::{
10    InMemoryLiveReplayStore, InMemoryLiveReplayStoreConfig, LiveReplayGap, LiveReplayGapReason,
11    LiveReplayResult, LiveReplayStore, LiveReplayStoreError, LiveReplaySubscribeResult,
12    LiveReplaySubscription, SessionCursor, SessionCursorError, SessionObservation,
13    SessionObservationEvent, SessionObservationEventPayload, SessionObservationSubscription,
14    SessionProcessEventKind, SessionQueueEventKind, SessionResume, SessionRevision,
15};
16
17#[derive(Clone)]
18pub struct RuntimeObservation {
19    pub session_id: Arc<str>,
20    pub revision: SessionRevision,
21    pub cursor: SessionCursor,
22    pub policy: crate::SessionPolicy,
23    pub read_view: crate::SessionReadView,
24    pub persisted_state: super::RuntimeSessionState,
25    pub usage_report: super::SessionUsageReport,
26    pub tool_state: Option<crate::ToolState>,
27    pub tool_catalog: Arc<Vec<serde_json::Value>>,
28    pub tool_catalog_error: Option<String>,
29    pub plugin_session: Option<Arc<crate::PluginSession>>,
30    pub session_read_service: Option<Arc<dyn crate::plugin::SessionReadService>>,
31    pub process_read_service: Option<Arc<dyn crate::plugin::ProcessReadService>>,
32    pub process_registry: Option<Arc<dyn ProcessRegistry>>,
33    pub queue_store: Option<Arc<dyn crate::RuntimePersistence>>,
34    pub queued_work_driver: Option<super::QueuedWorkDriver>,
35}
36
37impl RuntimeObservation {
38    fn from_runtime(
39        runtime: &LashRuntime,
40        cursor: SessionCursor,
41        previous: Option<&RuntimeObservation>,
42        revision: SessionRevision,
43        read_view: crate::SessionReadView,
44        persisted_state: super::RuntimeSessionState,
45        usage_report: super::SessionUsageReport,
46    ) -> Self {
47        let (tool_catalog, tool_catalog_error) = match runtime.active_tool_catalog_shared() {
48            Ok(catalog) => (catalog, None),
49            Err(err) => (Arc::new(Vec::new()), Some(err.to_string())),
50        };
51        let tool_state_generation = runtime
52            .session
53            .as_ref()
54            .map(|session| session.plugins().tool_registry().generation());
55        let tool_state = match (
56            tool_state_generation,
57            previous.and_then(|observation| observation.tool_state.as_ref()),
58        ) {
59            (Some(generation), Some(snapshot)) if snapshot.generation() == generation => {
60                Some(snapshot.clone())
61            }
62            (Some(_), _) => match runtime.tool_state() {
63                Ok(state) => Some(state),
64                Err(err) => {
65                    tracing::warn!(
66                        session_id = %runtime.session_id(),
67                        error = %err,
68                        "failed to capture tool state for observation; omitting the snapshot",
69                    );
70                    None
71                }
72            },
73            (None, _) => None,
74        };
75        let (plugin_session, session_read_service, process_read_service) =
76            match (runtime.session.as_ref(), runtime.runtime_session_services()) {
77                (Some(session), Ok(services)) => (
78                    Some(Arc::clone(session.plugins())),
79                    Some(services.read_service()),
80                    Some(services.process_read_service()),
81                ),
82                (_, Err(err)) => {
83                    tracing::warn!(
84                        session_id = %runtime.session_id(),
85                        error = %err,
86                        "failed to capture plugin query services for observation",
87                    );
88                    (None, None, None)
89                }
90                (None, _) => (None, None, None),
91            };
92        Self {
93            session_id: Arc::from(runtime.session_id()),
94            revision,
95            cursor,
96            policy: read_view.policy().clone(),
97            read_view,
98            persisted_state,
99            usage_report,
100            tool_state,
101            tool_catalog,
102            tool_catalog_error,
103            plugin_session,
104            session_read_service,
105            process_read_service,
106            process_registry: runtime.host.process_registry.clone(),
107            queue_store: runtime
108                .session
109                .as_ref()
110                .and_then(|session| session.history_store()),
111            queued_work_driver: runtime.host.queued_work_driver.clone(),
112        }
113    }
114
115    pub fn session_id(&self) -> &str {
116        &self.session_id
117    }
118
119    pub fn session_revision(&self) -> SessionRevision {
120        self.revision
121    }
122
123    pub fn cursor(&self) -> &SessionCursor {
124        &self.cursor
125    }
126
127    pub fn session_observation(&self) -> SessionObservation {
128        SessionObservation {
129            read_view: self.read_view.clone(),
130            cursor: self.cursor.clone(),
131        }
132    }
133
134    pub fn process_scope(&self) -> crate::SessionScope {
135        crate::SessionScope::new(self.session_id.as_ref())
136    }
137
138    pub fn process_scope_id(&self) -> crate::SessionScopeId {
139        self.process_scope().id()
140    }
141
142    pub async fn query_plugin(
143        &self,
144        name: &str,
145        args: serde_json::Value,
146        session_id: Option<String>,
147    ) -> Result<(String, serde_json::Value), crate::PluginOperationInvokeError> {
148        let Some(plugin_session) = self.plugin_session.as_ref().cloned() else {
149            return Err(crate::PluginOperationInvokeError::Unknown(
150                "runtime session not available".to_string(),
151            ));
152        };
153        let Some(session_read_service) = self.session_read_service.as_ref().cloned() else {
154            return Err(crate::PluginOperationInvokeError::Unknown(
155                "runtime session read service not available".to_string(),
156            ));
157        };
158        let Some(process_read_service) = self.process_read_service.as_ref().cloned() else {
159            return Err(crate::PluginOperationInvokeError::Unknown(
160                "runtime process read service not available".to_string(),
161            ));
162        };
163        plugin_session
164            .query_plugin(
165                name,
166                args,
167                session_id,
168                true,
169                session_read_service,
170                process_read_service,
171            )
172            .await
173    }
174
175    pub async fn list_process_handles(&self) -> Vec<ProcessHandleSummary> {
176        let Some(executor) = self.process_registry.as_ref() else {
177            return Vec::new();
178        };
179        self.list_process_handles_with_mode(executor, crate::ProcessListMode::Live)
180            .await
181    }
182
183    pub async fn list_all_process_handles(&self) -> Vec<ProcessHandleSummary> {
184        let Some(executor) = self.process_registry.as_ref() else {
185            return Vec::new();
186        };
187        self.list_process_handles_with_mode(executor, crate::ProcessListMode::All)
188            .await
189    }
190
191    async fn list_process_handles_with_mode(
192        &self,
193        executor: &Arc<dyn crate::ProcessRegistry>,
194        mode: crate::ProcessListMode,
195    ) -> Vec<ProcessHandleSummary> {
196        let root_scope = self.process_scope();
197        let mut entries = list_scope_process_handles(executor, &root_scope, mode).await;
198        let agent_frame_id = self.persisted_state.current_agent_frame_id.as_str();
199        if !agent_frame_id.is_empty() {
200            let frame_scope =
201                crate::SessionScope::for_agent_frame(self.session_id.as_ref(), agent_frame_id);
202            if frame_scope.id() != root_scope.id() {
203                entries.extend(list_scope_process_handles(executor, &frame_scope, mode).await);
204                entries.sort_by(|(left, _), (right, _)| left.process_id.cmp(&right.process_id));
205                entries.dedup_by(|(left, _), (right, _)| left.process_id == right.process_id);
206            }
207        }
208        entries
209            .into_iter()
210            .map(ProcessHandleSummary::from)
211            .collect()
212    }
213}
214
215fn export_observation_state(
216    runtime: &LashRuntime,
217) -> (
218    super::RuntimeSessionState,
219    crate::SessionReadView,
220    super::SessionUsageReport,
221) {
222    let mut state = runtime.export_persisted_state();
223    let read_view = runtime.read_view();
224    let shared_ledger = runtime
225        .shared_token_ledger
226        .lock()
227        .expect("token ledger lock");
228    for entry in shared_ledger.iter().cloned() {
229        super::merge_ledger_entry(&mut state.token_ledger, entry);
230    }
231    let usage_report = super::SessionUsageReport::from_entries(&state.token_ledger);
232    (state, read_view, usage_report)
233}
234
235async fn list_scope_process_handles(
236    executor: &Arc<dyn crate::ProcessRegistry>,
237    scope: &crate::SessionScope,
238    mode: crate::ProcessListMode,
239) -> Vec<ProcessHandleGrantEntry> {
240    match mode {
241        crate::ProcessListMode::Live => executor.list_live_handle_grants(scope).await,
242        crate::ProcessListMode::All => executor.list_handle_grants(scope).await,
243    }
244    .unwrap_or_default()
245}
246
247#[derive(Clone)]
248pub struct RuntimeHandle {
249    pub(in crate::runtime) runtime: Arc<Mutex<LashRuntime>>,
250    observation: Arc<ArcSwap<RuntimeObservation>>,
251    live_replay_store: Arc<dyn LiveReplayStore>,
252}
253
254impl RuntimeHandle {
255    pub fn new(runtime: LashRuntime) -> Self {
256        Self::with_live_replay_store(runtime, Arc::new(InMemoryLiveReplayStore::default()))
257    }
258
259    pub fn with_live_replay_store(
260        runtime: LashRuntime,
261        live_replay_store: Arc<dyn LiveReplayStore>,
262    ) -> Self {
263        let revision = SessionRevision::from_runtime(&runtime);
264        let cursor = live_replay_store.current_cursor(runtime.session_id(), revision);
265        let (state, read_view, usage_report) = export_observation_state(&runtime);
266        let observation = RuntimeObservation::from_runtime(
267            &runtime,
268            cursor,
269            None,
270            revision,
271            read_view,
272            state,
273            usage_report,
274        );
275        Self {
276            runtime: Arc::new(Mutex::new(runtime)),
277            observation: Arc::new(ArcSwap::from_pointee(observation)),
278            live_replay_store,
279        }
280    }
281
282    pub fn writer(&self) -> Arc<Mutex<LashRuntime>> {
283        Arc::clone(&self.runtime)
284    }
285
286    pub fn observe(&self) -> Arc<RuntimeObservation> {
287        self.observation.load_full()
288    }
289
290    pub fn publish_from(&self, runtime: &LashRuntime) {
291        let revision = SessionRevision::from_runtime(runtime);
292        let previous = self.observation.load_full();
293        let (state, read_view, usage_report) = export_observation_state(runtime);
294        if previous.persisted_state.current_agent_frame_id != state.current_agent_frame_id
295            && !state.current_agent_frame_id.is_empty()
296            && let Err(err) = self.live_replay_store.append(
297                runtime.session_id(),
298                revision,
299                SessionObservationEventPayload::AgentFrameSwitched {
300                    frame_id: state.current_agent_frame_id.clone(),
301                },
302            )
303        {
304            tracing::warn!(
305                session_id = %runtime.session_id(),
306                error = %err,
307                "failed to append agent-frame observation event; reconnect may require gap recovery",
308            );
309        }
310        let cursor = match self.live_replay_store.append(
311            runtime.session_id(),
312            revision,
313            SessionObservationEventPayload::Committed {
314                read_view: read_view.clone(),
315            },
316        ) {
317            Ok(event) => event.cursor.clone(),
318            Err(err) => {
319                tracing::warn!(
320                    session_id = %runtime.session_id(),
321                    error = %err,
322                    "failed to append session observation commit event; reconnect will fall back to gap recovery",
323                );
324                self.live_replay_store
325                    .current_cursor(runtime.session_id(), revision)
326            }
327        };
328        self.observation
329            .store(Arc::new(RuntimeObservation::from_runtime(
330                runtime,
331                cursor,
332                Some(previous.as_ref()),
333                revision,
334                read_view,
335                state,
336                usage_report,
337            )));
338    }
339
340    pub fn record_turn_activity(&self, activity: crate::TurnActivity) {
341        let observation = self.observe();
342        if let Err(err) = self.live_replay_store.append(
343            observation.session_id(),
344            observation.session_revision(),
345            SessionObservationEventPayload::TurnActivity(activity),
346        ) {
347            tracing::warn!(
348                session_id = %observation.session_id(),
349                error = %err,
350                "failed to append live turn activity to session observation replay; reconnect may require gap recovery",
351            );
352        }
353    }
354
355    pub fn record_queue_changed(&self, kind: SessionQueueEventKind, batch_ids: Vec<String>) {
356        let observation = self.observe();
357        if let Err(err) = self.live_replay_store.append(
358            observation.session_id(),
359            observation.session_revision(),
360            SessionObservationEventPayload::QueueChanged { kind, batch_ids },
361        ) {
362            tracing::warn!(
363                session_id = %observation.session_id(),
364                error = %err,
365                "failed to append queue observation event; reconnect may require gap recovery",
366            );
367        }
368    }
369
370    pub fn record_process_changed(&self, kind: SessionProcessEventKind, process_ids: Vec<String>) {
371        let observation = self.observe();
372        if let Err(err) = self.live_replay_store.append(
373            observation.session_id(),
374            observation.session_revision(),
375            SessionObservationEventPayload::ProcessChanged { kind, process_ids },
376        ) {
377            tracing::warn!(
378                session_id = %observation.session_id(),
379                error = %err,
380                "failed to append process observation event; reconnect may require gap recovery",
381            );
382        }
383    }
384
385    pub fn current_session_observation(&self) -> SessionObservation {
386        self.observe().session_observation()
387    }
388
389    pub fn resume_session_observation(
390        &self,
391        cursor: &SessionCursor,
392    ) -> Result<SessionResume, LiveReplayStoreError> {
393        let observation = self.observe();
394        cursor.parse_for_session(observation.session_id())?;
395        match self.live_replay_store.replay_after_cursor(cursor)? {
396            LiveReplayResult::Replayed(events) => Ok(SessionResume::Replayed { events }),
397            LiveReplayResult::Gap(reason) => Ok(SessionResume::Gap {
398                gap: self.live_replay_gap(cursor, reason, observation.as_ref()),
399                observation: observation.session_observation(),
400            }),
401        }
402    }
403
404    pub fn subscribe_session_observation(
405        &self,
406        cursor: &SessionCursor,
407    ) -> Result<SessionObservationSubscription, LiveReplayStoreError> {
408        let observation = self.observe();
409        cursor.parse_for_session(observation.session_id())?;
410        match self.live_replay_store.subscribe_after_cursor(cursor)? {
411            LiveReplaySubscribeResult::Subscribed(subscription) => {
412                Ok(SessionObservationSubscription::Subscribed(subscription))
413            }
414            LiveReplaySubscribeResult::Gap(reason) => Ok(SessionObservationSubscription::Gap {
415                gap: self.live_replay_gap(cursor, reason, observation.as_ref()),
416                observation: observation.session_observation(),
417            }),
418        }
419    }
420
421    fn live_replay_gap(
422        &self,
423        requested_cursor: &SessionCursor,
424        reason: LiveReplayGapReason,
425        observation: &RuntimeObservation,
426    ) -> LiveReplayGap {
427        LiveReplayGap {
428            session_id: observation.session_id().to_string(),
429            requested_cursor: requested_cursor.clone(),
430            latest_cursor: observation.cursor().clone(),
431            latest_revision: observation.session_revision(),
432            reason,
433        }
434    }
435
436    pub async fn enqueue_turn_input(
437        &self,
438        input: crate::TurnInput,
439        ingress: crate::TurnInputIngress,
440        source_key: Option<String>,
441    ) -> Result<crate::PendingTurnInput, crate::RuntimeError> {
442        let observation = self.observe();
443        let store = observation
444            .queue_store
445            .clone()
446            .ok_or_else(super::session_api::queued_turn_input_store_required)?;
447        let is_next_turn = matches!(ingress, crate::TurnInputIngress::NextTurn);
448        super::session_api::enqueue_turn_input_to_store(
449            observation.session_id.as_ref().to_string(),
450            store,
451            observation.queued_work_driver.clone(),
452            input,
453            ingress,
454            source_key,
455        )
456        .await
457        .inspect(|input| {
458            self.record_queue_changed(
459                SessionQueueEventKind::Enqueued,
460                if is_next_turn {
461                    vec![input.input_id.clone()]
462                } else {
463                    Vec::new()
464                },
465            );
466        })
467    }
468
469    pub async fn cancel_pending_turn_input(
470        &self,
471        session_id: &str,
472        input_id: &str,
473    ) -> Result<crate::PendingTurnInputCancelOutcome, crate::RuntimeError> {
474        let observation = self.observe();
475        let store = observation
476            .queue_store
477            .clone()
478            .ok_or_else(super::session_api::queued_turn_input_store_required)?;
479        store
480            .cancel_pending_turn_input(session_id, input_id)
481            .await
482            .map_err(|err| {
483                crate::RuntimeError::new(
484                    crate::RuntimeErrorCode::StoreCommitFailed,
485                    err.to_string(),
486                )
487            })
488            .inspect(|outcome| {
489                if outcome.is_cancelled() {
490                    self.record_queue_changed(
491                        SessionQueueEventKind::Cancelled,
492                        vec![input_id.to_string()],
493                    );
494                }
495            })
496    }
497
498    pub async fn cancel_pending_turn_inputs(
499        &self,
500        session_id: &str,
501        targets: &[crate::PendingTurnInputCancelTarget],
502    ) -> Result<Vec<crate::PendingTurnInputCancelResult>, crate::RuntimeError> {
503        let observation = self.observe();
504        let store = observation
505            .queue_store
506            .clone()
507            .ok_or_else(super::session_api::queued_turn_input_store_required)?;
508        store
509            .cancel_pending_turn_inputs(session_id, targets)
510            .await
511            .map_err(|err| {
512                crate::RuntimeError::new(
513                    crate::RuntimeErrorCode::StoreCommitFailed,
514                    err.to_string(),
515                )
516            })
517            .inspect(|results| {
518                let cancelled_ids = results
519                    .iter()
520                    .filter_map(|result| match &result.outcome {
521                        crate::PendingTurnInputCancelOutcome::Cancelled(input) => {
522                            Some(input.input_id.clone())
523                        }
524                        _ => None,
525                    })
526                    .collect::<Vec<_>>();
527                if !cancelled_ids.is_empty() {
528                    self.record_queue_changed(SessionQueueEventKind::Cancelled, cancelled_ids);
529                }
530            })
531    }
532
533    pub async fn cancel_pending_turn_input_suffix(
534        &self,
535        session_id: &str,
536        anchor: &crate::PendingTurnInputCancelTarget,
537    ) -> Result<crate::PendingTurnInputSuffixCancelOutcome, crate::RuntimeError> {
538        let observation = self.observe();
539        let store = observation
540            .queue_store
541            .clone()
542            .ok_or_else(super::session_api::queued_turn_input_store_required)?;
543        store
544            .cancel_pending_turn_input_suffix(session_id, anchor)
545            .await
546            .map_err(|err| {
547                crate::RuntimeError::new(
548                    crate::RuntimeErrorCode::StoreCommitFailed,
549                    err.to_string(),
550                )
551            })
552            .inspect(|outcome| {
553                let crate::PendingTurnInputSuffixCancelOutcome::Outcomes { outcomes, .. } = outcome
554                else {
555                    return;
556                };
557                let cancelled_ids = outcomes
558                    .iter()
559                    .filter_map(|outcome| match outcome {
560                        crate::PendingTurnInputCancelOutcome::Cancelled(input) => {
561                            Some(input.input_id.clone())
562                        }
563                        _ => None,
564                    })
565                    .collect::<Vec<_>>();
566                if !cancelled_ids.is_empty() {
567                    self.record_queue_changed(SessionQueueEventKind::Cancelled, cancelled_ids);
568                }
569            })
570    }
571
572    /// Release a held queued-work claim without completing it, returning its
573    /// batches to the pending queue immediately.
574    ///
575    /// This is the host lever behind stopping an external queued-work driver
576    /// mid-claim: instead of letting the claim age out over its lease TTL, the
577    /// host hands the claim back and the work becomes claimable again at once.
578    pub async fn abandon_queued_work_claim(
579        &self,
580        claim: &crate::QueuedWorkClaim,
581    ) -> Result<(), crate::RuntimeError> {
582        let observation = self.observe();
583        let store = observation
584            .queue_store
585            .clone()
586            .ok_or_else(super::session_api::queued_turn_input_store_required)?;
587        store
588            .abandon_queued_work_claim(claim)
589            .await
590            .map_err(|err| {
591                crate::RuntimeError::new(
592                    crate::RuntimeErrorCode::StoreCommitFailed,
593                    err.to_string(),
594                )
595            })?;
596        self.record_queue_changed(
597            SessionQueueEventKind::Enqueued,
598            claim
599                .batches
600                .iter()
601                .map(|batch| batch.batch_id.clone())
602                .collect(),
603        );
604        Ok(())
605    }
606
607    /// Release a held pending-turn-input claim without completing it, returning
608    /// its inputs to the pending queue immediately. The turn-input counterpart
609    /// of [`abandon_queued_work_claim`](Self::abandon_queued_work_claim).
610    pub async fn abandon_turn_input_claim(
611        &self,
612        claim: &crate::TurnInputClaim,
613    ) -> Result<(), crate::RuntimeError> {
614        let observation = self.observe();
615        let store = observation
616            .queue_store
617            .clone()
618            .ok_or_else(super::session_api::queued_turn_input_store_required)?;
619        store.abandon_turn_input_claim(claim).await.map_err(|err| {
620            crate::RuntimeError::new(crate::RuntimeErrorCode::StoreCommitFailed, err.to_string())
621        })?;
622        self.record_queue_changed(
623            SessionQueueEventKind::Enqueued,
624            claim
625                .inputs
626                .iter()
627                .map(|input| input.input_id.clone())
628                .collect(),
629        );
630        Ok(())
631    }
632
633    pub async fn cancel_queued_work_batch(
634        &self,
635        session_id: &str,
636        batch_id: &str,
637    ) -> Result<Option<crate::QueuedWorkBatch>, crate::RuntimeError> {
638        let observation = self.observe();
639        let store = observation
640            .queue_store
641            .clone()
642            .ok_or_else(super::session_api::queued_turn_input_store_required)?;
643        store
644            .cancel_queued_work_batch(session_id, batch_id)
645            .await
646            .map_err(|err| {
647                crate::RuntimeError::new(
648                    crate::RuntimeErrorCode::StoreCommitFailed,
649                    err.to_string(),
650                )
651            })
652            .inspect(|batch| {
653                if batch.is_some() {
654                    self.record_queue_changed(
655                        SessionQueueEventKind::Cancelled,
656                        vec![batch_id.to_string()],
657                    );
658                }
659            })
660    }
661
662    pub fn try_into_runtime(self) -> Result<LashRuntime, Self> {
663        match Arc::try_unwrap(self.runtime) {
664            Ok(mutex) => Ok(mutex.into_inner()),
665            Err(runtime) => Err(Self {
666                runtime,
667                observation: self.observation,
668                live_replay_store: self.live_replay_store,
669            }),
670        }
671    }
672}
673
674#[cfg(test)]
675mod tests {
676    use super::*;
677    use std::time::Instant;
678
679    struct PanicLiveReplayStore;
680
681    impl LiveReplayStore for PanicLiveReplayStore {
682        fn append(
683            &self,
684            _session_id: &str,
685            _revision: SessionRevision,
686            _payload: SessionObservationEventPayload,
687        ) -> Result<Arc<SessionObservationEvent>, LiveReplayStoreError> {
688            panic!("append should not be called by cursor rejection tests")
689        }
690
691        fn replay_after_cursor(
692            &self,
693            _cursor: &SessionCursor,
694        ) -> Result<LiveReplayResult, LiveReplayStoreError> {
695            panic!("replay_after_cursor should not be called for rejected cursors")
696        }
697
698        fn subscribe_after_cursor(
699            &self,
700            _cursor: &SessionCursor,
701        ) -> Result<LiveReplaySubscribeResult, LiveReplayStoreError> {
702            panic!("subscribe_after_cursor should not be called for rejected cursors")
703        }
704
705        fn current_cursor(&self, session_id: &str, revision: SessionRevision) -> SessionCursor {
706            SessionCursor::new(session_id, revision, 0)
707        }
708
709        fn trim_session(&self, _session_id: &str) -> Result<(), LiveReplayStoreError> {
710            Ok(())
711        }
712    }
713
714    #[tokio::test]
715    async fn runtime_rejects_bad_cursors_before_replay_store_gap_handling() {
716        let runtime = LashRuntime::builder()
717            .with_session_id("session-a")
718            .with_policy(crate::SessionPolicy {
719                model: crate::ModelSpec::from_token_limits(
720                    "test-model",
721                    Default::default(),
722                    1024,
723                    None,
724                )
725                .expect("model"),
726                ..Default::default()
727            })
728            .build()
729            .await
730            .expect("runtime");
731        let handle = RuntimeHandle::with_live_replay_store(runtime, Arc::new(PanicLiveReplayStore));
732        let wrong_session = SessionCursor::new("session-b", SessionRevision(0), 99);
733        let malformed = SessionCursor::from_raw_for_testing("bad");
734
735        assert!(matches!(
736            handle.resume_session_observation(&wrong_session),
737            Err(LiveReplayStoreError::Cursor(
738                SessionCursorError::WrongSession { .. }
739            ))
740        ));
741        assert!(matches!(
742            handle.subscribe_session_observation(&wrong_session),
743            Err(LiveReplayStoreError::Cursor(
744                SessionCursorError::WrongSession { .. }
745            ))
746        ));
747        assert!(matches!(
748            handle.resume_session_observation(&malformed),
749            Err(LiveReplayStoreError::Cursor(
750                SessionCursorError::Malformed { .. }
751            ))
752        ));
753        assert!(matches!(
754            handle.subscribe_session_observation(&malformed),
755            Err(LiveReplayStoreError::Cursor(
756                SessionCursorError::Malformed { .. }
757            ))
758        ));
759    }
760
761    #[tokio::test]
762    async fn publish_revision_matches_the_single_export_across_a_commit() {
763        let runtime = LashRuntime::builder()
764            .with_session_id("revision-equivalence")
765            .with_policy(crate::SessionPolicy {
766                model: crate::ModelSpec::from_token_limits(
767                    "test-model",
768                    Default::default(),
769                    1024,
770                    None,
771                )
772                .expect("model"),
773                ..Default::default()
774            })
775            .build()
776            .await
777            .expect("runtime");
778        let handle = RuntimeHandle::new(runtime);
779        let writer = handle.writer();
780        let mut runtime = writer.lock().await;
781        runtime.state.turn_index = 9;
782        runtime.state.head_revision = Some(17);
783
784        let exported = runtime.export_persisted_state();
785        let exported_revision = SessionRevision::from_state(&exported);
786        let accessor_revision = SessionRevision::from_runtime(&runtime);
787        assert_eq!(accessor_revision, exported_revision);
788
789        handle.publish_from(&runtime);
790        assert_eq!(handle.observe().session_revision(), exported_revision);
791    }
792
793    #[tokio::test]
794    async fn publish_keeps_frame_switch_immediately_before_commit() {
795        let runtime = LashRuntime::builder()
796            .with_session_id("publish-order")
797            .with_policy(crate::SessionPolicy {
798                model: crate::ModelSpec::from_token_limits(
799                    "test-model",
800                    Default::default(),
801                    1024,
802                    None,
803                )
804                .expect("model"),
805                ..Default::default()
806            })
807            .build()
808            .await
809            .expect("runtime");
810        let handle = RuntimeHandle::new(runtime);
811        let cursor = handle.observe().cursor().clone();
812        let writer = handle.writer();
813        let mut runtime = writer.lock().await;
814        runtime.state.current_agent_frame_id = "next-frame".to_string();
815
816        handle.publish_from(&runtime);
817        let SessionResume::Replayed { events } = handle
818            .resume_session_observation(&cursor)
819            .expect("replay publication")
820        else {
821            panic!("publication should remain replayable");
822        };
823        assert_eq!(events.len(), 2);
824        assert!(matches!(
825            events[0].payload,
826            SessionObservationEventPayload::AgentFrameSwitched { .. }
827        ));
828        assert!(matches!(
829            events[1].payload,
830            SessionObservationEventPayload::Committed { .. }
831        ));
832    }
833
834    #[tokio::test]
835    #[ignore = "manual lane-O publish_from timing measurement"]
836    async fn measure_publish_from_wall_clock() {
837        const COMMITS: usize = 5_000;
838        let runtime = LashRuntime::builder()
839            .with_session_id("publish-perf")
840            .with_policy(crate::SessionPolicy {
841                model: crate::ModelSpec::from_token_limits(
842                    "test-model",
843                    Default::default(),
844                    1024,
845                    None,
846                )
847                .expect("model"),
848                ..Default::default()
849            })
850            .build()
851            .await
852            .expect("runtime");
853        let handle = RuntimeHandle::with_live_replay_store(
854            runtime,
855            Arc::new(InMemoryLiveReplayStore::with_bounds(
856                COMMITS + 1,
857                std::time::Duration::from_secs(120),
858            )),
859        );
860        let writer = handle.writer();
861        let runtime = writer.lock().await;
862        for _ in 0..100 {
863            handle.publish_from(&runtime);
864        }
865        let started = Instant::now();
866        for _ in 0..COMMITS {
867            handle.publish_from(&runtime);
868        }
869        let elapsed = started.elapsed();
870        eprintln!(
871            "publish_from: commits={COMMITS} elapsed_ns={} ns_per_commit={:.3}",
872            elapsed.as_nanos(),
873            elapsed.as_nanos() as f64 / COMMITS as f64,
874        );
875    }
876}