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 QueuedWorkClaim, TurnInputAcceptanceReceipt, TurnInputClaim, 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 Box::pin(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 Box::pin(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 reconcile_loaded_state_policy(&mut state, &policy);
146 Box::pin(self.open_resolved(policy, state, store)).await
147 }
148
149 fn session_policy(&self) -> SessionPolicy {
150 let mut policy = self.spec.resolve_against(&self.core.policy);
151 policy.session_id = Some(self.session_id.clone());
152 policy
153 }
154
155 async fn load_or_default_state(
156 &self,
157 policy: &SessionPolicy,
158 store: Option<&dyn RuntimePersistence>,
159 ) -> Result<RuntimeSessionState> {
160 let state = match store {
161 Some(store) => {
162 let loaded = self.load_persisted_state_for_residency(store).await?;
163 let mut state = loaded.unwrap_or_else(|| RuntimeSessionState {
164 session_id: self.session_id.clone(),
165 policy: policy.clone(),
166 ..RuntimeSessionState::default()
167 });
168 if state.session_id != self.session_id {
169 return Err(EmbedError::StoreSessionMismatch {
170 loaded: state.session_id,
171 requested: self.session_id.clone(),
172 });
173 }
174 reconcile_loaded_state_policy(&mut state, policy);
175 state
176 }
177 None => RuntimeSessionState {
178 session_id: self.session_id.clone(),
179 policy: policy.clone(),
180 ..RuntimeSessionState::default()
181 },
182 };
183 Ok(state)
184 }
185
186 async fn load_persisted_state_for_residency(
187 &self,
188 store: &dyn RuntimePersistence,
189 ) -> Result<Option<RuntimeSessionState>> {
190 load_persisted_state_for_residency(self.core.env.residency, store).await
191 }
192
193 async fn open_resolved(
194 self,
195 policy: SessionPolicy,
196 state: RuntimeSessionState,
197 store: Option<Arc<dyn RuntimePersistence>>,
198 ) -> Result<LashSession> {
199 let mut env = self.core.env.clone();
200 if let Some(provider) = self.provider.clone().or_else(|| self.core.provider.clone()) {
201 env.core.providers.provider_resolver =
202 Arc::new(lash_core::SingleProviderResolver::new(provider));
203 }
204 let plugin_host = build_plugin_host(
205 self.core.protocol_factory.as_ref(),
206 self.core.plugin_factories.as_ref(),
207 self.plugin_factories,
208 )?;
209 env.core = plugin_host.install_process_engine_contributions(
210 env.core.clone(),
211 self.core.process_lifecycle_available,
212 )?;
213 env.plugin_host = Some(Arc::new(plugin_host));
214 let effect_host = Arc::clone(&env.core.control.effect_host);
215 let drivers = self.core.work_driver.drivers().await;
216 env.process_work_driver = drivers.process.clone();
217 env.queued_work_driver = drivers.queued.clone();
218 let mut runtime = LashRuntime::from_environment(&env, policy, state, store).await?;
219 runtime.configure_protocol_on_materialize(
223 &self.plugin_options,
224 self.parent_session_id.is_none(),
225 )?;
226 if let Some(owner) = self.session_execution_owner {
227 runtime.set_runtime_lease_owner(owner);
228 }
229 if drivers.drive_process_on_open
230 && let Some(driver) = drivers.process.as_ref()
231 {
232 driver.claim_and_run_pending("session_open").await?;
233 }
234 let handle = RuntimeHandle::with_live_replay_store(
235 runtime,
236 Arc::clone(&self.core.live_replay_store),
237 );
238 Ok(LashSession {
239 runtime: handle,
240 effect_host,
241 parent_session_id: self.parent_session_id,
242 active_plugins: self.active_plugins,
243 process_phase_probe_slot: self.core.work_driver.phase_probe_slot(),
244 turn_cancels: crate::turn::TurnCancelRegistry::default(),
245 })
246 }
247
248 async fn create_store(
249 &self,
250 policy: &SessionPolicy,
251 ) -> Result<Option<Arc<dyn RuntimePersistence>>> {
252 if let Some(store) = self.store.as_ref() {
253 return Ok(Some(Arc::clone(store)));
254 }
255 let Some(factory) = self.core.store_factory.as_ref() else {
256 return Ok(None);
257 };
258 let request = SessionStoreCreateRequest {
259 session_id: self.session_id.clone(),
260 relation: self
261 .parent_session_id
262 .as_ref()
263 .map(|parent_session_id| lash_core::SessionRelation::Child {
264 parent_session_id: parent_session_id.clone(),
265 caused_by: None,
266 })
267 .unwrap_or_default(),
268 policy: policy.clone(),
269 };
270 factory
271 .create_store(&request)
272 .await
273 .map(Some)
274 .map_err(|message| EmbedError::StoreFactory {
275 session_id: self.session_id.clone(),
276 message,
277 })
278 }
279}
280
281pub(crate) async fn load_state_for_residency(
282 residency: Residency,
283 session_id: &str,
284 policy: &SessionPolicy,
285 store: &dyn RuntimePersistence,
286) -> Result<RuntimeSessionState> {
287 let mut state = load_persisted_state_for_residency(residency, store)
288 .await?
289 .unwrap_or_else(|| RuntimeSessionState {
290 session_id: session_id.to_string(),
291 policy: policy.clone(),
292 ..RuntimeSessionState::default()
293 });
294 if state.session_id != session_id {
295 return Err(EmbedError::StoreSessionMismatch {
296 loaded: state.session_id,
297 requested: session_id.to_string(),
298 });
299 }
300 reconcile_loaded_state_policy(&mut state, policy);
301 Ok(state)
302}
303
304fn reconcile_loaded_state_policy(state: &mut RuntimeSessionState, policy: &SessionPolicy) {
305 let recorded_provider_id = state.policy.recorded_provider_id().to_string();
306 state.policy = policy.clone();
307 state.policy.provider_id = recorded_provider_id;
308 let reconciled_policy = state.policy.clone();
309 if let Some(frame) = state.current_agent_frame_mut() {
310 frame.assignment.policy = reconciled_policy;
311 }
312}
313
314async fn load_persisted_state_for_residency(
315 residency: Residency,
316 store: &dyn RuntimePersistence,
317) -> Result<Option<RuntimeSessionState>> {
318 match residency {
319 Residency::KeepAll => {
320 let loaded = lash_core::store::load_persisted_session_state(store)
321 .await
322 .map_err(|err| SessionError::Protocol(format!("failed to load store: {err}")))?;
323 Ok(loaded)
324 }
325 Residency::ActivePathOnly => {
326 let active = lash_core::store::load_persisted_session_state_active_path(store, None)
327 .await
328 .map_err(|err| {
329 SessionError::Protocol(format!("failed to load active-path store: {err}"))
330 })?;
331 if active
332 .as_ref()
333 .is_some_and(|state| state.session_graph.nodes.is_empty())
334 {
335 let mut full = lash_core::store::load_persisted_session_state(store)
336 .await
337 .map_err(|err| {
338 SessionError::Protocol(format!(
339 "failed to heal active-path store from full graph: {err}"
340 ))
341 })?;
342 if let Some(state) = full.as_mut() {
343 state.graph_replace_required = true;
344 }
345 return Ok(full);
346 }
347 Ok(active)
348 }
349 }
350}
351
352impl PromptLayerSink for SessionBuilder {
353 fn prompt_layer_mut(&mut self) -> &mut PromptLayer {
354 self.spec.prompt.get_or_insert_with(PromptLayer::new)
355 }
356}
357
358#[derive(Clone)]
359pub struct LashSession {
360 pub(crate) runtime: RuntimeHandle,
361 pub(crate) effect_host: Arc<dyn EffectHost>,
362 pub(crate) parent_session_id: Option<String>,
363 pub(crate) active_plugins: Vec<ActivePluginBinding>,
364 pub(crate) process_phase_probe_slot: Option<lash_core::runtime::RuntimeTurnPhaseProbeSlot>,
365 pub(crate) turn_cancels: crate::turn::TurnCancelRegistry,
366}
367
368#[derive(Clone, Debug, Default)]
369pub struct SessionConfigPatch {
370 pub provider: Option<ProviderHandle>,
371 pub model: Option<ModelSpec>,
372 pub prompt: Option<PromptLayer>,
373}
374
375pub struct ParkedSession {
388 pub(crate) inner: lash_core::ParkedSession,
389}
390
391impl ParkedSession {
392 pub fn session_id(&self) -> &str {
395 self.inner.session_id()
396 }
397}
398
399impl LashSession {
400 pub async fn close(self) -> Result<()> {
418 let persistent = self.runtime.observe().queue_store.is_some();
422 let runtime = self.into_owned_runtime()?;
423 runtime.unregister_plugin_session()?;
424 if persistent {
425 runtime.park().await?;
428 }
429 Ok(())
431 }
432
433 pub async fn park(self) -> Result<ParkedSession> {
458 let runtime = self.into_owned_runtime()?;
459 runtime.unregister_plugin_session()?;
462 let parked = runtime.park().await?;
463 Ok(ParkedSession { inner: parked })
464 }
465
466 fn into_owned_runtime(self) -> Result<LashRuntime> {
472 let LashSession { runtime, .. } = self;
473 let writer = runtime.writer();
478 drop(runtime);
479 Arc::try_unwrap(writer)
480 .map(|mutex| mutex.into_inner())
481 .map_err(|_| EmbedError::SessionStillInUse)
482 }
483
484 pub fn session_id(&self) -> String {
485 self.runtime.observe().session_id().to_string()
486 }
487
488 pub fn policy_snapshot(&self) -> SessionPolicy {
489 self.runtime.observe().policy.clone()
490 }
491
492 pub fn observe(&self) -> ObservableSession {
493 ObservableSession {
494 runtime: self.runtime.clone(),
495 }
496 }
497
498 pub fn parent_session_id(&self) -> Option<&str> {
499 self.parent_session_id.as_deref()
500 }
501
502 pub fn effect_host(&self) -> Arc<dyn EffectHost> {
503 Arc::clone(&self.effect_host)
504 }
505
506 pub fn turn(&self, input: TurnInput) -> TurnBuilder {
507 TurnBuilder {
508 runtime: self.runtime.clone(),
509 effect_host: Arc::clone(&self.effect_host),
510 active_plugins: self.active_plugins.clone(),
511 input,
512 cancel: CancellationToken::new(),
513 cancel_origin_hint: lash_core::TurnCancelOriginHint::default(),
514 cancels: self.turn_cancels.clone(),
515 protocol_turn_options: None,
516 provider: None,
517 turn_id: None,
518 }
519 }
520
521 pub fn queued_turn(&self) -> QueuedTurnBuilder {
522 QueuedTurnBuilder {
523 runtime: self.runtime.clone(),
524 effect_host: Arc::clone(&self.effect_host),
525 cancel: CancellationToken::new(),
526 cancel_origin_hint: lash_core::TurnCancelOriginHint::default(),
527 cancels: self.turn_cancels.clone(),
528 batch_ids: Vec::new(),
529 drain_id: None,
530 }
531 }
532
533 pub async fn request_turn_cancel(
544 &self,
545 turn_id: &str,
546 request_id: impl Into<String>,
547 origin: Option<String>,
548 reason: Option<String>,
549 ) -> Result<lash_core::TurnCancelReceipt> {
550 let mut request = lash_core::TurnCancelRequest::new(
551 lash_core::TurnAddress::new(self.session_id(), turn_id),
552 request_id,
553 origin,
554 );
555 request.reason = reason;
556 lash_core::TurnWorkDriver::new(self.effect_host())
557 .request_cancel(request)
558 .await
559 .map_err(EmbedError::Runtime)
560 }
561
562 pub fn cancel_running_turns(&self) -> usize {
579 self.cancel_running_turns_with_origin(None)
580 }
581
582 pub fn cancel_running_turns_with_origin(&self, origin: Option<String>) -> usize {
584 self.turn_cancels.cancel_all(origin)
585 }
586
587 pub fn admin(&self) -> SessionAdmin {
588 SessionAdmin {
589 runtime: self.runtime.clone(),
590 }
591 }
592
593 pub async fn configure(&self, patch: SessionConfigPatch) -> Result<()> {
594 self.admin().config().update(patch).await
595 }
596
597 pub fn tools(&self) -> ToolAdmin {
598 ToolAdmin::new(self.admin())
599 }
600
601 pub fn commands(&self) -> SessionCommandAdmin {
602 self.admin().commands()
603 }
604
605 pub fn triggers(&self) -> SessionTriggerAdmin {
606 self.admin().triggers()
607 }
608
609 pub fn processes(&self) -> SessionProcessAdmin {
610 SessionProcessAdmin::new(self.admin())
611 }
612
613 pub async fn refresh_background_graph(&self) -> Result<()> {
620 self.admin().refresh_background_graph().await
621 }
622
623 pub fn plugin_operations(&self) -> PluginOperations {
624 PluginOperations {
625 control: self.admin(),
626 }
627 }
628
629 pub fn enqueue(&self, input: TurnInput) -> EnqueueTurnBuilder<'_> {
630 EnqueueTurnBuilder {
631 session: self,
632 input,
633 id: None,
634 ingress: TurnInputIngress::NextTurn,
635 }
636 }
637
638 pub async fn queued_work(&self) -> Result<Vec<QueuedWorkBatch>> {
645 let observation = self.runtime.observe();
646 let store = observation.queue_store.as_ref().ok_or_else(|| {
647 EmbedError::Runtime(lash_core::RuntimeError::new(
648 lash_core::RuntimeErrorCode::StoreCommitFailed,
649 "queued work inspection requires a persistent runtime store",
650 ))
651 })?;
652 store
653 .list_pending_queued_work(observation.session_id())
654 .await
655 .map_err(|err| {
656 EmbedError::Runtime(lash_core::RuntimeError::new(
657 lash_core::RuntimeErrorCode::StoreCommitFailed,
658 err.to_string(),
659 ))
660 })
661 }
662
663 pub async fn pending_turn_inputs(&self) -> Result<Vec<PendingTurnInput>> {
664 let observation = self.runtime.observe();
665 let store = observation.queue_store.as_ref().ok_or_else(|| {
666 EmbedError::Runtime(lash_core::RuntimeError::new(
667 lash_core::RuntimeErrorCode::StoreCommitFailed,
668 "pending turn input inspection requires a persistent runtime store",
669 ))
670 })?;
671 store
672 .list_pending_turn_inputs(observation.session_id())
673 .await
674 .map_err(|err| {
675 EmbedError::Runtime(lash_core::RuntimeError::new(
676 lash_core::RuntimeErrorCode::StoreCommitFailed,
677 err.to_string(),
678 ))
679 })
680 }
681
682 pub async fn turn_input_applications(&self) -> Result<Vec<lash_core::TurnInputApplication>> {
687 let observation = self.runtime.observe();
688 let store = observation.queue_store.as_ref().ok_or_else(|| {
689 EmbedError::Runtime(lash_core::RuntimeError::new(
690 lash_core::RuntimeErrorCode::StoreCommitFailed,
691 "turn input application reconciliation requires a persistent runtime store",
692 ))
693 })?;
694 store
695 .list_turn_input_applications(observation.session_id())
696 .await
697 .map_err(|err| {
698 EmbedError::Runtime(lash_core::RuntimeError::new(
699 lash_core::RuntimeErrorCode::StoreCommitFailed,
700 err.to_string(),
701 ))
702 })
703 }
704
705 pub async fn remote_turn_input_applications(
707 &self,
708 ) -> Result<Vec<lash_remote_protocol::RemoteTurnInputApplication>> {
709 Ok(self
710 .turn_input_applications()
711 .await?
712 .iter()
713 .map(Into::into)
714 .collect())
715 }
716
717 pub async fn cancel_pending_turn_input(
718 &self,
719 input_id: &str,
720 ) -> Result<PendingTurnInputCancelOutcome> {
721 let session_id = self.session_id();
722 self.runtime
723 .cancel_pending_turn_input(&session_id, input_id)
724 .await
725 .map_err(EmbedError::Runtime)
726 }
727
728 pub async fn cancel_pending_turn_inputs(
736 &self,
737 targets: impl IntoIterator<Item = PendingTurnInputCancelTarget>,
738 ) -> Result<Vec<PendingTurnInputCancelResult>> {
739 let session_id = self.session_id();
740 let targets = targets.into_iter().collect::<Vec<_>>();
741 self.runtime
742 .cancel_pending_turn_inputs(&session_id, &targets)
743 .await
744 .map_err(EmbedError::Runtime)
745 }
746
747 pub async fn cancel_pending_turn_input_suffix(
756 &self,
757 anchor: PendingTurnInputCancelTarget,
758 ) -> Result<PendingTurnInputSuffixCancelOutcome> {
759 let session_id = self.session_id();
760 self.runtime
761 .cancel_pending_turn_input_suffix(&session_id, &anchor)
762 .await
763 .map_err(EmbedError::Runtime)
764 }
765
766 pub async fn cancel_queued_work_batch(
767 &self,
768 batch_id: &str,
769 ) -> Result<Option<QueuedWorkBatch>> {
770 let session_id = self.session_id();
771 self.runtime
772 .cancel_queued_work_batch(&session_id, batch_id)
773 .await
774 .map_err(EmbedError::Runtime)
775 }
776
777 pub async fn abandon_queued_work_claim(&self, claim: &QueuedWorkClaim) -> Result<()> {
784 self.runtime
785 .abandon_queued_work_claim(claim)
786 .await
787 .map_err(EmbedError::Runtime)
788 }
789
790 pub async fn abandon_turn_input_claim(&self, claim: &TurnInputClaim) -> Result<()> {
794 self.runtime
795 .abandon_turn_input_claim(claim)
796 .await
797 .map_err(EmbedError::Runtime)
798 }
799
800 pub async fn revoke_durable_waits(&self) -> Result<()> {
810 let session_id = self.session_id();
811 self.effect_host
812 .cancel_await_events_for_session(&session_id)
813 .await
814 .map_err(EmbedError::Runtime)
815 }
816
817 pub async fn await_queued_work_batch(&self, batch_id: &str) -> Result<()> {
829 let observation = self.runtime.observe();
830 let store = observation.queue_store.clone().ok_or_else(|| {
831 EmbedError::Runtime(lash_core::RuntimeError::new(
832 lash_core::RuntimeErrorCode::StoreCommitFailed,
833 "queued work inspection requires a persistent runtime store",
834 ))
835 })?;
836 let session_id = observation.session_id().to_string();
837 drop(observation);
838 let mut delay = std::time::Duration::from_millis(25);
839 loop {
840 let pending = store
841 .list_pending_queued_work(&session_id)
842 .await
843 .map_err(|err| {
844 EmbedError::Runtime(lash_core::RuntimeError::new(
845 lash_core::RuntimeErrorCode::StoreCommitFailed,
846 err.to_string(),
847 ))
848 })?;
849 if !pending.iter().any(|batch| batch.batch_id == batch_id) {
850 return Ok(());
851 }
852 tokio::time::sleep(delay).await;
853 delay = (delay * 2).min(std::time::Duration::from_millis(400));
854 }
855 }
856
857 pub fn read_view(&self) -> SessionReadView {
858 self.runtime.observe().read_view.clone()
859 }
860
861 pub fn usage_report(&self) -> SessionUsageReport {
862 self.runtime.observe().usage_report.clone()
863 }
864
865 pub async fn set_turn_phase_probe(
866 &self,
867 probe: Arc<dyn lash_core::runtime::RuntimeTurnPhaseProbe>,
868 ) {
869 let writer = self.runtime.writer();
870 let mut runtime = writer.lock().await;
871 runtime.set_turn_phase_probe(Arc::clone(&probe));
872 self.runtime.publish_from(&runtime);
873 if let Some(slot) = &self.process_phase_probe_slot {
874 let observation = self.runtime.observe();
875 slot.set_for_session(observation.session_id(), Arc::clone(&probe));
876 let current_frame = observation.persisted_state.current_agent_frame_id.as_str();
877 if !current_frame.is_empty() {
878 let scope = lash_core::SessionScope::for_agent_frame(
879 observation.session_id(),
880 current_frame,
881 );
882 slot.set_for_scope(&scope, probe);
883 }
884 }
885 }
886}
887
888#[derive(Clone)]
889pub struct ObservableSession {
890 pub(crate) runtime: RuntimeHandle,
891}
892
893impl ObservableSession {
894 fn snapshot(&self) -> Arc<RuntimeObservation> {
895 self.runtime.observe()
896 }
897
898 pub fn current_observation(&self) -> SessionObservation {
899 self.runtime.current_session_observation()
900 }
901
902 pub fn current_remote_observation(&self) -> RemoteSessionObservation {
903 RemoteSessionObservation::from_core(self.current_observation())
904 }
905
906 pub fn resume_from_cursor(&self, cursor: &SessionCursor) -> Result<SessionResume> {
907 self.runtime
908 .resume_session_observation(cursor)
909 .map_err(live_replay_error)
910 }
911
912 pub fn subscribe_from_cursor(
913 &self,
914 cursor: &SessionCursor,
915 ) -> Result<SessionObservationSubscription> {
916 self.runtime
917 .subscribe_session_observation(cursor)
918 .map_err(live_replay_error)
919 }
920
921 pub fn subscribe_from_remote_cursor(
922 &self,
923 cursor: &RemoteSessionCursor,
924 ) -> Result<RemoteSessionObservationSubscription> {
925 cursor.validate()?;
926 let cursor = lash_core::SessionCursor::try_from(cursor.clone())?;
927 match self.subscribe_from_cursor(&cursor)? {
928 SessionObservationSubscription::Subscribed(subscription) => {
929 Ok(RemoteSessionObservationSubscription::Subscribed(
930 RemoteSessionObservationEventStream::new(subscription),
931 ))
932 }
933 SessionObservationSubscription::Gap { observation, gap } => {
934 Ok(RemoteSessionObservationSubscription::Gap {
935 observation: observation.into(),
936 gap: gap.into(),
937 })
938 }
939 }
940 }
941
942 pub fn subscribe_and_recover(&self, cursor: SessionCursor) -> SessionObservationStream {
951 SessionObservationStream {
952 observable: self.clone(),
953 cursor,
954 subscription: None,
955 done: false,
956 }
957 }
958
959 pub fn subscribe_and_recover_remote(
962 &self,
963 cursor: RemoteSessionCursor,
964 ) -> Result<RemoteSessionObservationStream> {
965 cursor.validate()?;
966 let cursor = lash_core::SessionCursor::try_from(cursor)?;
967 Ok(RemoteSessionObservationStream {
968 inner: self.subscribe_and_recover(cursor),
969 next_sequence: 0,
970 })
971 }
972
973 pub fn session_id(&self) -> String {
974 self.snapshot().session_id().to_string()
975 }
976
977 pub fn policy_snapshot(&self) -> SessionPolicy {
978 self.snapshot().policy.clone()
979 }
980
981 pub fn read_view(&self) -> SessionReadView {
982 self.snapshot().read_view.clone()
983 }
984
985 pub fn usage_report(&self) -> SessionUsageReport {
986 self.snapshot().usage_report.clone()
987 }
988
989 pub fn tool_state(&self) -> Option<ToolState> {
990 self.snapshot().tool_state.clone()
991 }
992
993 pub fn active_tool_manifests(&self) -> Vec<ToolManifest> {
994 self.snapshot()
995 .tool_state
996 .as_ref()
997 .map(ToolState::tool_manifests)
998 .unwrap_or_default()
999 }
1000
1001 pub async fn list_process_handles(&self) -> Vec<ProcessHandleSummary> {
1002 self.snapshot().list_process_handles().await
1003 }
1004
1005 pub async fn list_all_process_handles(&self) -> Vec<ProcessHandleSummary> {
1006 self.snapshot().list_all_process_handles().await
1007 }
1008
1009 pub fn process_scope(&self) -> SessionScope {
1010 self.snapshot().process_scope()
1011 }
1012}
1013
1014#[allow(clippy::large_enum_variant)]
1021#[derive(Clone, Debug)]
1022pub enum SessionObservationStreamItem {
1023 Event(Arc<SessionObservationEvent>),
1025 Gap {
1027 observation: SessionObservation,
1028 gap: LiveReplayGap,
1029 },
1030}
1031
1032pub enum RemoteSessionObservationSubscription {
1033 Subscribed(RemoteSessionObservationEventStream),
1034 Gap {
1035 observation: RemoteSessionObservation,
1036 gap: RemoteLiveReplayGap,
1037 },
1038}
1039
1040#[derive(Clone, Debug)]
1041pub enum RemoteSessionObservationStreamItem {
1042 Event(RemoteSessionObservationEvent),
1044 Gap {
1046 observation: RemoteSessionObservation,
1047 gap: RemoteLiveReplayGap,
1048 },
1049}
1050
1051pub struct RemoteSessionObservationEventStream {
1052 inner: lash_core::LiveReplaySubscription,
1053 next_sequence: u64,
1054}
1055
1056impl RemoteSessionObservationEventStream {
1057 fn new(inner: lash_core::LiveReplaySubscription) -> Self {
1058 Self {
1059 inner,
1060 next_sequence: 0,
1061 }
1062 }
1063
1064 pub async fn next_event(&mut self) -> Result<RemoteSessionObservationEvent> {
1065 futures_util::future::poll_fn(|cx| Pin::new(&mut *self).poll_next(cx))
1066 .await
1067 .transpose()?
1068 .ok_or_else(|| live_replay_error(LiveReplayStoreError::Closed))
1069 }
1070}
1071
1072impl Stream for RemoteSessionObservationEventStream {
1073 type Item = Result<RemoteSessionObservationEvent>;
1074
1075 fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
1076 match Pin::new(&mut self.inner).poll_next(cx) {
1077 Poll::Pending => Poll::Pending,
1078 Poll::Ready(Some(Ok(event))) => {
1079 let remote = RemoteSessionObservationEvent::from_core(self.next_sequence, event);
1080 self.next_sequence = self.next_sequence.saturating_add(1);
1081 Poll::Ready(Some(Ok(remote)))
1082 }
1083 Poll::Ready(Some(Err(err))) => Poll::Ready(Some(Err(live_replay_error(err)))),
1084 Poll::Ready(None) => Poll::Ready(None),
1085 }
1086 }
1087}
1088
1089pub struct RemoteSessionObservationStream {
1091 inner: SessionObservationStream,
1092 next_sequence: u64,
1093}
1094
1095impl RemoteSessionObservationStream {
1096 pub fn cursor(&self) -> RemoteSessionCursor {
1097 RemoteSessionCursor::from(self.inner.cursor())
1098 }
1099}
1100
1101impl Stream for RemoteSessionObservationStream {
1102 type Item = Result<RemoteSessionObservationStreamItem>;
1103
1104 fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
1105 match Pin::new(&mut self.inner).poll_next(cx) {
1106 Poll::Pending => Poll::Pending,
1107 Poll::Ready(Some(Ok(SessionObservationStreamItem::Event(event)))) => {
1108 let remote = RemoteSessionObservationEvent::from_core(self.next_sequence, event);
1109 self.next_sequence = self.next_sequence.saturating_add(1);
1110 Poll::Ready(Some(Ok(RemoteSessionObservationStreamItem::Event(remote))))
1111 }
1112 Poll::Ready(Some(Ok(SessionObservationStreamItem::Gap { observation, gap }))) => {
1113 Poll::Ready(Some(Ok(RemoteSessionObservationStreamItem::Gap {
1114 observation: observation.into(),
1115 gap: gap.into(),
1116 })))
1117 }
1118 Poll::Ready(Some(Err(err))) => Poll::Ready(Some(Err(err))),
1119 Poll::Ready(None) => Poll::Ready(None),
1120 }
1121 }
1122}
1123
1124pub struct SessionObservationStream {
1126 observable: ObservableSession,
1127 cursor: SessionCursor,
1128 subscription: Option<lash_core::LiveReplaySubscription>,
1129 done: bool,
1130}
1131
1132impl SessionObservationStream {
1133 pub fn cursor(&self) -> &SessionCursor {
1134 &self.cursor
1135 }
1136}
1137
1138impl Stream for SessionObservationStream {
1139 type Item = Result<SessionObservationStreamItem>;
1140
1141 fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
1142 loop {
1143 if self.done {
1144 return Poll::Ready(None);
1145 }
1146 if self.subscription.is_none() {
1147 match self.observable.subscribe_from_cursor(&self.cursor) {
1148 Ok(SessionObservationSubscription::Subscribed(subscription)) => {
1149 self.subscription = Some(subscription);
1150 }
1151 Ok(SessionObservationSubscription::Gap { observation, gap }) => {
1152 self.cursor = gap.latest_cursor.clone();
1153 return Poll::Ready(Some(Ok(SessionObservationStreamItem::Gap {
1154 observation,
1155 gap,
1156 })));
1157 }
1158 Err(err) => {
1159 self.done = true;
1160 return Poll::Ready(Some(Err(err)));
1161 }
1162 }
1163 }
1164
1165 let Some(subscription) = self.subscription.as_mut() else {
1166 continue;
1167 };
1168 match Pin::new(subscription).poll_next(cx) {
1169 Poll::Pending => return Poll::Pending,
1170 Poll::Ready(Some(Ok(event))) => {
1171 self.cursor = event.cursor.clone();
1172 return Poll::Ready(Some(Ok(SessionObservationStreamItem::Event(event))));
1173 }
1174 Poll::Ready(Some(Err(LiveReplayStoreError::SubscriberLagged(_)))) => {
1175 self.subscription = None;
1176 continue;
1177 }
1178 Poll::Ready(Some(Err(err))) => {
1179 self.done = true;
1180 return Poll::Ready(Some(Err(live_replay_error(err))));
1181 }
1182 Poll::Ready(None) => {
1183 self.done = true;
1184 return Poll::Ready(None);
1185 }
1186 }
1187 }
1188 }
1189}
1190
1191fn live_replay_error(err: lash_core::LiveReplayStoreError) -> EmbedError {
1192 EmbedError::Runtime(lash_core::RuntimeError::new(
1193 RuntimeErrorCode::Other("live_replay".to_string()),
1194 err.to_string(),
1195 ))
1196}
1197
1198pub struct EnqueueTurnBuilder<'a> {
1199 session: &'a LashSession,
1200 input: TurnInput,
1201 id: Option<String>,
1202 ingress: TurnInputIngress,
1203}
1204
1205impl<'a> EnqueueTurnBuilder<'a> {
1206 pub fn id(mut self, id: impl Into<String>) -> Self {
1207 self.id = Some(id.into());
1208 self
1209 }
1210
1211 pub fn ingress(mut self, ingress: TurnInputIngress) -> Self {
1212 self.ingress = ingress;
1213 self
1214 }
1215
1216 pub async fn send(self) -> Result<TurnInputAcceptanceReceipt> {
1223 let source_key = self.id.map(|id| format!("host:{id}"));
1224 self.session
1225 .runtime
1226 .enqueue_turn_input(self.input, self.ingress, source_key)
1227 .await
1228 .map(|pending| TurnInputAcceptanceReceipt::from(&pending))
1229 .map_err(EmbedError::Runtime)
1230 }
1231}
1232
1233impl<'a> std::future::IntoFuture for EnqueueTurnBuilder<'a> {
1234 type Output = Result<TurnInputAcceptanceReceipt>;
1235 type IntoFuture = std::pin::Pin<Box<dyn std::future::Future<Output = Self::Output> + 'a>>;
1236
1237 fn into_future(self) -> Self::IntoFuture {
1238 Box::pin(self.send())
1239 }
1240}