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