1use crate::support::*;
2use lash_core::runtime::{DeliveryPolicy, QueuedWorkBatch, SlotPolicy};
3
4pub struct SessionBuilder {
5 pub(crate) core: LashCore,
6 pub(crate) session_id: String,
7 pub(crate) spec: SessionSpec,
8 pub(crate) parent_session_id: Option<String>,
9 pub(crate) store: Option<Arc<dyn RuntimePersistence>>,
10 pub(crate) provider: Option<ProviderHandle>,
11 pub(crate) active_plugins: Vec<ActivePluginBinding>,
12 pub(crate) plugin_factories: Vec<Arc<dyn PluginFactory>>,
13}
14
15#[cfg(feature = "rlm")]
16pub struct RlmSessionBuilder {
17 pub(crate) builder: SessionBuilder,
18 pub(crate) rlm_final_answer_format: Option<lash_rlm_types::RlmFinalAnswerFormat>,
19}
20
21impl SessionBuilder {
22 pub fn provider(mut self, provider: ProviderHandle) -> Self {
23 self.spec = self.spec.provider_id(provider.kind());
24 self.provider = Some(provider);
25 self
26 }
27
28 pub fn session_spec(mut self, spec: SessionSpec) -> Self {
29 self.spec = spec;
30 self
31 }
32
33 pub fn parent(mut self, parent_session_id: impl Into<String>) -> Self {
34 self.parent_session_id = Some(parent_session_id.into());
35 self
36 }
37
38 pub fn store(mut self, store: Arc<dyn RuntimePersistence>) -> Self {
45 self.store = Some(store);
46 self
47 }
48
49 pub fn plugin<P: PluginBinding>(mut self, config: P::SessionConfig) -> Self {
50 self.active_plugins.push(ActivePluginBinding {
51 id: P::ID,
52 requires_turn_input: P::requires_turn_input(&config),
53 });
54 self.plugin_factories.push(P::factory(&config));
55 self
56 }
57
58 pub async fn open(self) -> Result<LashSession> {
59 let policy = self.session_policy();
60 let store = self.create_store(&policy).await?;
61 let state = self
62 .load_or_default_state(&policy, store.as_deref())
63 .await?;
64 self.open_resolved(policy, state, store).await
65 }
66
67 pub async fn open_fresh(self) -> Result<LashSession> {
77 let policy = self.session_policy();
78 let store = self.create_store(&policy).await?;
79 let state = RuntimeSessionState {
80 session_id: self.session_id.clone(),
81 policy: policy.clone(),
82 graph_replace_required: true,
83 ..RuntimeSessionState::default()
84 };
85 self.open_resolved(policy, state, store).await
86 }
87
88 pub async fn open_with_state(self, mut state: RuntimeSessionState) -> Result<LashSession> {
95 let policy = self.session_policy();
96 let store = self.create_store(&policy).await?;
97 if state.session_id != self.session_id {
98 return Err(EmbedError::StoreSessionMismatch {
99 loaded: state.session_id,
100 requested: self.session_id,
101 });
102 }
103 let recorded_provider_id = state.policy.recorded_provider_id().to_string();
104 state.policy = policy.clone();
105 state.policy.provider_id = recorded_provider_id;
106 self.open_resolved(policy, state, store).await
107 }
108
109 fn session_policy(&self) -> SessionPolicy {
110 let mut policy = self.spec.resolve_against(&self.core.policy);
111 policy.session_id = Some(self.session_id.clone());
112 policy
113 }
114
115 async fn load_or_default_state(
116 &self,
117 policy: &SessionPolicy,
118 store: Option<&dyn RuntimePersistence>,
119 ) -> Result<RuntimeSessionState> {
120 let state = match store {
121 Some(store) => {
122 let loaded = self.load_persisted_state_for_residency(store).await?;
123 let mut state = loaded.unwrap_or_else(|| RuntimeSessionState {
124 session_id: self.session_id.clone(),
125 policy: policy.clone(),
126 ..RuntimeSessionState::default()
127 });
128 if state.session_id != self.session_id {
129 return Err(EmbedError::StoreSessionMismatch {
130 loaded: state.session_id,
131 requested: self.session_id.clone(),
132 });
133 }
134 let recorded_provider_id = state.policy.recorded_provider_id().to_string();
135 state.policy = policy.clone();
136 state.policy.provider_id = recorded_provider_id;
137 state
138 }
139 None => RuntimeSessionState {
140 session_id: self.session_id.clone(),
141 policy: policy.clone(),
142 ..RuntimeSessionState::default()
143 },
144 };
145 Ok(state)
146 }
147
148 async fn load_persisted_state_for_residency(
149 &self,
150 store: &dyn RuntimePersistence,
151 ) -> Result<Option<RuntimeSessionState>> {
152 load_persisted_state_for_residency(self.core.env.residency, store).await
153 }
154
155 async fn open_resolved(
156 self,
157 policy: SessionPolicy,
158 state: RuntimeSessionState,
159 store: Option<Arc<dyn RuntimePersistence>>,
160 ) -> Result<LashSession> {
161 let mut env = self.core.env.clone();
162 if let Some(provider) = self.provider.clone().or_else(|| self.core.provider.clone()) {
163 env.core.providers.provider_resolver =
164 Arc::new(lash_core::SingleProviderResolver::new(provider));
165 }
166 let plugin_host = build_plugin_host(
167 self.core.protocol_factory.as_ref(),
168 self.core.plugin_factories.as_ref(),
169 self.plugin_factories,
170 )?;
171 env.core = self
172 .core
173 .runtime_host_for_plugin_host(env.core.clone(), &plugin_host)?;
174 env.plugin_host = Some(Arc::new(plugin_host));
175 let effect_host = Arc::clone(&env.core.control.effect_host);
176 let drivers = self.core.work_driver.drivers().await;
177 env.process_work_driver = drivers.process.clone();
178 env.queued_work_driver = drivers.queued.clone();
179 let runtime = LashRuntime::from_environment(&env, policy, state, store).await?;
180 if drivers.drive_process_on_open
181 && let Some(driver) = drivers.process.as_ref()
182 {
183 driver.claim_and_run_pending("session_open").await?;
184 }
185 let handle = RuntimeHandle::with_live_replay_store(
186 runtime,
187 Arc::clone(&self.core.live_replay_store),
188 );
189 Ok(LashSession {
190 runtime: handle,
191 effect_host,
192 parent_session_id: self.parent_session_id,
193 active_plugins: self.active_plugins,
194 process_phase_probe_slot: self.core.work_driver.phase_probe_slot(),
195 turn_cancels: crate::turn::TurnCancelRegistry::default(),
196 })
197 }
198
199 async fn create_store(
200 &self,
201 policy: &SessionPolicy,
202 ) -> Result<Option<Arc<dyn RuntimePersistence>>> {
203 if let Some(store) = self.store.as_ref() {
204 return Ok(Some(Arc::clone(store)));
205 }
206 let Some(factory) = self.core.store_factory.as_ref() else {
207 return Ok(None);
208 };
209 let request = SessionStoreCreateRequest {
210 session_id: self.session_id.clone(),
211 relation: self
212 .parent_session_id
213 .as_ref()
214 .map(|parent_session_id| lash_core::SessionRelation::Child {
215 parent_session_id: parent_session_id.clone(),
216 caused_by: None,
217 })
218 .unwrap_or_default(),
219 policy: policy.clone(),
220 };
221 factory
222 .create_store(&request)
223 .await
224 .map(Some)
225 .map_err(|message| EmbedError::StoreFactory {
226 session_id: self.session_id.clone(),
227 message,
228 })
229 }
230}
231
232pub(crate) async fn load_state_for_residency(
233 residency: Residency,
234 session_id: &str,
235 policy: &SessionPolicy,
236 store: &dyn RuntimePersistence,
237) -> Result<RuntimeSessionState> {
238 let mut state = load_persisted_state_for_residency(residency, store)
239 .await?
240 .unwrap_or_else(|| RuntimeSessionState {
241 session_id: session_id.to_string(),
242 policy: policy.clone(),
243 ..RuntimeSessionState::default()
244 });
245 if state.session_id != session_id {
246 return Err(EmbedError::StoreSessionMismatch {
247 loaded: state.session_id,
248 requested: session_id.to_string(),
249 });
250 }
251 let recorded_provider_id = state.policy.recorded_provider_id().to_string();
252 state.policy = policy.clone();
253 state.policy.provider_id = recorded_provider_id;
254 Ok(state)
255}
256
257async fn load_persisted_state_for_residency(
258 residency: Residency,
259 store: &dyn RuntimePersistence,
260) -> Result<Option<RuntimeSessionState>> {
261 match residency {
262 Residency::KeepAll => {
263 let loaded = lash_core::store::load_persisted_session_state(store)
264 .await
265 .map_err(|err| SessionError::Protocol(format!("failed to load store: {err}")))?;
266 Ok(loaded)
267 }
268 Residency::ActivePathOnly => {
269 let active = lash_core::store::load_persisted_session_state_active_path(store, None)
270 .await
271 .map_err(|err| {
272 SessionError::Protocol(format!("failed to load active-path store: {err}"))
273 })?;
274 if active
275 .as_ref()
276 .is_some_and(|state| state.session_graph.nodes.is_empty())
277 {
278 let mut full = lash_core::store::load_persisted_session_state(store)
279 .await
280 .map_err(|err| {
281 SessionError::Protocol(format!(
282 "failed to heal active-path store from full graph: {err}"
283 ))
284 })?;
285 if let Some(state) = full.as_mut() {
286 state.graph_replace_required = true;
287 }
288 return Ok(full);
289 }
290 Ok(active)
291 }
292 }
293}
294
295impl PromptLayerSink for SessionBuilder {
296 fn prompt_layer_mut(&mut self) -> &mut PromptLayer {
297 self.spec.prompt.get_or_insert_with(PromptLayer::new)
298 }
299}
300
301#[cfg(feature = "rlm")]
302impl RlmSessionBuilder {
303 pub fn provider(mut self, provider: ProviderHandle) -> Self {
304 self.builder = self.builder.provider(provider);
305 self
306 }
307
308 pub fn session_spec(mut self, spec: SessionSpec) -> Self {
309 self.builder = self.builder.session_spec(spec);
310 self
311 }
312
313 pub fn parent(mut self, parent_session_id: impl Into<String>) -> Self {
314 self.builder = self.builder.parent(parent_session_id);
315 self
316 }
317
318 pub fn store(mut self, store: Arc<dyn RuntimePersistence>) -> Self {
319 self.builder = self.builder.store(store);
320 self
321 }
322
323 pub fn plugin<P: PluginBinding>(mut self, config: P::SessionConfig) -> Self {
324 self.builder = self.builder.plugin::<P>(config);
325 self
326 }
327
328 pub async fn open(self) -> Result<LashSession> {
329 self.open_resolved(RlmOpenState::Resume).await
330 }
331
332 pub async fn open_fresh(self) -> Result<LashSession> {
333 self.open_resolved(RlmOpenState::Fresh).await
334 }
335
336 pub async fn open_with_state(self, state: RuntimeSessionState) -> Result<LashSession> {
337 self.open_resolved(RlmOpenState::Explicit(state)).await
338 }
339
340 async fn open_resolved(self, open_state: RlmOpenState) -> Result<LashSession> {
341 let Self {
342 builder,
343 rlm_final_answer_format,
344 } = self;
345 let policy = builder.session_policy();
346 let store = builder.create_store(&policy).await?;
347 let mut state = match open_state {
348 RlmOpenState::Resume => {
349 builder
350 .load_or_default_state(&policy, store.as_deref())
351 .await?
352 }
353 RlmOpenState::Fresh => RuntimeSessionState {
354 session_id: builder.session_id.clone(),
355 policy: policy.clone(),
356 graph_replace_required: true,
357 ..RuntimeSessionState::default()
358 },
359 RlmOpenState::Explicit(mut state) => {
360 if state.session_id != builder.session_id {
361 return Err(EmbedError::StoreSessionMismatch {
362 loaded: state.session_id,
363 requested: builder.session_id.clone(),
364 });
365 }
366 let recorded_provider_id = state.policy.recorded_provider_id().to_string();
367 state.policy = policy.clone();
368 state.policy.provider_id = recorded_provider_id;
369 state
370 }
371 };
372 apply_rlm_session_options(
373 builder.parent_session_id.is_none(),
374 rlm_final_answer_format,
375 &mut state,
376 )?;
377 builder.open_resolved(policy, state, store).await
378 }
379}
380
381#[cfg(feature = "rlm")]
382impl PromptLayerSink for RlmSessionBuilder {
383 fn prompt_layer_mut(&mut self) -> &mut PromptLayer {
384 self.builder.prompt_layer_mut()
385 }
386}
387
388#[cfg(feature = "rlm")]
389enum RlmOpenState {
390 Resume,
391 Fresh,
392 Explicit(RuntimeSessionState),
393}
394
395#[cfg(feature = "rlm")]
396fn apply_rlm_session_options(
397 is_root_session: bool,
398 explicit_format: Option<lash_rlm_types::RlmFinalAnswerFormat>,
399 state: &mut RuntimeSessionState,
400) -> Result<()> {
401 let final_answer_format = explicit_format.unwrap_or_else(|| {
402 if is_root_session {
403 lash_rlm_types::RlmFinalAnswerFormat::Markdown
404 } else {
405 lash_rlm_types::RlmFinalAnswerFormat::RawSubmitValue
406 }
407 });
408 let mut extras = if state.protocol_turn_options.is_empty() {
409 lash_rlm_types::RlmCreateExtras::default()
410 } else {
411 state.protocol_turn_options.decode()?
412 };
413 extras.final_answer_format = Some(final_answer_format);
414 let options = ProtocolTurnOptions::typed(extras)?;
415 state.protocol_turn_options = options.clone();
416 for frame in &mut state.agent_frames {
417 frame.protocol_turn_options = options.clone();
418 }
419 Ok(())
420}
421
422#[derive(Clone)]
423pub struct LashSession {
424 pub(crate) runtime: RuntimeHandle,
425 pub(crate) effect_host: Arc<dyn EffectHost>,
426 pub(crate) parent_session_id: Option<String>,
427 pub(crate) active_plugins: Vec<ActivePluginBinding>,
428 pub(crate) process_phase_probe_slot: Option<lash_core::runtime::RuntimeTurnPhaseProbeSlot>,
429 pub(crate) turn_cancels: crate::turn::TurnCancelRegistry,
430}
431
432#[derive(Clone, Debug, Default)]
433pub struct SessionConfigPatch {
434 pub provider: Option<ProviderHandle>,
435 pub model: Option<ModelSpec>,
436 pub prompt: Option<PromptLayer>,
437}
438
439impl LashSession {
440 pub async fn close(self) -> Result<()> {
441 let runtime = self.runtime.writer();
442 let runtime = runtime.lock().await;
443 runtime.unregister_plugin_session()?;
444 Ok(())
445 }
446
447 pub fn session_id(&self) -> String {
448 self.runtime.observe().session_id().to_string()
449 }
450
451 pub fn policy_snapshot(&self) -> SessionPolicy {
452 self.runtime.observe().policy.clone()
453 }
454
455 pub fn observe(&self) -> ObservableSession {
456 ObservableSession {
457 runtime: self.runtime.clone(),
458 }
459 }
460
461 pub fn parent_session_id(&self) -> Option<&str> {
462 self.parent_session_id.as_deref()
463 }
464
465 pub fn effect_host(&self) -> Arc<dyn EffectHost> {
466 Arc::clone(&self.effect_host)
467 }
468
469 pub fn turn(&self, input: TurnInput) -> TurnBuilder {
470 TurnBuilder {
471 runtime: self.runtime.clone(),
472 effect_host: Arc::clone(&self.effect_host),
473 active_plugins: self.active_plugins.clone(),
474 input,
475 cancel: CancellationToken::new(),
476 cancels: self.turn_cancels.clone(),
477 protocol_turn_options: None,
478 provider: None,
479 model: None,
480 turn_id: None,
481 }
482 }
483
484 pub fn queued_turn(&self) -> QueuedTurnBuilder {
485 QueuedTurnBuilder {
486 runtime: self.runtime.clone(),
487 effect_host: Arc::clone(&self.effect_host),
488 cancel: CancellationToken::new(),
489 cancels: self.turn_cancels.clone(),
490 batch_ids: Vec::new(),
491 drain_id: None,
492 }
493 }
494
495 pub fn cancel_running_turns(&self) -> usize {
510 self.turn_cancels.cancel_all()
511 }
512
513 pub fn admin(&self) -> SessionAdmin {
514 SessionAdmin {
515 runtime: self.runtime.clone(),
516 }
517 }
518
519 pub async fn configure(&self, patch: SessionConfigPatch) -> Result<()> {
520 self.admin().config().update(patch).await
521 }
522
523 pub fn tools(&self) -> ToolAdmin {
524 ToolAdmin::new(self.admin())
525 }
526
527 pub fn commands(&self) -> SessionCommandAdmin {
528 self.admin().commands()
529 }
530
531 pub fn triggers(&self) -> SessionTriggerAdmin {
532 self.admin().triggers()
533 }
534
535 pub fn processes(&self) -> SessionProcessAdmin {
536 SessionProcessAdmin::new(self.admin())
537 }
538
539 pub fn plugin_actions(&self) -> PluginActions {
540 PluginActions {
541 control: self.admin(),
542 }
543 }
544
545 pub fn enqueue(&self, input: TurnInput) -> EnqueueTurnBuilder<'_> {
546 EnqueueTurnBuilder {
547 session: self,
548 input,
549 id: None,
550 delivery_policy: DeliveryPolicy::AfterCurrentTurnCommit,
551 slot_policy: SlotPolicy::Exclusive,
552 }
553 }
554
555 pub async fn queued_work(&self) -> Result<Vec<QueuedWorkBatch>> {
556 let observation = self.runtime.observe();
557 let store = observation.queue_store.as_ref().ok_or_else(|| {
558 EmbedError::Runtime(lash_core::RuntimeError::new(
559 lash_core::RuntimeErrorCode::StoreCommitFailed,
560 "queued work inspection requires a persistent runtime store",
561 ))
562 })?;
563 store
564 .list_pending_queued_work(observation.session_id())
565 .await
566 .map_err(|err| {
567 EmbedError::Runtime(lash_core::RuntimeError::new(
568 lash_core::RuntimeErrorCode::StoreCommitFailed,
569 err.to_string(),
570 ))
571 })
572 }
573
574 pub async fn cancel_queued_work_batch(
575 &self,
576 batch_id: &str,
577 ) -> Result<Option<QueuedWorkBatch>> {
578 let session_id = self.session_id();
579 self.runtime
580 .cancel_queued_work_batch(&session_id, batch_id)
581 .await
582 .map_err(EmbedError::Runtime)
583 }
584
585 pub async fn await_queued_work_batch(&self, batch_id: &str) -> Result<()> {
597 let observation = self.runtime.observe();
598 let store = observation.queue_store.clone().ok_or_else(|| {
599 EmbedError::Runtime(lash_core::RuntimeError::new(
600 lash_core::RuntimeErrorCode::StoreCommitFailed,
601 "queued work inspection requires a persistent runtime store",
602 ))
603 })?;
604 let session_id = observation.session_id().to_string();
605 drop(observation);
606 let mut delay = std::time::Duration::from_millis(25);
607 loop {
608 let pending = store
609 .list_pending_queued_work(&session_id)
610 .await
611 .map_err(|err| {
612 EmbedError::Runtime(lash_core::RuntimeError::new(
613 lash_core::RuntimeErrorCode::StoreCommitFailed,
614 err.to_string(),
615 ))
616 })?;
617 if !pending.iter().any(|batch| batch.batch_id == batch_id) {
618 return Ok(());
619 }
620 tokio::time::sleep(delay).await;
621 delay = (delay * 2).min(std::time::Duration::from_millis(400));
622 }
623 }
624
625 pub fn read_view(&self) -> SessionReadView {
626 self.runtime.observe().read_view.clone()
627 }
628
629 pub fn usage_report(&self) -> SessionUsageReport {
630 self.runtime.observe().usage_report.clone()
631 }
632
633 pub async fn set_turn_phase_probe(
634 &self,
635 probe: Arc<dyn lash_core::runtime::RuntimeTurnPhaseProbe>,
636 ) {
637 let writer = self.runtime.writer();
638 let mut runtime = writer.lock().await;
639 runtime.set_turn_phase_probe(Arc::clone(&probe));
640 self.runtime.publish_from(&runtime);
641 if let Some(slot) = &self.process_phase_probe_slot {
642 let observation = self.runtime.observe();
643 slot.set_for_session(observation.session_id(), Arc::clone(&probe));
644 let current_frame = observation.persisted_state.current_agent_frame_id.as_str();
645 if !current_frame.is_empty() {
646 let scope = lash_core::SessionScope::for_agent_frame(
647 observation.session_id(),
648 current_frame,
649 );
650 slot.set_for_scope(&scope, probe);
651 }
652 }
653 }
654}
655
656#[derive(Clone)]
657pub struct ObservableSession {
658 pub(crate) runtime: RuntimeHandle,
659}
660
661impl ObservableSession {
662 fn snapshot(&self) -> Arc<RuntimeObservation> {
663 self.runtime.observe()
664 }
665
666 pub fn current_observation(&self) -> SessionObservation {
667 self.runtime.current_session_observation()
668 }
669
670 pub fn resume_from_cursor(&self, cursor: &SessionCursor) -> Result<SessionResume> {
671 self.runtime
672 .resume_session_observation(cursor)
673 .map_err(live_replay_error)
674 }
675
676 pub fn subscribe_from_cursor(
677 &self,
678 cursor: &SessionCursor,
679 ) -> Result<SessionObservationSubscription> {
680 self.runtime
681 .subscribe_session_observation(cursor)
682 .map_err(live_replay_error)
683 }
684
685 pub fn session_id(&self) -> String {
686 self.snapshot().session_id().to_string()
687 }
688
689 pub fn policy_snapshot(&self) -> SessionPolicy {
690 self.snapshot().policy.clone()
691 }
692
693 pub fn read_view(&self) -> SessionReadView {
694 self.snapshot().read_view.clone()
695 }
696
697 pub fn usage_report(&self) -> SessionUsageReport {
698 self.snapshot().usage_report.clone()
699 }
700
701 pub fn tool_state(&self) -> Option<ToolState> {
702 self.snapshot().tool_state.clone()
703 }
704
705 pub fn active_tool_manifests(&self) -> Vec<ToolManifest> {
706 self.snapshot()
707 .tool_state
708 .as_ref()
709 .map(ToolState::tool_manifests)
710 .unwrap_or_default()
711 }
712
713 pub async fn list_process_handles(&self) -> Vec<ProcessHandleSummary> {
714 self.snapshot().list_process_handles().await
715 }
716
717 pub async fn list_all_process_handles(&self) -> Vec<ProcessHandleSummary> {
718 self.snapshot().list_all_process_handles().await
719 }
720
721 pub fn process_scope(&self) -> SessionScope {
722 self.snapshot().process_scope()
723 }
724}
725
726fn live_replay_error(err: lash_core::LiveReplayStoreError) -> EmbedError {
727 EmbedError::Runtime(lash_core::RuntimeError::new(
728 RuntimeErrorCode::Other("live_replay".to_string()),
729 err.to_string(),
730 ))
731}
732
733pub struct EnqueueTurnBuilder<'a> {
734 session: &'a LashSession,
735 input: TurnInput,
736 id: Option<String>,
737 delivery_policy: DeliveryPolicy,
738 slot_policy: SlotPolicy,
739}
740
741impl<'a> EnqueueTurnBuilder<'a> {
742 pub fn id(mut self, id: impl Into<String>) -> Self {
743 self.id = Some(id.into());
744 self
745 }
746
747 pub fn delivery_policy(mut self, policy: DeliveryPolicy) -> Self {
748 self.delivery_policy = policy;
749 self
750 }
751
752 pub fn slot_policy(mut self, policy: SlotPolicy) -> Self {
753 self.slot_policy = policy;
754 self
755 }
756
757 pub async fn send(self) -> Result<QueuedWorkBatch> {
758 let source_key = self.id.map(|id| format!("host:{id}"));
759 self.session
760 .runtime
761 .enqueue_turn_input(
762 self.input,
763 self.delivery_policy,
764 self.slot_policy,
765 source_key,
766 )
767 .await
768 .map_err(EmbedError::Runtime)
769 }
770}
771
772impl<'a> std::future::IntoFuture for EnqueueTurnBuilder<'a> {
773 type Output = Result<QueuedWorkBatch>;
774 type IntoFuture =
775 std::pin::Pin<Box<dyn std::future::Future<Output = Result<QueuedWorkBatch>> + 'a>>;
776
777 fn into_future(self) -> Self::IntoFuture {
778 Box::pin(self.send())
779 }
780}