Skip to main content

lash/
session.rs

1use std::pin::Pin;
2use std::task::{Context, Poll};
3
4use crate::support::*;
5use futures_util::Stream;
6use lash_core::runtime::{
7    PendingTurnInput, PendingTurnInputCancelOutcome, PendingTurnInputCancelResult,
8    PendingTurnInputCancelTarget, PendingTurnInputSuffixCancelOutcome, QueuedWorkBatch,
9    TurnInputIngress,
10};
11use lash_core::{LiveReplayGap, LiveReplayStoreError, SessionObservationEvent};
12use lash_remote_protocol::{
13    RemoteLiveReplayGap, RemoteSessionCursor, RemoteSessionObservation,
14    RemoteSessionObservationEvent,
15};
16
17pub struct SessionBuilder {
18    pub(crate) core: LashCore,
19    pub(crate) session_id: String,
20    pub(crate) spec: SessionSpec,
21    pub(crate) parent_session_id: Option<String>,
22    pub(crate) session_execution_owner: Option<lash_core::LeaseOwnerIdentity>,
23    pub(crate) store: Option<Arc<dyn RuntimePersistence>>,
24    pub(crate) provider: Option<ProviderHandle>,
25    pub(crate) active_plugins: Vec<ActivePluginBinding>,
26    pub(crate) plugin_factories: Vec<Arc<dyn PluginFactory>>,
27    /// Plugin-keyed, serializable open-time options. They ride the protocol
28    /// materialization seam (the same `PluginOptions` bag a child
29    /// `SessionCreateRequest` carries) so every plugin gets open-time options
30    /// through one hook.
31    pub(crate) plugin_options: PluginOptions,
32}
33
34impl SessionBuilder {
35    /// Set the plugin-keyed open-time options bag wholesale.
36    pub fn plugin_options(mut self, plugin_options: PluginOptions) -> Self {
37        self.plugin_options = plugin_options;
38        self
39    }
40
41    /// Merge a single plugin's typed options into the open-time options bag,
42    /// preserving any options already set for other plugin keys.
43    pub fn plugin_option<T: serde::Serialize>(
44        mut self,
45        plugin_id: impl Into<String>,
46        extras: T,
47    ) -> Result<Self> {
48        self.plugin_options
49            .insert_typed(plugin_id, extras)
50            .map_err(EmbedError::ProtocolTurnOptions)?;
51        Ok(self)
52    }
53
54    pub fn provider(mut self, provider: ProviderHandle) -> Self {
55        self.spec = self.spec.provider_id(provider.kind());
56        self.provider = Some(provider);
57        self
58    }
59
60    pub fn session_spec(mut self, spec: SessionSpec) -> Self {
61        self.spec = spec;
62        self
63    }
64
65    pub fn parent(mut self, parent_session_id: impl Into<String>) -> Self {
66        self.parent_session_id = Some(parent_session_id.into());
67        self
68    }
69
70    /// Use an explicit owner identity for durable session execution leases.
71    ///
72    /// This is only for hosts that already serialize one logical execution lane
73    /// and intentionally choose stable owner + incarnation values. Normal
74    /// embedders should keep the default per-open identity.
75    pub fn session_execution_owner(mut self, owner: lash_core::LeaseOwnerIdentity) -> Self {
76        self.session_execution_owner = Some(owner);
77        self
78    }
79
80    /// Use a specific persistence store for this root session.
81    ///
82    /// This is the right API for a host-owned, pre-opened session database.
83    /// Managed child sessions never reuse this store; configure
84    /// `LashCoreBuilder::child_store_factory` when child sessions should also
85    /// persist.
86    pub fn store(mut self, store: Arc<dyn RuntimePersistence>) -> Self {
87        self.store = Some(store);
88        self
89    }
90
91    pub fn plugin<P: PluginBinding>(mut self, config: P::SessionConfig) -> Self {
92        self.active_plugins.push(ActivePluginBinding {
93            id: P::ID,
94            requires_turn_input: P::requires_turn_input(&config),
95        });
96        self.plugin_factories.push(P::factory(&config));
97        self
98    }
99
100    pub async fn open(self) -> Result<LashSession> {
101        let policy = self.session_policy();
102        let store = self.create_store(&policy).await?;
103        let state = self
104            .load_or_default_state(&policy, store.as_deref())
105            .await?;
106        self.open_resolved(policy, state, store).await
107    }
108
109    /// Open this session with a fresh resident graph, ignoring any persisted
110    /// session graph/checkpoint state that may already exist for the same
111    /// session id.
112    ///
113    /// The next successful commit writes a full replacement graph, so normal
114    /// embedders can use this to start over without manually calling
115    /// `load_persisted_session_state` or constructing a `RuntimeSessionState`.
116    /// Use [`Self::open`] for resume and [`Self::open_with_state`] only when
117    /// restoring explicit host-owned state.
118    pub async fn open_fresh(self) -> Result<LashSession> {
119        let policy = self.session_policy();
120        let store = self.create_store(&policy).await?;
121        let state = RuntimeSessionState {
122            session_id: self.session_id.clone(),
123            policy: policy.clone(),
124            graph_replace_required: true,
125            ..RuntimeSessionState::default()
126        };
127        self.open_resolved(policy, state, store).await
128    }
129
130    /// Open with an explicitly supplied runtime state.
131    ///
132    /// This is for advanced hosts that already own a complete state snapshot.
133    /// Normal embedders should use [`Self::open`] to resume according to Lash's
134    /// residency policy or [`Self::open_fresh`] to start over and replace prior
135    /// persisted state on the next commit.
136    pub async fn open_with_state(self, mut state: RuntimeSessionState) -> Result<LashSession> {
137        let policy = self.session_policy();
138        let store = self.create_store(&policy).await?;
139        if state.session_id != self.session_id {
140            return Err(EmbedError::StoreSessionMismatch {
141                loaded: state.session_id,
142                requested: self.session_id,
143            });
144        }
145        let recorded_provider_id = state.policy.recorded_provider_id().to_string();
146        state.policy = policy.clone();
147        state.policy.provider_id = recorded_provider_id;
148        self.open_resolved(policy, state, store).await
149    }
150
151    fn session_policy(&self) -> SessionPolicy {
152        let mut policy = self.spec.resolve_against(&self.core.policy);
153        policy.session_id = Some(self.session_id.clone());
154        policy
155    }
156
157    async fn load_or_default_state(
158        &self,
159        policy: &SessionPolicy,
160        store: Option<&dyn RuntimePersistence>,
161    ) -> Result<RuntimeSessionState> {
162        let state = match store {
163            Some(store) => {
164                let loaded = self.load_persisted_state_for_residency(store).await?;
165                let mut state = loaded.unwrap_or_else(|| RuntimeSessionState {
166                    session_id: self.session_id.clone(),
167                    policy: policy.clone(),
168                    ..RuntimeSessionState::default()
169                });
170                if state.session_id != self.session_id {
171                    return Err(EmbedError::StoreSessionMismatch {
172                        loaded: state.session_id,
173                        requested: self.session_id.clone(),
174                    });
175                }
176                let recorded_provider_id = state.policy.recorded_provider_id().to_string();
177                state.policy = policy.clone();
178                state.policy.provider_id = recorded_provider_id;
179                state
180            }
181            None => RuntimeSessionState {
182                session_id: self.session_id.clone(),
183                policy: policy.clone(),
184                ..RuntimeSessionState::default()
185            },
186        };
187        Ok(state)
188    }
189
190    async fn load_persisted_state_for_residency(
191        &self,
192        store: &dyn RuntimePersistence,
193    ) -> Result<Option<RuntimeSessionState>> {
194        load_persisted_state_for_residency(self.core.env.residency, store).await
195    }
196
197    async fn open_resolved(
198        self,
199        policy: SessionPolicy,
200        state: RuntimeSessionState,
201        store: Option<Arc<dyn RuntimePersistence>>,
202    ) -> Result<LashSession> {
203        let mut env = self.core.env.clone();
204        if let Some(provider) = self.provider.clone().or_else(|| self.core.provider.clone()) {
205            env.core.providers.provider_resolver =
206                Arc::new(lash_core::SingleProviderResolver::new(provider));
207        }
208        let plugin_host = build_plugin_host(
209            self.core.protocol_factory.as_ref(),
210            self.core.plugin_factories.as_ref(),
211            self.plugin_factories,
212        )?;
213        env.core = plugin_host.install_process_engine_contributions(
214            env.core.clone(),
215            self.core.process_lifecycle_available,
216        )?;
217        env.plugin_host = Some(Arc::new(plugin_host));
218        let effect_host = Arc::clone(&env.core.control.effect_host);
219        let drivers = self.core.work_driver.drivers().await;
220        env.process_work_driver = drivers.process.clone();
221        env.queued_work_driver = drivers.queued.clone();
222        let mut runtime = LashRuntime::from_environment(&env, policy, state, store).await?;
223        // Fire the protocol materialization hook for this root/builder open
224        // (including resume): the protocol plugin applies and defaults its
225        // per-session options at open time.
226        runtime.configure_protocol_on_materialize(
227            &self.plugin_options,
228            self.parent_session_id.is_none(),
229        )?;
230        if let Some(owner) = self.session_execution_owner {
231            runtime.set_runtime_lease_owner(owner);
232        }
233        if drivers.drive_process_on_open
234            && let Some(driver) = drivers.process.as_ref()
235        {
236            driver.claim_and_run_pending("session_open").await?;
237        }
238        let handle = RuntimeHandle::with_live_replay_store(
239            runtime,
240            Arc::clone(&self.core.live_replay_store),
241        );
242        Ok(LashSession {
243            runtime: handle,
244            effect_host,
245            parent_session_id: self.parent_session_id,
246            active_plugins: self.active_plugins,
247            process_phase_probe_slot: self.core.work_driver.phase_probe_slot(),
248            turn_cancels: crate::turn::TurnCancelRegistry::default(),
249        })
250    }
251
252    async fn create_store(
253        &self,
254        policy: &SessionPolicy,
255    ) -> Result<Option<Arc<dyn RuntimePersistence>>> {
256        if let Some(store) = self.store.as_ref() {
257            return Ok(Some(Arc::clone(store)));
258        }
259        let Some(factory) = self.core.store_factory.as_ref() else {
260            return Ok(None);
261        };
262        let request = SessionStoreCreateRequest {
263            session_id: self.session_id.clone(),
264            relation: self
265                .parent_session_id
266                .as_ref()
267                .map(|parent_session_id| lash_core::SessionRelation::Child {
268                    parent_session_id: parent_session_id.clone(),
269                    caused_by: None,
270                })
271                .unwrap_or_default(),
272            policy: policy.clone(),
273        };
274        factory
275            .create_store(&request)
276            .await
277            .map(Some)
278            .map_err(|message| EmbedError::StoreFactory {
279                session_id: self.session_id.clone(),
280                message,
281            })
282    }
283}
284
285pub(crate) async fn load_state_for_residency(
286    residency: Residency,
287    session_id: &str,
288    policy: &SessionPolicy,
289    store: &dyn RuntimePersistence,
290) -> Result<RuntimeSessionState> {
291    let mut state = load_persisted_state_for_residency(residency, store)
292        .await?
293        .unwrap_or_else(|| RuntimeSessionState {
294            session_id: session_id.to_string(),
295            policy: policy.clone(),
296            ..RuntimeSessionState::default()
297        });
298    if state.session_id != session_id {
299        return Err(EmbedError::StoreSessionMismatch {
300            loaded: state.session_id,
301            requested: session_id.to_string(),
302        });
303    }
304    let recorded_provider_id = state.policy.recorded_provider_id().to_string();
305    state.policy = policy.clone();
306    state.policy.provider_id = recorded_provider_id;
307    Ok(state)
308}
309
310async fn load_persisted_state_for_residency(
311    residency: Residency,
312    store: &dyn RuntimePersistence,
313) -> Result<Option<RuntimeSessionState>> {
314    match residency {
315        Residency::KeepAll => {
316            let loaded = lash_core::store::load_persisted_session_state(store)
317                .await
318                .map_err(|err| SessionError::Protocol(format!("failed to load store: {err}")))?;
319            Ok(loaded)
320        }
321        Residency::ActivePathOnly => {
322            let active = lash_core::store::load_persisted_session_state_active_path(store, None)
323                .await
324                .map_err(|err| {
325                    SessionError::Protocol(format!("failed to load active-path store: {err}"))
326                })?;
327            if active
328                .as_ref()
329                .is_some_and(|state| state.session_graph.nodes.is_empty())
330            {
331                let mut full = lash_core::store::load_persisted_session_state(store)
332                    .await
333                    .map_err(|err| {
334                        SessionError::Protocol(format!(
335                            "failed to heal active-path store from full graph: {err}"
336                        ))
337                    })?;
338                if let Some(state) = full.as_mut() {
339                    state.graph_replace_required = true;
340                }
341                return Ok(full);
342            }
343            Ok(active)
344        }
345    }
346}
347
348impl PromptLayerSink for SessionBuilder {
349    fn prompt_layer_mut(&mut self) -> &mut PromptLayer {
350        self.spec.prompt.get_or_insert_with(PromptLayer::new)
351    }
352}
353
354#[derive(Clone)]
355pub struct LashSession {
356    pub(crate) runtime: RuntimeHandle,
357    pub(crate) effect_host: Arc<dyn EffectHost>,
358    pub(crate) parent_session_id: Option<String>,
359    pub(crate) active_plugins: Vec<ActivePluginBinding>,
360    pub(crate) process_phase_probe_slot: Option<lash_core::runtime::RuntimeTurnPhaseProbeSlot>,
361    pub(crate) turn_cancels: crate::turn::TurnCancelRegistry,
362}
363
364#[derive(Clone, Debug, Default)]
365pub struct SessionConfigPatch {
366    pub provider: Option<ProviderHandle>,
367    pub model: Option<ModelSpec>,
368    pub prompt: Option<PromptLayer>,
369}
370
371impl LashSession {
372    pub async fn close(self) -> Result<()> {
373        let runtime = self.runtime.writer();
374        let runtime = runtime.lock().await;
375        runtime.unregister_plugin_session()?;
376        Ok(())
377    }
378
379    pub fn session_id(&self) -> String {
380        self.runtime.observe().session_id().to_string()
381    }
382
383    pub fn policy_snapshot(&self) -> SessionPolicy {
384        self.runtime.observe().policy.clone()
385    }
386
387    pub fn observe(&self) -> ObservableSession {
388        ObservableSession {
389            runtime: self.runtime.clone(),
390        }
391    }
392
393    pub fn parent_session_id(&self) -> Option<&str> {
394        self.parent_session_id.as_deref()
395    }
396
397    pub fn effect_host(&self) -> Arc<dyn EffectHost> {
398        Arc::clone(&self.effect_host)
399    }
400
401    pub fn turn(&self, input: TurnInput) -> TurnBuilder {
402        TurnBuilder {
403            runtime: self.runtime.clone(),
404            effect_host: Arc::clone(&self.effect_host),
405            active_plugins: self.active_plugins.clone(),
406            input,
407            cancel: CancellationToken::new(),
408            cancels: self.turn_cancels.clone(),
409            protocol_turn_options: None,
410            provider: None,
411            model: None,
412            turn_id: None,
413        }
414    }
415
416    pub fn queued_turn(&self) -> QueuedTurnBuilder {
417        QueuedTurnBuilder {
418            runtime: self.runtime.clone(),
419            effect_host: Arc::clone(&self.effect_host),
420            cancel: CancellationToken::new(),
421            cancels: self.turn_cancels.clone(),
422            batch_ids: Vec::new(),
423            drain_id: None,
424        }
425    }
426
427    /// Cancel every turn currently executing through this opened session
428    /// (including its clones) and report how many were signalled.
429    ///
430    /// This is the affordance behind a UI "stop" control: hold a clone of the
431    /// session wherever the stop arrives and call this, instead of threading a
432    /// [`CancellationToken`](crate::CancellationToken) into every turn call
433    /// ([`TurnBuilder::cancel`](crate::TurnBuilder::cancel) remains the
434    /// per-turn hook when you need one). A cancelled turn finishes with
435    /// `TurnOutcome::Stopped(TurnStop::Cancelled)` and commits like any other
436    /// turn; the session stays usable.
437    ///
438    /// Scope: turns started from this `LashSession` instance and its clones.
439    /// A handle opened separately for the same session id has its own
440    /// registry and is not reached.
441    pub fn cancel_running_turns(&self) -> usize {
442        self.turn_cancels.cancel_all()
443    }
444
445    pub fn admin(&self) -> SessionAdmin {
446        SessionAdmin {
447            runtime: self.runtime.clone(),
448        }
449    }
450
451    pub async fn configure(&self, patch: SessionConfigPatch) -> Result<()> {
452        self.admin().config().update(patch).await
453    }
454
455    pub fn tools(&self) -> ToolAdmin {
456        ToolAdmin::new(self.admin())
457    }
458
459    pub fn commands(&self) -> SessionCommandAdmin {
460        self.admin().commands()
461    }
462
463    pub fn triggers(&self) -> SessionTriggerAdmin {
464        self.admin().triggers()
465    }
466
467    pub fn processes(&self) -> SessionProcessAdmin {
468        SessionProcessAdmin::new(self.admin())
469    }
470
471    pub fn plugin_operations(&self) -> PluginOperations {
472        PluginOperations {
473            control: self.admin(),
474        }
475    }
476
477    pub fn enqueue(&self, input: TurnInput) -> EnqueueTurnBuilder<'_> {
478        EnqueueTurnBuilder {
479            session: self,
480            input,
481            id: None,
482            ingress: TurnInputIngress::NextTurn,
483        }
484    }
485
486    /// Return all pending durable queued-work batches for this session.
487    ///
488    /// This is an admin/introspection view for non-user queued work such as
489    /// process wakes and session commands. User-visible model input is stored
490    /// separately as pending turn input and is exposed by
491    /// [`pending_turn_inputs`](Self::pending_turn_inputs).
492    pub async fn queued_work(&self) -> Result<Vec<QueuedWorkBatch>> {
493        let observation = self.runtime.observe();
494        let store = observation.queue_store.as_ref().ok_or_else(|| {
495            EmbedError::Runtime(lash_core::RuntimeError::new(
496                lash_core::RuntimeErrorCode::StoreCommitFailed,
497                "queued work inspection requires a persistent runtime store",
498            ))
499        })?;
500        store
501            .list_pending_queued_work(observation.session_id())
502            .await
503            .map_err(|err| {
504                EmbedError::Runtime(lash_core::RuntimeError::new(
505                    lash_core::RuntimeErrorCode::StoreCommitFailed,
506                    err.to_string(),
507                ))
508            })
509    }
510
511    pub async fn pending_turn_inputs(&self) -> Result<Vec<PendingTurnInput>> {
512        let observation = self.runtime.observe();
513        let store = observation.queue_store.as_ref().ok_or_else(|| {
514            EmbedError::Runtime(lash_core::RuntimeError::new(
515                lash_core::RuntimeErrorCode::StoreCommitFailed,
516                "pending turn input inspection requires a persistent runtime store",
517            ))
518        })?;
519        store
520            .list_pending_turn_inputs(observation.session_id())
521            .await
522            .map_err(|err| {
523                EmbedError::Runtime(lash_core::RuntimeError::new(
524                    lash_core::RuntimeErrorCode::StoreCommitFailed,
525                    err.to_string(),
526                ))
527            })
528    }
529
530    pub async fn cancel_pending_turn_input(
531        &self,
532        input_id: &str,
533    ) -> Result<PendingTurnInputCancelOutcome> {
534        let session_id = self.session_id();
535        self.runtime
536            .cancel_pending_turn_input(&session_id, input_id)
537            .await
538            .map_err(EmbedError::Runtime)
539    }
540
541    /// Atomically cancel a set of pending user inputs by runtime input id or
542    /// app source key.
543    ///
544    /// This is the app reconciliation path for explicit selections such as
545    /// "remove these pending drafts". Returned outcomes distinguish newly
546    /// cancelled input from input that was already claimed, completed,
547    /// cancelled, or missing.
548    pub async fn cancel_pending_turn_inputs(
549        &self,
550        targets: impl IntoIterator<Item = PendingTurnInputCancelTarget>,
551    ) -> Result<Vec<PendingTurnInputCancelResult>> {
552        let session_id = self.session_id();
553        let targets = targets.into_iter().collect::<Vec<_>>();
554        self.runtime
555            .cancel_pending_turn_inputs(&session_id, &targets)
556            .await
557            .map_err(EmbedError::Runtime)
558    }
559
560    /// Atomically cancel the same-session pending-input suffix from `anchor`.
561    ///
562    /// Apps that let users edit previously submitted product messages should
563    /// map the edited message to the stored pending-input `input_id` or
564    /// `source_key`, call this method, and only restore/edit drafts that return
565    /// [`PendingTurnInputCancelOutcome::Cancelled`]. Claimed or completed
566    /// inputs have already crossed the runtime boundary and should be treated
567    /// as reconciliation state, not local editable drafts.
568    pub async fn cancel_pending_turn_input_suffix(
569        &self,
570        anchor: PendingTurnInputCancelTarget,
571    ) -> Result<PendingTurnInputSuffixCancelOutcome> {
572        let session_id = self.session_id();
573        self.runtime
574            .cancel_pending_turn_input_suffix(&session_id, &anchor)
575            .await
576            .map_err(EmbedError::Runtime)
577    }
578
579    pub async fn cancel_queued_work_batch(
580        &self,
581        batch_id: &str,
582    ) -> Result<Option<QueuedWorkBatch>> {
583        let session_id = self.session_id();
584        self.runtime
585            .cancel_queued_work_batch(&session_id, batch_id)
586            .await
587            .map_err(EmbedError::Runtime)
588    }
589
590    /// Resolve once `batch_id` is no longer pending in the queue store —
591    /// drained by whoever runs queued work (a queued-work runner, a durable
592    /// worker, or another handle's [`queued_turn`](Self::queued_turn)) or
593    /// cancelled. This is the enqueue-and-observe side of the queue: the
594    /// caller never claims the work itself.
595    ///
596    /// Completion is read from the persistent queue store, so it observes
597    /// drains performed by other session handles and other processes alike.
598    /// There is no built-in deadline — nothing resolves if nothing drains the
599    /// queue, so bound it with `tokio::time::timeout` when the worker may be
600    /// unavailable. A batch id the store has never seen resolves immediately.
601    pub async fn await_queued_work_batch(&self, batch_id: &str) -> Result<()> {
602        let observation = self.runtime.observe();
603        let store = observation.queue_store.clone().ok_or_else(|| {
604            EmbedError::Runtime(lash_core::RuntimeError::new(
605                lash_core::RuntimeErrorCode::StoreCommitFailed,
606                "queued work inspection requires a persistent runtime store",
607            ))
608        })?;
609        let session_id = observation.session_id().to_string();
610        drop(observation);
611        let mut delay = std::time::Duration::from_millis(25);
612        loop {
613            let pending = store
614                .list_pending_queued_work(&session_id)
615                .await
616                .map_err(|err| {
617                    EmbedError::Runtime(lash_core::RuntimeError::new(
618                        lash_core::RuntimeErrorCode::StoreCommitFailed,
619                        err.to_string(),
620                    ))
621                })?;
622            if !pending.iter().any(|batch| batch.batch_id == batch_id) {
623                return Ok(());
624            }
625            tokio::time::sleep(delay).await;
626            delay = (delay * 2).min(std::time::Duration::from_millis(400));
627        }
628    }
629
630    pub fn read_view(&self) -> SessionReadView {
631        self.runtime.observe().read_view.clone()
632    }
633
634    pub fn usage_report(&self) -> SessionUsageReport {
635        self.runtime.observe().usage_report.clone()
636    }
637
638    pub async fn set_turn_phase_probe(
639        &self,
640        probe: Arc<dyn lash_core::runtime::RuntimeTurnPhaseProbe>,
641    ) {
642        let writer = self.runtime.writer();
643        let mut runtime = writer.lock().await;
644        runtime.set_turn_phase_probe(Arc::clone(&probe));
645        self.runtime.publish_from(&runtime);
646        if let Some(slot) = &self.process_phase_probe_slot {
647            let observation = self.runtime.observe();
648            slot.set_for_session(observation.session_id(), Arc::clone(&probe));
649            let current_frame = observation.persisted_state.current_agent_frame_id.as_str();
650            if !current_frame.is_empty() {
651                let scope = lash_core::SessionScope::for_agent_frame(
652                    observation.session_id(),
653                    current_frame,
654                );
655                slot.set_for_scope(&scope, probe);
656            }
657        }
658    }
659}
660
661#[derive(Clone)]
662pub struct ObservableSession {
663    pub(crate) runtime: RuntimeHandle,
664}
665
666impl ObservableSession {
667    fn snapshot(&self) -> Arc<RuntimeObservation> {
668        self.runtime.observe()
669    }
670
671    pub fn current_observation(&self) -> SessionObservation {
672        self.runtime.current_session_observation()
673    }
674
675    pub fn current_remote_observation(&self) -> RemoteSessionObservation {
676        RemoteSessionObservation::from_core(self.current_observation())
677    }
678
679    pub fn resume_from_cursor(&self, cursor: &SessionCursor) -> Result<SessionResume> {
680        self.runtime
681            .resume_session_observation(cursor)
682            .map_err(live_replay_error)
683    }
684
685    pub fn subscribe_from_cursor(
686        &self,
687        cursor: &SessionCursor,
688    ) -> Result<SessionObservationSubscription> {
689        self.runtime
690            .subscribe_session_observation(cursor)
691            .map_err(live_replay_error)
692    }
693
694    pub fn subscribe_from_remote_cursor(
695        &self,
696        cursor: &RemoteSessionCursor,
697    ) -> Result<RemoteSessionObservationSubscription> {
698        cursor.validate()?;
699        let cursor = lash_core::SessionCursor::try_from(cursor.clone())?;
700        match self.subscribe_from_cursor(&cursor)? {
701            SessionObservationSubscription::Subscribed(subscription) => {
702                Ok(RemoteSessionObservationSubscription::Subscribed(
703                    RemoteSessionObservationEventStream::new(subscription),
704                ))
705            }
706            SessionObservationSubscription::Gap { observation, gap } => {
707                Ok(RemoteSessionObservationSubscription::Gap {
708                    observation: observation.into(),
709                    gap: gap.into(),
710                })
711            }
712        }
713    }
714
715    /// Subscribe to session observation events and keep the subscription alive
716    /// across recoverable live-replay gaps.
717    ///
718    /// The returned stream yields [`SessionObservationStreamItem::Gap`] when
719    /// the cursor missed the bounded replay window. Callers should replace
720    /// their UI/projection from the included fresh observation, persist
721    /// `gap.latest_cursor`, and keep polling the same stream; it resubscribes
722    /// from that cursor internally.
723    pub fn subscribe_and_recover(&self, cursor: SessionCursor) -> SessionObservationStream {
724        SessionObservationStream {
725            observable: self.clone(),
726            cursor,
727            subscription: None,
728            done: false,
729        }
730    }
731
732    /// Subscribe to remote DTO session observation events and keep the
733    /// subscription alive across recoverable live-replay gaps.
734    pub fn subscribe_and_recover_remote(
735        &self,
736        cursor: RemoteSessionCursor,
737    ) -> Result<RemoteSessionObservationStream> {
738        cursor.validate()?;
739        let cursor = lash_core::SessionCursor::try_from(cursor)?;
740        Ok(RemoteSessionObservationStream {
741            inner: self.subscribe_and_recover(cursor),
742            next_sequence: 0,
743        })
744    }
745
746    pub fn session_id(&self) -> String {
747        self.snapshot().session_id().to_string()
748    }
749
750    pub fn policy_snapshot(&self) -> SessionPolicy {
751        self.snapshot().policy.clone()
752    }
753
754    pub fn read_view(&self) -> SessionReadView {
755        self.snapshot().read_view.clone()
756    }
757
758    pub fn usage_report(&self) -> SessionUsageReport {
759        self.snapshot().usage_report.clone()
760    }
761
762    pub fn tool_state(&self) -> Option<ToolState> {
763        self.snapshot().tool_state.clone()
764    }
765
766    pub fn active_tool_manifests(&self) -> Vec<ToolManifest> {
767        self.snapshot()
768            .tool_state
769            .as_ref()
770            .map(ToolState::tool_manifests)
771            .unwrap_or_default()
772    }
773
774    pub async fn list_process_handles(&self) -> Vec<ProcessHandleSummary> {
775        self.snapshot().list_process_handles().await
776    }
777
778    pub async fn list_all_process_handles(&self) -> Vec<ProcessHandleSummary> {
779        self.snapshot().list_all_process_handles().await
780    }
781
782    pub fn process_scope(&self) -> SessionScope {
783        self.snapshot().process_scope()
784    }
785}
786
787// A public streaming yield produced one item at a time by `Stream::poll_next`;
788// the variant-size spread is transient (never accumulated in a collection), so
789// boxing would only add a per-event heap allocation on the observation hot path
790// and force `*`-deref churn on every SDK consumer. The sibling
791// `RemoteSessionObservationStreamItem` keeps the same inline shape.
792#[allow(clippy::large_enum_variant)]
793#[derive(Clone, Debug)]
794pub enum SessionObservationStreamItem {
795    /// A replayed or live session observation event.
796    Event(SessionObservationEvent),
797    /// A recoverable replay gap with a fresh durable observation.
798    Gap {
799        observation: SessionObservation,
800        gap: LiveReplayGap,
801    },
802}
803
804pub enum RemoteSessionObservationSubscription {
805    Subscribed(RemoteSessionObservationEventStream),
806    Gap {
807        observation: RemoteSessionObservation,
808        gap: RemoteLiveReplayGap,
809    },
810}
811
812#[derive(Clone, Debug)]
813pub enum RemoteSessionObservationStreamItem {
814    /// A replayed or live session observation event encoded as remote DTOs.
815    Event(RemoteSessionObservationEvent),
816    /// A recoverable replay gap with a fresh remote observation snapshot.
817    Gap {
818        observation: RemoteSessionObservation,
819        gap: RemoteLiveReplayGap,
820    },
821}
822
823pub struct RemoteSessionObservationEventStream {
824    inner: lash_core::LiveReplaySubscription,
825    next_sequence: u64,
826}
827
828impl RemoteSessionObservationEventStream {
829    fn new(inner: lash_core::LiveReplaySubscription) -> Self {
830        Self {
831            inner,
832            next_sequence: 0,
833        }
834    }
835
836    pub async fn next_event(&mut self) -> Result<RemoteSessionObservationEvent> {
837        futures_util::future::poll_fn(|cx| Pin::new(&mut *self).poll_next(cx))
838            .await
839            .transpose()?
840            .ok_or_else(|| live_replay_error(LiveReplayStoreError::Closed))
841    }
842}
843
844impl Stream for RemoteSessionObservationEventStream {
845    type Item = Result<RemoteSessionObservationEvent>;
846
847    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
848        match Pin::new(&mut self.inner).poll_next(cx) {
849            Poll::Pending => Poll::Pending,
850            Poll::Ready(Some(Ok(event))) => {
851                let remote = RemoteSessionObservationEvent::from_core(self.next_sequence, event);
852                self.next_sequence = self.next_sequence.saturating_add(1);
853                Poll::Ready(Some(Ok(remote)))
854            }
855            Poll::Ready(Some(Err(err))) => Poll::Ready(Some(Err(live_replay_error(err)))),
856            Poll::Ready(None) => Poll::Ready(None),
857        }
858    }
859}
860
861/// Remote DTO stream returned by [`ObservableSession::subscribe_and_recover_remote`].
862pub struct RemoteSessionObservationStream {
863    inner: SessionObservationStream,
864    next_sequence: u64,
865}
866
867impl RemoteSessionObservationStream {
868    pub fn cursor(&self) -> RemoteSessionCursor {
869        RemoteSessionCursor::from(self.inner.cursor())
870    }
871}
872
873impl Stream for RemoteSessionObservationStream {
874    type Item = Result<RemoteSessionObservationStreamItem>;
875
876    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
877        match Pin::new(&mut self.inner).poll_next(cx) {
878            Poll::Pending => Poll::Pending,
879            Poll::Ready(Some(Ok(SessionObservationStreamItem::Event(event)))) => {
880                let remote = RemoteSessionObservationEvent::from_core(self.next_sequence, event);
881                self.next_sequence = self.next_sequence.saturating_add(1);
882                Poll::Ready(Some(Ok(RemoteSessionObservationStreamItem::Event(remote))))
883            }
884            Poll::Ready(Some(Ok(SessionObservationStreamItem::Gap { observation, gap }))) => {
885                Poll::Ready(Some(Ok(RemoteSessionObservationStreamItem::Gap {
886                    observation: observation.into(),
887                    gap: gap.into(),
888                })))
889            }
890            Poll::Ready(Some(Err(err))) => Poll::Ready(Some(Err(err))),
891            Poll::Ready(None) => Poll::Ready(None),
892        }
893    }
894}
895
896/// Stream returned by [`ObservableSession::subscribe_and_recover`].
897pub struct SessionObservationStream {
898    observable: ObservableSession,
899    cursor: SessionCursor,
900    subscription: Option<lash_core::LiveReplaySubscription>,
901    done: bool,
902}
903
904impl SessionObservationStream {
905    pub fn cursor(&self) -> &SessionCursor {
906        &self.cursor
907    }
908}
909
910impl Stream for SessionObservationStream {
911    type Item = Result<SessionObservationStreamItem>;
912
913    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
914        loop {
915            if self.done {
916                return Poll::Ready(None);
917            }
918            if self.subscription.is_none() {
919                match self.observable.subscribe_from_cursor(&self.cursor) {
920                    Ok(SessionObservationSubscription::Subscribed(subscription)) => {
921                        self.subscription = Some(subscription);
922                    }
923                    Ok(SessionObservationSubscription::Gap { observation, gap }) => {
924                        self.cursor = gap.latest_cursor.clone();
925                        return Poll::Ready(Some(Ok(SessionObservationStreamItem::Gap {
926                            observation,
927                            gap,
928                        })));
929                    }
930                    Err(err) => {
931                        self.done = true;
932                        return Poll::Ready(Some(Err(err)));
933                    }
934                }
935            }
936
937            let Some(subscription) = self.subscription.as_mut() else {
938                continue;
939            };
940            match Pin::new(subscription).poll_next(cx) {
941                Poll::Pending => return Poll::Pending,
942                Poll::Ready(Some(Ok(event))) => {
943                    self.cursor = event.cursor.clone();
944                    return Poll::Ready(Some(Ok(SessionObservationStreamItem::Event(event))));
945                }
946                Poll::Ready(Some(Err(LiveReplayStoreError::SubscriberLagged(_)))) => {
947                    self.subscription = None;
948                    continue;
949                }
950                Poll::Ready(Some(Err(err))) => {
951                    self.done = true;
952                    return Poll::Ready(Some(Err(live_replay_error(err))));
953                }
954                Poll::Ready(None) => {
955                    self.done = true;
956                    return Poll::Ready(None);
957                }
958            }
959        }
960    }
961}
962
963fn live_replay_error(err: lash_core::LiveReplayStoreError) -> EmbedError {
964    EmbedError::Runtime(lash_core::RuntimeError::new(
965        RuntimeErrorCode::Other("live_replay".to_string()),
966        err.to_string(),
967    ))
968}
969
970pub struct EnqueueTurnBuilder<'a> {
971    session: &'a LashSession,
972    input: TurnInput,
973    id: Option<String>,
974    ingress: TurnInputIngress,
975}
976
977impl<'a> EnqueueTurnBuilder<'a> {
978    pub fn id(mut self, id: impl Into<String>) -> Self {
979        self.id = Some(id.into());
980        self
981    }
982
983    pub fn ingress(mut self, ingress: TurnInputIngress) -> Self {
984        self.ingress = ingress;
985        self
986    }
987
988    pub async fn send(self) -> Result<PendingTurnInput> {
989        let source_key = self.id.map(|id| format!("host:{id}"));
990        self.session
991            .runtime
992            .enqueue_turn_input(self.input, self.ingress, source_key)
993            .await
994            .map_err(EmbedError::Runtime)
995    }
996}
997
998impl<'a> std::future::IntoFuture for EnqueueTurnBuilder<'a> {
999    type Output = Result<PendingTurnInput>;
1000    type IntoFuture =
1001        std::pin::Pin<Box<dyn std::future::Future<Output = Result<PendingTurnInput>> + 'a>>;
1002
1003    fn into_future(self) -> Self::IntoFuture {
1004        Box::pin(self.send())
1005    }
1006}