lash_core/
session.rs

1use std::sync::{Arc, OnceLock};
2
3use tokio::sync::mpsc::UnboundedSender;
4
5use crate::PluginMessage;
6use crate::tool_dispatch::ToolDispatchContext;
7use crate::{
8    PromptContribution, RuntimeServices, SandboxMessage, SessionStreamEvent, ToolProvider,
9};
10
11mod execution_context;
12pub(crate) mod process_handles;
13mod tool_execution;
14
15pub use execution_context::RuntimeExecutionContext;
16pub(crate) use execution_context::RuntimeExecutionTracing;
17pub use tool_execution::{ToolInvocation, ToolInvocationReply};
18
19#[derive(Clone, Debug, PartialEq, Eq)]
20struct ToolCatalogCacheKey {
21    include_base_tools: bool,
22    context_overlay_revision: u64,
23    tool_generation: u64,
24    plugin_revision: u64,
25}
26
27#[derive(Debug, Default)]
28struct ToolCatalogDerived {
29    catalog: OnceLock<Arc<Vec<serde_json::Value>>>,
30}
31
32struct ToolCatalogArtifact {
33    tool_catalog: Arc<crate::ToolCatalog>,
34    preamble: Arc<crate::TurnDriverPreamble>,
35    derived: ToolCatalogDerived,
36}
37
38#[derive(Clone)]
39pub(crate) struct ToolCatalogHandle(Arc<ToolCatalogArtifact>);
40
41impl ToolCatalogHandle {
42    fn tool_catalog(&self) -> Arc<crate::ToolCatalog> {
43        Arc::clone(&self.0.tool_catalog)
44    }
45
46    fn preamble(&self) -> Arc<crate::TurnDriverPreamble> {
47        Arc::clone(&self.0.preamble)
48    }
49
50    fn catalog(&self) -> Arc<Vec<serde_json::Value>> {
51        Arc::clone(self.0.derived.catalog.get_or_init(|| {
52            Arc::new(crate::tool_registry::project_tool_catalog(
53                self.0.tool_catalog.tools.iter().cloned(),
54            ))
55        }))
56    }
57}
58
59#[derive(Clone, Debug)]
60pub struct InjectedTurnInput {
61    pub id: Option<String>,
62    pub message: PluginMessage,
63}
64
65#[derive(Debug, thiserror::Error)]
66pub enum SessionError {
67    #[error("I/O error: {0}")]
68    Io(#[from] std::io::Error),
69    #[error("JSON error: {0}")]
70    Json(#[from] serde_json::Error),
71    #[error("code execution is not available in this session")]
72    CodeExecutionUnavailable,
73    #[error("code execution runtime exited unexpectedly")]
74    CodeExecutionRuntimeStopped,
75    #[error(
76        "provider mismatch for session `{session_id}`: persisted provider `{expected}` does not match live provider `{actual}`"
77    )]
78    ProviderMismatch {
79        expected: String,
80        actual: String,
81        session_id: String,
82    },
83    #[error("provider is not configured for session `{session_id}`")]
84    ProviderUnconfigured { session_id: String },
85    #[error("provider `{provider_id}` is not registered for session `{session_id}`")]
86    ProviderUnavailable {
87        provider_id: String,
88        session_id: String,
89    },
90    #[error("protocol error: {0}")]
91    Protocol(String),
92}
93
94#[derive(Clone, Debug)]
95pub struct ExecRequest {
96    pub language: String,
97    pub code: String,
98    pub accept_finish: bool,
99}
100
101pub struct Session {
102    session_id: String,
103    services: RuntimeServices,
104    include_base_tools: bool,
105    context_overlay_revision: u64,
106    context_tools: Vec<Arc<dyn ToolProvider>>,
107    tool_registry: Arc<crate::ToolRegistry>,
108    context_prompt_contributions: Vec<PromptContribution>,
109    message_tx: Option<UnboundedSender<SandboxMessage>>,
110    tool_catalog_cache: Arc<std::sync::Mutex<Vec<(ToolCatalogCacheKey, ToolCatalogHandle)>>>,
111    /// Memoizes the rendered system prompt across turns. Most consecutive
112    /// turns reuse the same template + context overlay, so the cache hits
113    /// and we skip the section/Vec-join work in
114    /// `lash_sansio::PromptTemplate::render`.
115    prompt_cache: Arc<lash_sansio::PromptCache>,
116}
117
118impl Session {
119    pub async fn new(services: RuntimeServices, session_id: &str) -> Result<Self, SessionError> {
120        let tool_registry = services.plugins.tool_registry();
121        let mut session = Self {
122            session_id: session_id.to_string(),
123            services,
124            include_base_tools: true,
125            context_overlay_revision: 0,
126            context_tools: Vec::new(),
127            tool_registry,
128            context_prompt_contributions: Vec::new(),
129            message_tx: None,
130            tool_catalog_cache: Arc::new(std::sync::Mutex::new(Vec::new())),
131            prompt_cache: Arc::new(lash_sansio::PromptCache::new()),
132        };
133
134        let protocol_session = Arc::clone(session.plugins().protocol_session());
135        protocol_session
136            .initialize_session(crate::plugin::ProtocolSessionContext::new(
137                &mut session,
138                session_id,
139            ))
140            .await?;
141
142        Ok(session)
143    }
144
145    pub(crate) fn clone_for_effect(&self) -> Self {
146        Self {
147            session_id: self.session_id.clone(),
148            services: self.services.clone(),
149            include_base_tools: self.include_base_tools,
150            context_overlay_revision: self.context_overlay_revision,
151            context_tools: self.context_tools.clone(),
152            tool_registry: Arc::clone(&self.tool_registry),
153            context_prompt_contributions: self.context_prompt_contributions.clone(),
154            message_tx: self.message_tx.clone(),
155            tool_catalog_cache: Arc::clone(&self.tool_catalog_cache),
156            prompt_cache: Arc::clone(&self.prompt_cache),
157        }
158    }
159
160    pub fn session_id(&self) -> &str {
161        &self.session_id
162    }
163
164    pub(crate) fn protocol_extra_prompt_contributions(&self) -> Vec<PromptContribution> {
165        // Protocol-specific prompt contributions are owned by the protocol
166        // plugins via their
167        // `reg.prompt().contribute(...)` hooks. Nothing to add here.
168        Vec::new()
169    }
170
171    pub fn tools(&self) -> Arc<dyn ToolProvider> {
172        Arc::clone(&self.tool_registry) as Arc<dyn ToolProvider>
173    }
174
175    pub fn plugins(&self) -> &Arc<crate::PluginSession> {
176        &self.services.plugins
177    }
178
179    pub fn set_context_overlay(
180        &mut self,
181        tool_providers: Vec<Arc<dyn ToolProvider>>,
182        prompt_contributions: Vec<PromptContribution>,
183        include_base_tools: bool,
184    ) -> Result<(), crate::PluginError> {
185        let tool_providers_unchanged = self.context_tools.len() == tool_providers.len()
186            && self
187                .context_tools
188                .iter()
189                .zip(&tool_providers)
190                .all(|(current, next)| Arc::ptr_eq(current, next));
191        let overlay_unchanged = self.include_base_tools == include_base_tools
192            && self.context_prompt_contributions == prompt_contributions
193            && tool_providers_unchanged;
194        let registry = self
195            .services
196            .plugins
197            .tool_registry()
198            .compose_session_catalog(include_base_tools, tool_providers.clone())
199            .map(Arc::new)
200            .map_err(|err| {
201                crate::PluginError::Session(format!("failed to build session tool registry: {err}"))
202            })?;
203        self.include_base_tools = include_base_tools;
204        if !overlay_unchanged {
205            self.context_overlay_revision = self.context_overlay_revision.wrapping_add(1);
206        }
207        self.context_tools = tool_providers;
208        self.tool_registry = registry;
209        self.context_prompt_contributions = prompt_contributions;
210        self.tool_catalog_cache
211            .lock()
212            .expect("tool catalog cache lock")
213            .clear();
214        Ok(())
215    }
216
217    pub fn prompt_cache(&self) -> Arc<lash_sansio::PromptCache> {
218        Arc::clone(&self.prompt_cache)
219    }
220
221    pub fn context_prompt_contributions(&self) -> &[PromptContribution] {
222        &self.context_prompt_contributions
223    }
224
225    pub fn history_store(&self) -> Option<Arc<dyn crate::store::RuntimePersistence>> {
226        self.services.store.clone()
227    }
228
229    fn tool_catalog_cache_key(&self) -> ToolCatalogCacheKey {
230        ToolCatalogCacheKey {
231            include_base_tools: self.include_base_tools,
232            context_overlay_revision: self.context_overlay_revision,
233            tool_generation: self.tool_registry.generation(),
234            plugin_revision: self.plugins().snapshot_revision_fingerprint(),
235        }
236    }
237
238    fn build_tool_catalog_entry(
239        &self,
240        session_id: &str,
241    ) -> Result<ToolCatalogHandle, crate::PluginError> {
242        let provider = self.tools();
243        let tools = provider.tool_manifests();
244        let contract_provider = Arc::clone(&provider);
245        let resolve_contract: lash_sansio::ToolContractResolver =
246            Arc::new(move |name: &str| contract_provider.resolve_contract(name));
247        let tool_catalog = Arc::new(self.plugins().resolve_tool_catalog(
248            crate::plugin::ToolCatalogContext {
249                session_id: session_id.to_string(),
250                tools,
251                resolve_contract: Some(Arc::clone(&resolve_contract)),
252                tool_access: self.plugins().tool_access().clone(),
253                subagent: self.plugins().subagent_context().cloned(),
254                extensions: self.plugins().extensions().clone(),
255            },
256        )?);
257        let input = crate::ProtocolBuildInput {
258            tool_catalog: Arc::clone(&tool_catalog),
259            plugin_extensions: self.plugins().extensions().clone(),
260            trigger_events: self.plugins().triggers().clone(),
261            extra_prompt_contributions: self.protocol_extra_prompt_contributions(),
262        };
263        let driver = self.plugins().protocol_driver();
264        let preamble = driver.build_preamble(input);
265        Ok(ToolCatalogHandle(Arc::new(ToolCatalogArtifact {
266            tool_catalog,
267            preamble: Arc::new(preamble),
268            derived: ToolCatalogDerived::default(),
269        })))
270    }
271
272    fn tool_catalog_cache_entry(
273        &self,
274        session_id: &str,
275    ) -> Result<ToolCatalogHandle, crate::PluginError> {
276        let key = self.tool_catalog_cache_key();
277        let mut cache = self
278            .tool_catalog_cache
279            .lock()
280            .expect("tool catalog cache lock");
281        if let Some((_, entry)) = cache.iter().find(|(entry_key, _)| *entry_key == key) {
282            return Ok(entry.clone());
283        }
284        let entry = self.build_tool_catalog_entry(session_id)?;
285        cache.push((key, entry.clone()));
286        Ok(entry)
287    }
288
289    pub fn resolved_tool_catalog(
290        &self,
291        session_id: &str,
292    ) -> Result<Arc<crate::ToolCatalog>, crate::PluginError> {
293        Ok(self.tool_catalog_cache_entry(session_id)?.tool_catalog())
294    }
295
296    pub(crate) fn turn_driver_preamble(
297        &self,
298        session_id: &str,
299    ) -> Result<Arc<crate::TurnDriverPreamble>, crate::PluginError> {
300        Ok(self.tool_catalog_cache_entry(session_id)?.preamble())
301    }
302
303    pub(crate) fn shared_tool_catalog(
304        &self,
305        session_id: &str,
306    ) -> Result<Arc<Vec<serde_json::Value>>, crate::PluginError> {
307        Ok(self.tool_catalog_cache_entry(session_id)?.catalog())
308    }
309
310    pub fn tool_catalog(
311        &self,
312        session_id: &str,
313    ) -> Result<Vec<serde_json::Value>, crate::PluginError> {
314        Ok(self.shared_tool_catalog(session_id)?.as_ref().clone())
315    }
316
317    #[allow(
318        clippy::too_many_arguments,
319        reason = "code execution bridge carries explicit per-turn runtime dependencies"
320    )]
321    pub(crate) fn code_execution_context<'run>(
322        &self,
323        session_id: &str,
324        agent_frame_id: &str,
325        sessions: Arc<dyn crate::plugin::SessionStateService>,
326        session_lifecycle: Arc<dyn crate::plugin::SessionLifecycleService>,
327        session_graph: Arc<dyn crate::plugin::SessionGraphService>,
328        processes: Arc<dyn crate::ProcessService>,
329        process_cancel_ability: Arc<dyn crate::ProcessCancelAbility>,
330        effect_controller: crate::runtime::RuntimeEffectControllerHandle<'run>,
331        direct_completions: crate::DirectCompletionClient<'run>,
332        trigger_router: Option<crate::TriggerRouter>,
333        event_tx: tokio::sync::mpsc::Sender<SessionStreamEvent>,
334        chronological_projection: Arc<crate::ChronologicalProjection>,
335        protocol_extension: Option<crate::ProtocolTurnExtensionHandle>,
336        turn_context: crate::TurnContext,
337        execution_env_spec: crate::ProcessExecutionEnvSpec,
338        checkpoint_messages: crate::tool_dispatch::CheckpointMessageBuffer,
339        attachment_source_policy: Arc<dyn crate::AttachmentSourcePolicy>,
340    ) -> Result<RuntimeExecutionContext<'run>, crate::PluginError> {
341        let dispatch = Arc::new(ToolDispatchContext {
342            plugins: Arc::clone(self.plugins()),
343            tools: self.tools(),
344            tool_catalog: self.resolved_tool_catalog(session_id)?,
345            sessions,
346            session_lifecycle,
347            session_graph,
348            processes,
349            process_cancel_ability,
350            trigger_router,
351            effect_controller,
352            direct_completions: direct_completions.clone(),
353            parent_invocation: None,
354            execution_env_spec: execution_env_spec.clone(),
355            session_id: session_id.to_string(),
356            agent_frame_id: agent_frame_id.to_string(),
357            event_tx,
358            checkpoint_messages,
359            trigger_outcomes: crate::tool_dispatch::ToolTriggerOutcomeBuffer::default(),
360            attachment_store: Arc::clone(&self.services.attachment_store),
361            attachment_source_policy,
362            turn_context: turn_context.clone(),
363            clock: Arc::clone(&self.services.clock),
364        });
365        Ok(RuntimeExecutionContext::new(
366            session_id.to_string(),
367            dispatch,
368            Arc::clone(&self.services.process_env_store),
369            Arc::clone(&self.services.attachment_store),
370            chronological_projection,
371            protocol_extension,
372            turn_context,
373        ))
374        .map(|context| context.with_execution_env_spec(execution_env_spec))
375    }
376
377    /// Set the message sender for streaming messages during execution.
378    pub fn set_message_sender(&mut self, tx: UnboundedSender<SandboxMessage>) {
379        self.message_tx = Some(tx);
380    }
381
382    /// Clear the message sender (drops the sender, causing receivers to terminate).
383    pub fn clear_message_sender(&mut self) {
384        self.message_tx = None;
385    }
386
387    pub fn invalidate_runtime_caches(&self) {
388        self.tool_catalog_cache
389            .lock()
390            .expect("tool catalog cache lock")
391            .clear();
392        self.prompt_cache.clear();
393    }
394
395    pub async fn refresh_tool_catalog(&mut self) -> Result<(), SessionError> {
396        self.tool_registry = self
397            .services
398            .plugins
399            .tool_registry()
400            .compose_session_catalog(self.include_base_tools, self.context_tools.clone())
401            .map(Arc::new)
402            .map_err(|err| SessionError::Protocol(format!("tool reconfigure failed: {err}")))?;
403        self.tool_catalog_cache
404            .lock()
405            .expect("tool catalog cache lock")
406            .clear();
407        Ok(())
408    }
409}
lash_core/session.rs

lash_core/
session.rs
