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::{PromptContribution, RuntimeServices, SandboxMessage, SessionEvent, ToolProvider};
8
9mod execution_context;
10pub(crate) mod process_handles;
11mod tool_execution;
12pub(crate) mod triggers;
13
14pub use execution_context::RuntimeExecutionContext;
15pub(crate) use execution_context::lashlang_surface_from_tool_surface;
16pub use tool_execution::{ToolInvocation, ToolInvocationReply};
17
18#[derive(Clone, Debug, PartialEq, Eq)]
19struct ToolSurfaceCacheKey {
20    include_base_tools: bool,
21    context_surface_revision: u64,
22    tool_generation: u64,
23    plugin_revision: u64,
24    lashlang_language_features: lashlang::LashlangLanguageFeatures,
25}
26
27#[derive(Debug, Default)]
28struct ToolSurfaceDerived {
29    catalog: OnceLock<Arc<Vec<serde_json::Value>>>,
30}
31
32struct ToolSurfaceArtifact {
33    surface: Arc<crate::ToolSurface>,
34    preamble: Arc<crate::TurnDriverPreamble>,
35    derived: ToolSurfaceDerived,
36}
37
38#[derive(Clone)]
39pub(crate) struct ToolSurfaceHandle(Arc<ToolSurfaceArtifact>);
40
41impl ToolSurfaceHandle {
42    fn surface(&self) -> Arc<crate::ToolSurface> {
43        Arc::clone(&self.0.surface)
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.surface.searchable_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 code: String,
97    pub accept_finish: bool,
98}
99
100pub struct Session {
101    session_id: String,
102    services: RuntimeServices,
103    include_base_tools: bool,
104    context_surface_revision: u64,
105    context_tools: Vec<Arc<dyn ToolProvider>>,
106    tool_registry: Arc<crate::ToolRegistry>,
107    context_prompt_contributions: Vec<PromptContribution>,
108    message_tx: Option<UnboundedSender<SandboxMessage>>,
109    tool_surface_cache: std::sync::Mutex<Vec<(ToolSurfaceCacheKey, ToolSurfaceHandle)>>,
110    /// Memoizes the rendered system prompt across turns. Most consecutive
111    /// turns reuse the same template + context surface, so the cache hits
112    /// and we skip the section/Vec-join work in
113    /// `lash_sansio::PromptTemplate::render`.
114    prompt_cache: Arc<lash_sansio::PromptCache>,
115}
116
117impl Session {
118    pub async fn new(services: RuntimeServices, session_id: &str) -> Result<Self, SessionError> {
119        let tool_registry = services.plugins.tool_registry();
120        let mut session = Self {
121            session_id: session_id.to_string(),
122            services,
123            include_base_tools: true,
124            context_surface_revision: 0,
125            context_tools: Vec::new(),
126            tool_registry,
127            context_prompt_contributions: Vec::new(),
128            message_tx: None,
129            tool_surface_cache: std::sync::Mutex::new(Vec::new()),
130            prompt_cache: Arc::new(lash_sansio::PromptCache::new()),
131        };
132
133        let protocol_session = Arc::clone(session.plugins().protocol_session());
134        protocol_session
135            .initialize_session(crate::plugin::ProtocolSessionContext::new(
136                &mut session,
137                session_id,
138            ))
139            .await?;
140
141        Ok(session)
142    }
143
144    pub fn session_id(&self) -> &str {
145        &self.session_id
146    }
147
148    pub(crate) fn protocol_extra_prompt_contributions(&self) -> Vec<PromptContribution> {
149        // Protocol-specific prompt contributions are owned by the protocol
150        // plugins via their
151        // `reg.prompt().contribute(...)` hooks. Nothing to add here.
152        Vec::new()
153    }
154
155    pub fn tools(&self) -> Arc<dyn ToolProvider> {
156        Arc::clone(&self.tool_registry) as Arc<dyn ToolProvider>
157    }
158
159    pub fn plugins(&self) -> &Arc<crate::PluginSession> {
160        &self.services.plugins
161    }
162
163    pub fn set_context_surface(
164        &mut self,
165        tool_providers: Vec<Arc<dyn ToolProvider>>,
166        prompt_contributions: Vec<PromptContribution>,
167        include_base_tools: bool,
168    ) -> Result<(), crate::PluginError> {
169        let tool_providers_unchanged = self.context_tools.len() == tool_providers.len()
170            && self
171                .context_tools
172                .iter()
173                .zip(&tool_providers)
174                .all(|(current, next)| Arc::ptr_eq(current, next));
175        if self.include_base_tools == include_base_tools
176            && self.context_prompt_contributions == prompt_contributions
177            && tool_providers_unchanged
178        {
179            return Ok(());
180        }
181        let registry = self
182            .services
183            .plugins
184            .tool_registry()
185            .compose_session_surface(include_base_tools, tool_providers.clone())
186            .map(Arc::new)
187            .map_err(|err| {
188                crate::PluginError::Session(format!("failed to build session tool registry: {err}"))
189            })?;
190        self.include_base_tools = include_base_tools;
191        self.context_surface_revision = self.context_surface_revision.wrapping_add(1);
192        self.context_tools = tool_providers;
193        self.tool_registry = registry;
194        self.context_prompt_contributions = prompt_contributions;
195        self.tool_surface_cache
196            .lock()
197            .expect("tool surface cache lock")
198            .clear();
199        Ok(())
200    }
201
202    pub fn prompt_cache(&self) -> Arc<lash_sansio::PromptCache> {
203        Arc::clone(&self.prompt_cache)
204    }
205
206    pub fn context_prompt_contributions(&self) -> &[PromptContribution] {
207        &self.context_prompt_contributions
208    }
209
210    pub fn history_store(&self) -> Option<Arc<dyn crate::store::RuntimePersistence>> {
211        self.services.store.clone()
212    }
213
214    fn tool_surface_cache_key(&self) -> ToolSurfaceCacheKey {
215        ToolSurfaceCacheKey {
216            include_base_tools: self.include_base_tools,
217            context_surface_revision: self.context_surface_revision,
218            tool_generation: self.tool_registry.generation(),
219            plugin_revision: self.plugins().snapshot_revision_fingerprint(),
220            lashlang_language_features: self.plugins().lashlang_language_features(),
221        }
222    }
223
224    fn build_tool_surface_entry(
225        &self,
226        session_id: &str,
227    ) -> Result<ToolSurfaceHandle, crate::PluginError> {
228        let provider = self.tools();
229        let tools = provider.tool_manifests();
230        let contract_provider = Arc::clone(&provider);
231        let resolve_contract: lash_sansio::ToolContractResolver =
232            Arc::new(move |name: &str| contract_provider.resolve_contract(name));
233        let surface = Arc::new(self.plugins().resolve_tool_surface(
234            crate::plugin::ToolSurfaceContext {
235                session_id: session_id.to_string(),
236                tools,
237                resolve_contract: Some(Arc::clone(&resolve_contract)),
238                tool_access: self.plugins().tool_access().clone(),
239                subagent: self.plugins().subagent_context().cloned(),
240                lashlang_abilities: self.plugins().lashlang_abilities(),
241            },
242        )?);
243        let input = crate::ProtocolBuildInput {
244            tool_surface: Arc::clone(&surface),
245            lashlang_surface: execution_context::lashlang_surface_from_tool_surface(
246                &surface,
247                self.plugins().lashlang_abilities(),
248                self.plugins().lashlang_language_features(),
249                self.plugins().lashlang_resources(),
250            ),
251            extra_prompt_contributions: self.protocol_extra_prompt_contributions(),
252        };
253        let driver = self.plugins().protocol_driver();
254        let preamble = driver.build_preamble(input);
255        Ok(ToolSurfaceHandle(Arc::new(ToolSurfaceArtifact {
256            surface,
257            preamble: Arc::new(preamble),
258            derived: ToolSurfaceDerived::default(),
259        })))
260    }
261
262    fn tool_surface_cache_entry(
263        &self,
264        session_id: &str,
265    ) -> Result<ToolSurfaceHandle, crate::PluginError> {
266        let key = self.tool_surface_cache_key();
267        let mut cache = self
268            .tool_surface_cache
269            .lock()
270            .expect("tool surface cache lock");
271        if let Some((_, entry)) = cache.iter().find(|(entry_key, _)| *entry_key == key) {
272            return Ok(entry.clone());
273        }
274        let entry = self.build_tool_surface_entry(session_id)?;
275        cache.push((key, entry.clone()));
276        Ok(entry)
277    }
278
279    pub fn tool_surface(
280        &self,
281        session_id: &str,
282    ) -> Result<Arc<crate::ToolSurface>, crate::PluginError> {
283        Ok(self.tool_surface_cache_entry(session_id)?.surface())
284    }
285
286    pub(crate) fn turn_driver_preamble(
287        &self,
288        session_id: &str,
289    ) -> Result<Arc<crate::TurnDriverPreamble>, crate::PluginError> {
290        Ok(self.tool_surface_cache_entry(session_id)?.preamble())
291    }
292
293    pub(crate) fn shared_tool_catalog(
294        &self,
295        session_id: &str,
296    ) -> Result<Arc<Vec<serde_json::Value>>, crate::PluginError> {
297        Ok(self.tool_surface_cache_entry(session_id)?.catalog())
298    }
299
300    pub fn tool_catalog(
301        &self,
302        session_id: &str,
303    ) -> Result<Vec<serde_json::Value>, crate::PluginError> {
304        Ok(self.shared_tool_catalog(session_id)?.as_ref().clone())
305    }
306
307    #[allow(
308        clippy::too_many_arguments,
309        reason = "code execution bridge carries explicit per-turn runtime dependencies"
310    )]
311    pub(crate) fn code_execution_context<'run>(
312        &self,
313        session_id: &str,
314        agent_frame_id: &str,
315        sessions: Arc<dyn crate::plugin::SessionStateService>,
316        session_lifecycle: Arc<dyn crate::plugin::SessionLifecycleService>,
317        session_graph: Arc<dyn crate::plugin::SessionGraphService>,
318        processes: Arc<dyn crate::ProcessService>,
319        process_cancel_ability: Arc<dyn crate::ProcessCancelAbility>,
320        effect_controller: crate::runtime::RuntimeEffectControllerHandle<'run>,
321        direct_completions: crate::DirectCompletionClient<'run>,
322        host_event_router: Option<crate::HostEventRouter>,
323        event_tx: tokio::sync::mpsc::Sender<SessionEvent>,
324        chronological_projection: Arc<crate::ChronologicalProjection>,
325        protocol_extension: Option<crate::ProtocolTurnExtensionHandle>,
326        turn_context: crate::TurnContext,
327        checkpoint_messages: crate::tool_dispatch::CheckpointMessageBuffer,
328    ) -> Result<RuntimeExecutionContext<'run>, crate::PluginError> {
329        let dispatch = Arc::new(ToolDispatchContext {
330            plugins: Arc::clone(self.plugins()),
331            tools: self.tools(),
332            surface: self.tool_surface(session_id)?,
333            sessions,
334            session_lifecycle,
335            session_graph,
336            processes,
337            process_cancel_ability,
338            host_event_router,
339            effect_controller,
340            direct_completions: direct_completions.clone(),
341            parent_invocation: None,
342            session_id: session_id.to_string(),
343            agent_frame_id: agent_frame_id.to_string(),
344            event_tx,
345            checkpoint_messages,
346            host_event_outcomes: crate::tool_dispatch::ToolHostEventOutcomeBuffer::default(),
347            attachment_store: Arc::clone(&self.services.attachment_store),
348            turn_context: turn_context.clone(),
349        });
350        Ok(RuntimeExecutionContext::new(
351            session_id.to_string(),
352            dispatch,
353            self.plugins().lashlang_abilities(),
354            self.plugins().lashlang_language_features(),
355            Arc::clone(&self.services.lashlang_artifact_store),
356            Arc::clone(&self.services.attachment_store),
357            chronological_projection,
358            protocol_extension,
359            turn_context,
360        ))
361    }
362
363    /// Set the message sender for streaming messages during execution.
364    pub fn set_message_sender(&mut self, tx: UnboundedSender<SandboxMessage>) {
365        self.message_tx = Some(tx);
366    }
367
368    /// Clear the message sender (drops the sender, causing receivers to terminate).
369    pub fn clear_message_sender(&mut self) {
370        self.message_tx = None;
371    }
372
373    pub fn invalidate_runtime_caches(&self) {
374        self.tool_surface_cache
375            .lock()
376            .expect("tool surface cache lock")
377            .clear();
378        self.prompt_cache.clear();
379    }
380
381    pub async fn refresh_tool_surface(&mut self) -> Result<(), SessionError> {
382        self.tool_registry = self
383            .services
384            .plugins
385            .tool_registry()
386            .compose_session_surface(self.include_base_tools, self.context_tools.clone())
387            .map(Arc::new)
388            .map_err(|err| SessionError::Protocol(format!("tool reconfigure failed: {err}")))?;
389        self.tool_surface_cache
390            .lock()
391            .expect("tool surface cache lock")
392            .clear();
393        Ok(())
394    }
395}
lash_core/session.rs

lash_core/
session.rs
