Skip to main content

lash_core/tool_provider/
session.rs

1use std::sync::Arc;
2
3use crate::plugin::{
4    PluginError, SessionHandle, SessionLifecycleService, SessionSnapshot, SessionStateService,
5};
6
7#[derive(Clone, Debug, PartialEq, Eq)]
8pub struct ToolSessionModel {
9    pub model: String,
10    pub model_variant: Option<String>,
11    pub model_capability: crate::provider::ModelCapability,
12}
13
14#[derive(Clone)]
15pub struct ToolSessionAdmin<'run> {
16    pub(super) session_id: String,
17    pub(super) sessions: Arc<dyn SessionStateService>,
18    pub(super) session_lifecycle: Arc<dyn SessionLifecycleService>,
19    pub(super) effect_controller: crate::runtime::RuntimeEffectControllerHandle<'run>,
20}
21
22impl<'run> ToolSessionAdmin<'run> {
23    pub async fn model(&self) -> Result<ToolSessionModel, PluginError> {
24        let snapshot = self.snapshot_current().await?;
25        Ok(ToolSessionModel {
26            model: snapshot.policy.model.id,
27            model_variant: snapshot.policy.model.variant,
28            model_capability: snapshot.policy.model.capability,
29        })
30    }
31
32    pub async fn snapshot_current(&self) -> Result<SessionSnapshot, PluginError> {
33        self.snapshot(&self.session_id).await
34    }
35
36    pub async fn snapshot(
37        &self,
38        session_id: impl AsRef<str>,
39    ) -> Result<SessionSnapshot, PluginError> {
40        self.sessions.snapshot_session(session_id.as_ref()).await
41    }
42
43    pub async fn create_session(
44        &self,
45        request: crate::SessionCreateRequest,
46    ) -> Result<SessionHandle, PluginError> {
47        self.session_lifecycle.create_session(request).await
48    }
49
50    pub async fn close_session(&self, session_id: &str) -> Result<(), PluginError> {
51        self.session_lifecycle.close_session(session_id).await
52    }
53
54    pub async fn start_turn(
55        &self,
56        session_id: &str,
57        turn_id: &str,
58        input: crate::TurnInput,
59    ) -> Result<crate::AssembledTurn, PluginError> {
60        let scoped_effect_controller = crate::ScopedEffectController::borrowed(
61            self.effect_controller.controller(),
62            crate::ExecutionScope::turn(session_id, turn_id),
63        )
64        .map_err(|err| PluginError::Session(err.to_string()))?;
65        let request =
66            crate::SessionTurnRequest::new(session_id, turn_id, input, scoped_effect_controller)?;
67        self.session_lifecycle.start_turn(request).await
68    }
69
70    pub async fn tool_catalog(&self) -> Result<Vec<serde_json::Value>, PluginError> {
71        self.sessions.tool_catalog(&self.session_id).await
72    }
73
74    pub async fn shared_tool_catalog(&self) -> Result<Arc<Vec<serde_json::Value>>, PluginError> {
75        self.sessions.shared_tool_catalog(&self.session_id).await
76    }
77
78    pub async fn set_tool_membership(
79        &self,
80        names: &[String],
81        present: bool,
82    ) -> Result<u64, PluginError> {
83        self.sessions
84            .set_tool_membership(&self.session_id, names, present)
85            .await
86    }
87}