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, 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 cancel_pending_turn_input(
683 &self,
684 input_id: &str,
685 ) -> Result<PendingTurnInputCancelOutcome> {
686 let session_id = self.session_id();
687 self.runtime
688 .cancel_pending_turn_input(&session_id, input_id)
689 .await
690 .map_err(EmbedError::Runtime)
691 }
692
693 pub async fn cancel_pending_turn_inputs(
701 &self,
702 targets: impl IntoIterator<Item = PendingTurnInputCancelTarget>,
703 ) -> Result<Vec<PendingTurnInputCancelResult>> {
704 let session_id = self.session_id();
705 let targets = targets.into_iter().collect::<Vec<_>>();
706 self.runtime
707 .cancel_pending_turn_inputs(&session_id, &targets)
708 .await
709 .map_err(EmbedError::Runtime)
710 }
711
712 pub async fn cancel_pending_turn_input_suffix(
721 &self,
722 anchor: PendingTurnInputCancelTarget,
723 ) -> Result<PendingTurnInputSuffixCancelOutcome> {
724 let session_id = self.session_id();
725 self.runtime
726 .cancel_pending_turn_input_suffix(&session_id, &anchor)
727 .await
728 .map_err(EmbedError::Runtime)
729 }
730
731 pub async fn cancel_queued_work_batch(
732 &self,
733 batch_id: &str,
734 ) -> Result<Option<QueuedWorkBatch>> {
735 let session_id = self.session_id();
736 self.runtime
737 .cancel_queued_work_batch(&session_id, batch_id)
738 .await
739 .map_err(EmbedError::Runtime)
740 }
741
742 pub async fn abandon_queued_work_claim(&self, claim: &QueuedWorkClaim) -> Result<()> {
749 self.runtime
750 .abandon_queued_work_claim(claim)
751 .await
752 .map_err(EmbedError::Runtime)
753 }
754
755 pub async fn abandon_turn_input_claim(&self, claim: &TurnInputClaim) -> Result<()> {
759 self.runtime
760 .abandon_turn_input_claim(claim)
761 .await
762 .map_err(EmbedError::Runtime)
763 }
764
765 pub async fn revoke_durable_waits(&self) -> Result<()> {
775 let session_id = self.session_id();
776 self.effect_host
777 .cancel_await_events_for_session(&session_id)
778 .await
779 .map_err(EmbedError::Runtime)
780 }
781
782 pub async fn await_queued_work_batch(&self, batch_id: &str) -> Result<()> {
794 let observation = self.runtime.observe();
795 let store = observation.queue_store.clone().ok_or_else(|| {
796 EmbedError::Runtime(lash_core::RuntimeError::new(
797 lash_core::RuntimeErrorCode::StoreCommitFailed,
798 "queued work inspection requires a persistent runtime store",
799 ))
800 })?;
801 let session_id = observation.session_id().to_string();
802 drop(observation);
803 let mut delay = std::time::Duration::from_millis(25);
804 loop {
805 let pending = store
806 .list_pending_queued_work(&session_id)
807 .await
808 .map_err(|err| {
809 EmbedError::Runtime(lash_core::RuntimeError::new(
810 lash_core::RuntimeErrorCode::StoreCommitFailed,
811 err.to_string(),
812 ))
813 })?;
814 if !pending.iter().any(|batch| batch.batch_id == batch_id) {
815 return Ok(());
816 }
817 tokio::time::sleep(delay).await;
818 delay = (delay * 2).min(std::time::Duration::from_millis(400));
819 }
820 }
821
822 pub fn read_view(&self) -> SessionReadView {
823 self.runtime.observe().read_view.clone()
824 }
825
826 pub fn usage_report(&self) -> SessionUsageReport {
827 self.runtime.observe().usage_report.clone()
828 }
829
830 pub async fn set_turn_phase_probe(
831 &self,
832 probe: Arc<dyn lash_core::runtime::RuntimeTurnPhaseProbe>,
833 ) {
834 let writer = self.runtime.writer();
835 let mut runtime = writer.lock().await;
836 runtime.set_turn_phase_probe(Arc::clone(&probe));
837 self.runtime.publish_from(&runtime);
838 if let Some(slot) = &self.process_phase_probe_slot {
839 let observation = self.runtime.observe();
840 slot.set_for_session(observation.session_id(), Arc::clone(&probe));
841 let current_frame = observation.persisted_state.current_agent_frame_id.as_str();
842 if !current_frame.is_empty() {
843 let scope = lash_core::SessionScope::for_agent_frame(
844 observation.session_id(),
845 current_frame,
846 );
847 slot.set_for_scope(&scope, probe);
848 }
849 }
850 }
851}
852
853#[derive(Clone)]
854pub struct ObservableSession {
855 pub(crate) runtime: RuntimeHandle,
856}
857
858impl ObservableSession {
859 fn snapshot(&self) -> Arc<RuntimeObservation> {
860 self.runtime.observe()
861 }
862
863 pub fn current_observation(&self) -> SessionObservation {
864 self.runtime.current_session_observation()
865 }
866
867 pub fn current_remote_observation(&self) -> RemoteSessionObservation {
868 RemoteSessionObservation::from_core(self.current_observation())
869 }
870
871 pub fn resume_from_cursor(&self, cursor: &SessionCursor) -> Result<SessionResume> {
872 self.runtime
873 .resume_session_observation(cursor)
874 .map_err(live_replay_error)
875 }
876
877 pub fn subscribe_from_cursor(
878 &self,
879 cursor: &SessionCursor,
880 ) -> Result<SessionObservationSubscription> {
881 self.runtime
882 .subscribe_session_observation(cursor)
883 .map_err(live_replay_error)
884 }
885
886 pub fn subscribe_from_remote_cursor(
887 &self,
888 cursor: &RemoteSessionCursor,
889 ) -> Result<RemoteSessionObservationSubscription> {
890 cursor.validate()?;
891 let cursor = lash_core::SessionCursor::try_from(cursor.clone())?;
892 match self.subscribe_from_cursor(&cursor)? {
893 SessionObservationSubscription::Subscribed(subscription) => {
894 Ok(RemoteSessionObservationSubscription::Subscribed(
895 RemoteSessionObservationEventStream::new(subscription),
896 ))
897 }
898 SessionObservationSubscription::Gap { observation, gap } => {
899 Ok(RemoteSessionObservationSubscription::Gap {
900 observation: observation.into(),
901 gap: gap.into(),
902 })
903 }
904 }
905 }
906
907 pub fn subscribe_and_recover(&self, cursor: SessionCursor) -> SessionObservationStream {
916 SessionObservationStream {
917 observable: self.clone(),
918 cursor,
919 subscription: None,
920 done: false,
921 }
922 }
923
924 pub fn subscribe_and_recover_remote(
927 &self,
928 cursor: RemoteSessionCursor,
929 ) -> Result<RemoteSessionObservationStream> {
930 cursor.validate()?;
931 let cursor = lash_core::SessionCursor::try_from(cursor)?;
932 Ok(RemoteSessionObservationStream {
933 inner: self.subscribe_and_recover(cursor),
934 next_sequence: 0,
935 })
936 }
937
938 pub fn session_id(&self) -> String {
939 self.snapshot().session_id().to_string()
940 }
941
942 pub fn policy_snapshot(&self) -> SessionPolicy {
943 self.snapshot().policy.clone()
944 }
945
946 pub fn read_view(&self) -> SessionReadView {
947 self.snapshot().read_view.clone()
948 }
949
950 pub fn usage_report(&self) -> SessionUsageReport {
951 self.snapshot().usage_report.clone()
952 }
953
954 pub fn tool_state(&self) -> Option<ToolState> {
955 self.snapshot().tool_state.clone()
956 }
957
958 pub fn active_tool_manifests(&self) -> Vec<ToolManifest> {
959 self.snapshot()
960 .tool_state
961 .as_ref()
962 .map(ToolState::tool_manifests)
963 .unwrap_or_default()
964 }
965
966 pub async fn list_process_handles(&self) -> Vec<ProcessHandleSummary> {
967 self.snapshot().list_process_handles().await
968 }
969
970 pub async fn list_all_process_handles(&self) -> Vec<ProcessHandleSummary> {
971 self.snapshot().list_all_process_handles().await
972 }
973
974 pub fn process_scope(&self) -> SessionScope {
975 self.snapshot().process_scope()
976 }
977}
978
979#[allow(clippy::large_enum_variant)]
986#[derive(Clone, Debug)]
987pub enum SessionObservationStreamItem {
988 Event(Arc<SessionObservationEvent>),
990 Gap {
992 observation: SessionObservation,
993 gap: LiveReplayGap,
994 },
995}
996
997pub enum RemoteSessionObservationSubscription {
998 Subscribed(RemoteSessionObservationEventStream),
999 Gap {
1000 observation: RemoteSessionObservation,
1001 gap: RemoteLiveReplayGap,
1002 },
1003}
1004
1005#[derive(Clone, Debug)]
1006pub enum RemoteSessionObservationStreamItem {
1007 Event(RemoteSessionObservationEvent),
1009 Gap {
1011 observation: RemoteSessionObservation,
1012 gap: RemoteLiveReplayGap,
1013 },
1014}
1015
1016pub struct RemoteSessionObservationEventStream {
1017 inner: lash_core::LiveReplaySubscription,
1018 next_sequence: u64,
1019}
1020
1021impl RemoteSessionObservationEventStream {
1022 fn new(inner: lash_core::LiveReplaySubscription) -> Self {
1023 Self {
1024 inner,
1025 next_sequence: 0,
1026 }
1027 }
1028
1029 pub async fn next_event(&mut self) -> Result<RemoteSessionObservationEvent> {
1030 futures_util::future::poll_fn(|cx| Pin::new(&mut *self).poll_next(cx))
1031 .await
1032 .transpose()?
1033 .ok_or_else(|| live_replay_error(LiveReplayStoreError::Closed))
1034 }
1035}
1036
1037impl Stream for RemoteSessionObservationEventStream {
1038 type Item = Result<RemoteSessionObservationEvent>;
1039
1040 fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
1041 match Pin::new(&mut self.inner).poll_next(cx) {
1042 Poll::Pending => Poll::Pending,
1043 Poll::Ready(Some(Ok(event))) => {
1044 let remote = RemoteSessionObservationEvent::from_core(self.next_sequence, event);
1045 self.next_sequence = self.next_sequence.saturating_add(1);
1046 Poll::Ready(Some(Ok(remote)))
1047 }
1048 Poll::Ready(Some(Err(err))) => Poll::Ready(Some(Err(live_replay_error(err)))),
1049 Poll::Ready(None) => Poll::Ready(None),
1050 }
1051 }
1052}
1053
1054pub struct RemoteSessionObservationStream {
1056 inner: SessionObservationStream,
1057 next_sequence: u64,
1058}
1059
1060impl RemoteSessionObservationStream {
1061 pub fn cursor(&self) -> RemoteSessionCursor {
1062 RemoteSessionCursor::from(self.inner.cursor())
1063 }
1064}
1065
1066impl Stream for RemoteSessionObservationStream {
1067 type Item = Result<RemoteSessionObservationStreamItem>;
1068
1069 fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
1070 match Pin::new(&mut self.inner).poll_next(cx) {
1071 Poll::Pending => Poll::Pending,
1072 Poll::Ready(Some(Ok(SessionObservationStreamItem::Event(event)))) => {
1073 let remote = RemoteSessionObservationEvent::from_core(self.next_sequence, event);
1074 self.next_sequence = self.next_sequence.saturating_add(1);
1075 Poll::Ready(Some(Ok(RemoteSessionObservationStreamItem::Event(remote))))
1076 }
1077 Poll::Ready(Some(Ok(SessionObservationStreamItem::Gap { observation, gap }))) => {
1078 Poll::Ready(Some(Ok(RemoteSessionObservationStreamItem::Gap {
1079 observation: observation.into(),
1080 gap: gap.into(),
1081 })))
1082 }
1083 Poll::Ready(Some(Err(err))) => Poll::Ready(Some(Err(err))),
1084 Poll::Ready(None) => Poll::Ready(None),
1085 }
1086 }
1087}
1088
1089pub struct SessionObservationStream {
1091 observable: ObservableSession,
1092 cursor: SessionCursor,
1093 subscription: Option<lash_core::LiveReplaySubscription>,
1094 done: bool,
1095}
1096
1097impl SessionObservationStream {
1098 pub fn cursor(&self) -> &SessionCursor {
1099 &self.cursor
1100 }
1101}
1102
1103impl Stream for SessionObservationStream {
1104 type Item = Result<SessionObservationStreamItem>;
1105
1106 fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
1107 loop {
1108 if self.done {
1109 return Poll::Ready(None);
1110 }
1111 if self.subscription.is_none() {
1112 match self.observable.subscribe_from_cursor(&self.cursor) {
1113 Ok(SessionObservationSubscription::Subscribed(subscription)) => {
1114 self.subscription = Some(subscription);
1115 }
1116 Ok(SessionObservationSubscription::Gap { observation, gap }) => {
1117 self.cursor = gap.latest_cursor.clone();
1118 return Poll::Ready(Some(Ok(SessionObservationStreamItem::Gap {
1119 observation,
1120 gap,
1121 })));
1122 }
1123 Err(err) => {
1124 self.done = true;
1125 return Poll::Ready(Some(Err(err)));
1126 }
1127 }
1128 }
1129
1130 let Some(subscription) = self.subscription.as_mut() else {
1131 continue;
1132 };
1133 match Pin::new(subscription).poll_next(cx) {
1134 Poll::Pending => return Poll::Pending,
1135 Poll::Ready(Some(Ok(event))) => {
1136 self.cursor = event.cursor.clone();
1137 return Poll::Ready(Some(Ok(SessionObservationStreamItem::Event(event))));
1138 }
1139 Poll::Ready(Some(Err(LiveReplayStoreError::SubscriberLagged(_)))) => {
1140 self.subscription = None;
1141 continue;
1142 }
1143 Poll::Ready(Some(Err(err))) => {
1144 self.done = true;
1145 return Poll::Ready(Some(Err(live_replay_error(err))));
1146 }
1147 Poll::Ready(None) => {
1148 self.done = true;
1149 return Poll::Ready(None);
1150 }
1151 }
1152 }
1153 }
1154}
1155
1156fn live_replay_error(err: lash_core::LiveReplayStoreError) -> EmbedError {
1157 EmbedError::Runtime(lash_core::RuntimeError::new(
1158 RuntimeErrorCode::Other("live_replay".to_string()),
1159 err.to_string(),
1160 ))
1161}
1162
1163pub struct EnqueueTurnBuilder<'a> {
1164 session: &'a LashSession,
1165 input: TurnInput,
1166 id: Option<String>,
1167 ingress: TurnInputIngress,
1168}
1169
1170impl<'a> EnqueueTurnBuilder<'a> {
1171 pub fn id(mut self, id: impl Into<String>) -> Self {
1172 self.id = Some(id.into());
1173 self
1174 }
1175
1176 pub fn ingress(mut self, ingress: TurnInputIngress) -> Self {
1177 self.ingress = ingress;
1178 self
1179 }
1180
1181 pub async fn send(self) -> Result<PendingTurnInput> {
1182 let source_key = self.id.map(|id| format!("host:{id}"));
1183 self.session
1184 .runtime
1185 .enqueue_turn_input(self.input, self.ingress, source_key)
1186 .await
1187 .map_err(EmbedError::Runtime)
1188 }
1189}
1190
1191impl<'a> std::future::IntoFuture for EnqueueTurnBuilder<'a> {
1192 type Output = Result<PendingTurnInput>;
1193 type IntoFuture =
1194 std::pin::Pin<Box<dyn std::future::Future<Output = Result<PendingTurnInput>> + 'a>>;
1195
1196 fn into_future(self) -> Self::IntoFuture {
1197 Box::pin(self.send())
1198 }
1199}