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

lash_core/runtime/
session_api.rs
