lash_core/runtime/
session_api.rs

1use super::*;
2
3impl LashRuntime {
4    pub fn session_id(&self) -> &str {
5        &self.state.session_id
6    }
7
8    pub(super) fn stamp_live_plugin_state(&mut self) {
9        if let Some(session) = self.session.as_ref() {
10            let snapshot = session.plugins().tool_registry().export_state();
11            self.state.tool_state_generation = Some(snapshot.generation());
12            self.state.tool_state_snapshot = Some(snapshot);
13            let captured = session.plugins().snapshot();
14            crate::runtime::state::store_plugin_snapshot(&mut self.state.plugin_snapshot, captured);
15            self.state.plugin_snapshot_revision =
16                Some(session.plugins().snapshot_revision_fingerprint());
17        } else {
18            self.state.tool_state_generation = None;
19            self.state.tool_state_snapshot = None;
20            self.state.plugin_snapshot = None;
21            self.state.plugin_snapshot_revision = None;
22        }
23    }
24    pub(super) fn active_tool_catalog_shared(
25        &self,
26    ) -> Result<Arc<Vec<serde_json::Value>>, crate::PluginError> {
27        self.session
28            .as_ref()
29            .map(|session| session.shared_tool_catalog(&self.state.session_id))
30            .unwrap_or_else(|| Ok(Arc::new(Vec::new())))
31    }
32
33    pub fn tool_state(&self) -> Result<crate::ToolState, SessionError> {
34        let Some(session) = self.session.as_ref() else {
35            return Err(SessionError::Protocol(
36                "runtime session not available".to_string(),
37            ));
38        };
39        Ok(session.plugins().tool_registry().export_state())
40    }
41    /// Override protocol-owned turn options for this session.
42    pub fn set_protocol_turn_options(&mut self, options: crate::ProtocolTurnOptions) {
43        self.state.protocol_turn_options = options.clone();
44        if let Some(frame) = self.state.current_agent_frame_mut() {
45            frame.protocol_turn_options = options.clone();
46        }
47        self.protocol_turn_options = options;
48    }
49
50    /// The durable protocol turn options recorded on the session.
51    pub fn protocol_turn_options(&self) -> &crate::ProtocolTurnOptions {
52        &self.state.protocol_turn_options
53    }
54
55    /// Override protocol-owned turn options and mirror them to **every** agent
56    /// frame. Used by the protocol materialization hook (apply-at-open): the
57    /// last applied value is recorded on the session and on all frames.
58    pub fn set_protocol_turn_options_all_frames(&mut self, options: crate::ProtocolTurnOptions) {
59        self.state.protocol_turn_options = options.clone();
60        for frame in &mut self.state.agent_frames {
61            frame.protocol_turn_options = options.clone();
62        }
63        self.protocol_turn_options = options;
64    }
65
66    /// Run the protocol plugin's materialization hook against this runtime.
67    ///
68    /// Fires the [`ProtocolSessionPlugin::configure_runtime_on_materialize`]
69    /// hook, so both the child-create path and the root/builder-open path
70    /// converge on one seam. `plugin_options` are the plugin-keyed options that
71    /// reached this materialization (builder options for root opens, request
72    /// options for child create); `is_root_session` distinguishes root from
73    /// child.
74    pub fn configure_protocol_on_materialize(
75        &mut self,
76        plugin_options: &crate::PluginOptions,
77        is_root_session: bool,
78    ) -> Result<(), crate::PluginError> {
79        let protocol_session = self
80            .session
81            .as_ref()
82            .map(|session| Arc::clone(session.plugins().protocol_session()));
83        if let Some(protocol_session) = protocol_session {
84            let materialization = crate::plugin::ProtocolSessionMaterialization {
85                plugin_options,
86                is_root_session,
87            };
88            protocol_session
89                .configure_runtime_on_materialize(
90                    crate::plugin::ProtocolRuntimeContext::new(self),
91                    materialization,
92                )
93                .map_err(|err| crate::PluginError::Session(err.to_string()))?;
94        }
95        Ok(())
96    }
97
98    /// Export current session state for inspection/UI purposes.
99    /// This keeps persistence-heavy snapshots untouched; callers that need a
100    /// fully persisted view should use `export_persisted_state`.
101    pub fn export_state(&self) -> crate::SessionSnapshot {
102        self.state.to_snapshot()
103    }
104
105    pub fn read_view(&self) -> crate::SessionReadView {
106        crate::SessionReadView::from_runtime_state(
107            &self.state,
108            self.state.effective_policy().clone(),
109            self.state.effective_protocol_turn_options().clone(),
110        )
111    }
112
113    /// Export the narrow persistence snapshot used by stores and resume logic.
114    pub fn export_persistence_state(&self) -> RuntimeSessionState {
115        self.state.clone()
116    }
117
118    pub fn apply_persistence_state(
119        &mut self,
120        state: RuntimeSessionState,
121    ) -> Result<(), SessionError> {
122        self.set_persisted_state(state)
123    }
124
125    pub(crate) fn export_graph_first_state(&self) -> RuntimeSessionState {
126        self.state.clone()
127    }
128
129    /// Export a persistence-ready state envelope with dynamic/plugin snapshots
130    /// refreshed from the live session.
131    pub fn export_persisted_state(&self) -> RuntimeSessionState {
132        let mut state = self.state.clone();
133        state.protocol_turn_options = self.protocol_turn_options.clone();
134        if let Some(frame) = state.current_agent_frame_mut() {
135            frame.protocol_turn_options = self.protocol_turn_options.clone();
136        }
137        if let Some(session) = self.session.as_ref() {
138            let snapshot = session.plugins().tool_registry().export_state();
139            state.tool_state_generation = Some(snapshot.generation());
140            state.tool_state_snapshot = Some(snapshot);
141            let captured = session.plugins().snapshot();
142            crate::runtime::state::store_plugin_snapshot(&mut state.plugin_snapshot, captured);
143            state.plugin_snapshot_revision =
144                Some(session.plugins().snapshot_revision_fingerprint());
145        }
146        normalize_session_graph(&mut state);
147        state
148    }
149
150    pub fn usage_report(&self) -> SessionUsageReport {
151        let mut entries = self.state.token_ledger.clone();
152        let drained = self.shared_token_ledger.lock().expect("token ledger lock");
153        for entry in drained.iter().cloned() {
154            merge_ledger_entry(&mut entries, entry);
155        }
156        SessionUsageReport::from_entries(&entries)
157    }
158
159    pub async fn await_background_work(&mut self) -> Result<(), SessionError> {
160        if self.process_sync_needed.swap(false, Ordering::AcqRel) {
161            self.refresh_session_graph_from_store().await?;
162        }
163        Ok(())
164    }
165
166    pub(super) async fn refresh_session_graph_from_store(&mut self) -> Result<(), SessionError> {
167        // Fresh replacement opens intentionally start from an empty resident
168        // graph and commit a full replacement. Do not resurrect the old head
169        // before that first commit.
170        if self.state.graph_replace_required && self.state.head_revision.is_none() {
171            return Ok(());
172        }
173        let Some(store) = self
174            .session
175            .as_ref()
176            .and_then(|session| session.history_store())
177        else {
178            return Ok(());
179        };
180        let scope = match self.residency {
181            crate::Residency::KeepAll => crate::store::SessionReadScope::FullGraph,
182            crate::Residency::ActivePathOnly => crate::store::SessionReadScope::ActivePath {
183                leaf_node_id: self.state.session_graph.leaf_node_id.clone(),
184            },
185        };
186        let read = store.load_session(scope).await.map_err(|err| {
187            SessionError::Protocol(format!("failed to refresh session graph from store: {err}"))
188        })?;
189        self.graph_loaded_from_store = true;
190        let Some(read) = read else {
191            return Ok(());
192        };
193        let has_newer_graph = self.state.head_revision != Some(read.head_revision)
194            || read.graph.leaf_node_id != self.state.session_graph.leaf_node_id
195            || read.checkpoint_ref != self.state.checkpoint_ref;
196        if !has_newer_graph {
197            return Ok(());
198        }
199        let head = crate::store::SessionHead {
200            session_id: read.session_id.clone(),
201            head_revision: read.head_revision,
202            agent_frames: read.agent_frames.clone(),
203            current_agent_frame_id: read.current_agent_frame_id.clone(),
204            graph: read.graph,
205            config: read.config.clone(),
206            checkpoint_ref: read.checkpoint_ref.clone(),
207            token_ledger: merge_usage_delta_entries(read.token_ledger),
208        };
209        apply_session_head(&mut self.state, &head);
210        apply_session_checkpoint(&mut self.state, read.checkpoint);
211        self.policy = self.state.effective_policy().clone();
212        self.protocol_turn_options = self.state.effective_protocol_turn_options().clone();
213        Ok(())
214    }
215
216    pub(super) fn runtime_session_services(
217        &self,
218    ) -> Result<Arc<RuntimeSessionServices>, PluginOperationInvokeError> {
219        Ok(Arc::new(RuntimeSessionServices::new(self, true, None)?))
220    }
221
222    pub(super) fn runtime_session_services_for_turn(
223        &self,
224        child_usage_event_relay: Option<ChildUsageEventRelay>,
225    ) -> Result<Arc<RuntimeSessionServices>, PluginOperationInvokeError> {
226        Ok(Arc::new(RuntimeSessionServices::new(
227            self,
228            false,
229            child_usage_event_relay,
230        )?))
231    }
232
233    pub fn session_state_service(
234        &self,
235    ) -> Result<Arc<dyn crate::plugin::SessionStateService>, PluginOperationInvokeError> {
236        self.runtime_session_services()
237            .map(|services| services.state_service())
238    }
239
240    pub fn session_lifecycle_service(
241        &self,
242    ) -> Result<Arc<dyn crate::plugin::SessionLifecycleService>, PluginOperationInvokeError> {
243        self.runtime_session_services()
244            .map(|services| services.lifecycle_service())
245    }
246
247    pub fn session_graph_service(
248        &self,
249    ) -> Result<Arc<dyn crate::plugin::SessionGraphService>, PluginOperationInvokeError> {
250        self.runtime_session_services()
251            .map(|services| services.graph_service())
252    }
253
254    pub fn process_service(
255        &self,
256    ) -> Result<Arc<dyn crate::ProcessService>, PluginOperationInvokeError> {
257        self.runtime_session_services()
258            .map(|services| services.process_service())
259    }
260
261    pub fn process_cancel_ability(&self) -> Arc<dyn crate::ProcessCancelAbility> {
262        Arc::clone(&self.host.core.control.process_cancel_ability)
263    }
264
265    pub fn effect_host(&self) -> Arc<dyn crate::EffectHost> {
266        Arc::clone(&self.host.core.control.effect_host)
267    }
268
269    pub async fn enqueue_turn_input(
270        &self,
271        input: crate::TurnInput,
272        ingress: crate::TurnInputIngress,
273        source_key: Option<String>,
274    ) -> Result<crate::PendingTurnInput, RuntimeError> {
275        let store = self
276            .session
277            .as_ref()
278            .and_then(|session| session.history_store())
279            .ok_or_else(queued_turn_input_store_required)?;
280        enqueue_turn_input_to_store(
281            self.state.session_id.clone(),
282            store,
283            self.host.queued_work_driver.clone(),
284            input,
285            ingress,
286            source_key,
287        )
288        .await
289    }
290
291    pub async fn cancel_queued_work_batch(
292        &self,
293        session_id: &str,
294        batch_id: &str,
295    ) -> Result<Option<crate::QueuedWorkBatch>, RuntimeError> {
296        let store = self
297            .session
298            .as_ref()
299            .and_then(|session| session.history_store())
300            .ok_or_else(queued_turn_input_store_required)?;
301        store
302            .cancel_queued_work_batch(session_id, batch_id)
303            .await
304            .map_err(|err| RuntimeError::new(RuntimeErrorCode::StoreCommitFailed, err.to_string()))
305    }
306
307    /// The plugin session bound to the currently active runtime session, if any.
308    pub fn plugin_session(&self) -> Option<Arc<crate::PluginSession>> {
309        self.session.as_ref().map(|s| Arc::clone(s.plugins()))
310    }
311
312    pub fn open_agent_frame(
313        &mut self,
314        request: crate::OpenAgentFrameRequest,
315    ) -> crate::OpenAgentFrameResult {
316        open_agent_frame_in_state_with_clock(
317            &mut self.state,
318            request,
319            self.host.core.clock.as_ref(),
320        )
321    }
322
323    /// Run the registered compaction provider and commit the resulting
324    /// seed nodes into a fresh Agent Frame.
325    pub async fn compact_context(
326        &mut self,
327        instructions: Option<String>,
328        scoped_effect_controller: crate::ScopedEffectController<'_>,
329    ) -> Result<bool, PluginOperationInvokeError> {
330        let services = self.runtime_session_services()?;
331        let Some(plugin_session) = self.session.as_ref().map(|s| Arc::clone(s.plugins())) else {
332            return Err(PluginOperationInvokeError::Unknown(
333                "runtime session not available".to_string(),
334            ));
335        };
336        let ctx = crate::CompactionContext {
337            session_id: self.state.session_id.clone(),
338            state: self.read_view(),
339            instructions,
340            sessions: services.state_service(),
341            session_lifecycle: services.lifecycle_service(),
342            session_graph: services.graph_service(),
343            scoped_effect_controller,
344        };
345        let Some(compaction) = plugin_session.compact_context(&ctx).await.map_err(|err| {
346            PluginOperationInvokeError::Unknown(format!("context compaction failed: {err}"))
347        })?
348        else {
349            return Ok(false);
350        };
351        let frame_id = format!(
352            "{}:frame:compaction:{}",
353            self.state.session_id,
354            uuid::Uuid::new_v4()
355        );
356        let result = self.open_agent_frame(
357            crate::OpenAgentFrameRequest::new(frame_id, crate::AgentFrameReason::compaction())
358                .with_initial_nodes(compaction.initial_nodes),
359        );
360        if result.opened {
361            self.stamp_live_plugin_state();
362        }
363        Ok(result.opened)
364    }
365
366    pub(super) fn session_policy(&self) -> SessionPolicy {
367        self.policy.clone()
368    }
369
370    pub(super) async fn notify_session_config_changed(&self, previous: SessionPolicy) {
371        let Some(session) = self.session.as_ref() else {
372            return;
373        };
374        let current = self.session_policy();
375        if current == previous {
376            return;
377        }
378        let Ok(services) = self.runtime_session_services() else {
379            return;
380        };
381        session
382            .plugins()
383            .emit_runtime_event(crate::PluginLifecycleEvent::SessionConfigChanged(Box::new(
384                SessionConfigChangedContext {
385                    session_id: self.state.session_id.clone(),
386                    previous,
387                    current,
388                    sessions: services.state_service(),
389                },
390            )))
391            .await;
392    }
393
394    pub(super) async fn apply_session_config_mutations(&mut self, previous: SessionPolicy) {
395        let Some(session) = self.session.as_ref() else {
396            return;
397        };
398        let current = self.session_policy();
399        if current == previous {
400            return;
401        }
402        let Ok(services) = self.runtime_session_services() else {
403            return;
404        };
405        self.policy = session
406            .plugins()
407            .mutate_session_config(
408                SessionConfigChangedContext {
409                    session_id: self.state.session_id.clone(),
410                    previous,
411                    current,
412                    sessions: services.state_service(),
413                },
414                self.policy.clone(),
415            )
416            .await;
417        self.state.policy = self.policy.clone();
418    }
419}
420
421pub(in crate::runtime) async fn enqueue_turn_input_to_store(
422    session_id: String,
423    store: Arc<dyn crate::RuntimePersistence>,
424    queued_work_driver: Option<crate::QueuedWorkDriver>,
425    input: crate::TurnInput,
426    ingress: crate::TurnInputIngress,
427    source_key: Option<String>,
428) -> Result<crate::PendingTurnInput, RuntimeError> {
429    super::turn_loop::ensure_durable_effect_input(&input)?;
430    let is_next_turn = matches!(ingress, crate::TurnInputIngress::NextTurn);
431    let mut draft = crate::PendingTurnInputDraft::new(session_id, ingress, input);
432    draft.source_key = source_key;
433    let enqueued = store
434        .enqueue_pending_turn_input(draft)
435        .await
436        .map_err(|err| RuntimeError::new(RuntimeErrorCode::StoreCommitFailed, err.to_string()))?;
437    if is_next_turn && let Some(driver) = queued_work_driver.as_ref() {
438        driver
439            .claim_and_run_pending(Some(&enqueued.session_id), "queued_turn_input")
440            .await
441            .map_err(|err| {
442                RuntimeError::new(
443                    RuntimeErrorCode::Other("queued_work".to_string()),
444                    err.to_string(),
445                )
446            })?;
447    }
448    Ok(enqueued)
449}
450
451impl LashRuntime {
452    pub async fn submit_session_command(
453        &mut self,
454        command: crate::SessionCommand,
455        idempotency_key: impl Into<String>,
456    ) -> Result<crate::SessionCommandReceipt, RuntimeError> {
457        let idempotency_key = idempotency_key.into();
458        if idempotency_key.trim().is_empty() {
459            return Err(RuntimeError::new(
460                RuntimeErrorCode::Other("session_command_idempotency_key".to_string()),
461                "session command idempotency key cannot be empty",
462            ));
463        }
464        let source_key = command.source_key(&idempotency_key);
465        let session_id = self.state.session_id.clone();
466        let Some(store) = self
467            .session
468            .as_ref()
469            .and_then(|session| session.history_store())
470        else {
471            let batch_id = format!("inline-command:{}", uuid::Uuid::new_v4());
472            self.apply_session_command(command, None, None).await?;
473            return Ok(crate::SessionCommandReceipt {
474                session_id,
475                batch_id,
476                source_key,
477            });
478        };
479        let draft = crate::QueuedWorkBatchDraft::new(
480            session_id.clone(),
481            crate::DeliveryPolicy::AfterCurrentTurnCommit,
482            crate::SlotPolicy::Exclusive,
483            vec![crate::QueuedWorkPayload::session_command(command)],
484        )
485        .with_source_key(source_key.clone());
486        let enqueued = store.enqueue_queued_work(draft).await.map_err(|err| {
487            RuntimeError::new(RuntimeErrorCode::StoreCommitFailed, err.to_string())
488        })?;
489        if let Some(driver) = self.host.queued_work_driver.as_ref() {
490            driver
491                .claim_and_run_pending(Some(&session_id), "session_command")
492                .await
493                .map_err(|err| {
494                    RuntimeError::new(
495                        RuntimeErrorCode::Other("queued_work".to_string()),
496                        err.to_string(),
497                    )
498                })?;
499        }
500        Ok(crate::SessionCommandReceipt {
501            session_id,
502            batch_id: enqueued.batch_id,
503            source_key,
504        })
505    }
506
507    pub async fn drain_next_session_command(
508        &mut self,
509        session_execution_lease: &crate::SessionExecutionLeaseFence,
510    ) -> Result<Option<crate::SessionCommandReceipt>, RuntimeError> {
511        let Some(store) = self
512            .session
513            .as_ref()
514            .and_then(|session| session.history_store())
515        else {
516            return Ok(None);
517        };
518        let claim = store
519            .claim_leading_ready_session_command(
520                &self.state.session_id,
521                session_execution_lease,
522                &self.runtime_lease_owner,
523            )
524            .await
525            .map_err(super::runtime_error_from_store_commit)?;
526        let Some(claim) = claim else {
527            return Ok(None);
528        };
529        let Some((batch, command)) = claim.exclusive_session_command() else {
530            return Err(RuntimeError::new(
531                "session_command_claim",
532                format!(
533                    "queued-work claim `{}` did not contain exactly one session command batch",
534                    claim.claim_id
535                ),
536            ));
537        };
538        let batch_id = batch.batch_id.clone();
539        let source_key = batch.source_key.clone().unwrap_or_else(|| batch_id.clone());
540        let command = command.clone();
541        self.apply_session_command(
542            command,
543            Some(claim.completion()),
544            Some(session_execution_lease),
545        )
546        .await?;
547        Ok(Some(crate::SessionCommandReceipt {
548            session_id: self.state.session_id.clone(),
549            batch_id,
550            source_key,
551        }))
552    }
553
554    async fn apply_session_command(
555        &mut self,
556        command: crate::SessionCommand,
557        completion: Option<crate::QueuedWorkCompletion>,
558        session_execution_lease: Option<&crate::SessionExecutionLeaseFence>,
559    ) -> Result<(), RuntimeError> {
560        self.refresh_session_graph_from_store()
561            .await
562            .map_err(|err| RuntimeError::new("session_command_refresh", err.to_string()))?;
563        let graph = match command {
564            crate::SessionCommand::RefreshToolCatalog { .. } => {
565                self.refresh_session_tool_catalog().await.map_err(|err| {
566                    RuntimeError::new("session_command_refresh_tools", err.to_string())
567                })?;
568                crate::store::GraphCommitDelta::Unchanged {
569                    leaf_node_id: self.state.session_graph.leaf_node_id.clone(),
570                }
571            }
572            crate::SessionCommand::ResetSession { .. } => {
573                let mut state = crate::RuntimeSessionState {
574                    session_id: self.state.session_id.clone(),
575                    policy: self.policy.clone(),
576                    graph_replace_required: true,
577                    ..crate::RuntimeSessionState::default()
578                };
579                state.ensure_agent_frame_initialized();
580                self.set_persisted_state(state)
581                    .map_err(|err| RuntimeError::new("session_command_reset", err.to_string()))?;
582                crate::store::GraphCommitDelta::ReplaceFull(self.state.session_graph.clone())
583            }
584        };
585        let Some(store) = self
586            .session
587            .as_ref()
588            .and_then(|session| session.history_store())
589        else {
590            return Ok(());
591        };
592        let mut commit =
593            crate::store::RuntimeCommit::persisted_state_with_graph_commit(&self.state, graph, &[]);
594        let Some(session_execution_lease) = session_execution_lease else {
595            return Err(RuntimeError::new(
596                RuntimeErrorCode::StoreCommitFailed,
597                "session command commit requires a session execution lease",
598            ));
599        };
600        commit = commit.with_session_execution_lease(session_execution_lease.clone());
601        if let Some(completion) = completion {
602            commit = commit.completing_queue_claim(completion);
603        }
604        let result = store
605            .commit_runtime_state(commit)
606            .await
607            .map_err(super::runtime_error_from_store_commit)?;
608        self.state.apply_persisted_commit_result(result);
609        Ok(())
610    }
611}
612
613pub(in crate::runtime) fn queued_turn_input_store_required() -> RuntimeError {
614    RuntimeError::new(
615        RuntimeErrorCode::StoreCommitFailed,
616        "queued turn input requires a persistent runtime store",
617    )
618}
lash_core/runtime/session_api.rs

lash_core/runtime/
session_api.rs
