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 pub(crate) plugin_options: PluginOptions,
32}
33
34impl SessionBuilder {
35 pub fn plugin_options(mut self, plugin_options: PluginOptions) -> Self {
37 self.plugin_options = plugin_options;
38 self
39 }
40
41 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 pub fn session_execution_owner(mut self, owner: lash_core::LeaseOwnerIdentity) -> Self {
76 self.session_execution_owner = Some(owner);
77 self
78 }
79
80 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 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 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 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 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 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 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 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 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 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 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#[allow(clippy::large_enum_variant)]
793#[derive(Clone, Debug)]
794pub enum SessionObservationStreamItem {
795 Event(SessionObservationEvent),
797 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 Event(RemoteSessionObservationEvent),
816 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
861pub 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
896pub 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}