Skip to main content

lash_core/runtime/
builder.rs

1use std::path::PathBuf;
2use std::sync::Arc;
3
4use crate::plugin::{PluginFactory, PluginHost, PluginSession};
5use crate::{
6    BackgroundRuntimeHost, BackgroundTaskHost, EmbeddedRuntimeHost, LashRuntime,
7    PersistedSessionState, PersistentRuntimeServices, PluginStack, RuntimeCoreConfig,
8    RuntimePersistence, RuntimeServices, SessionError, SessionPolicy, SessionStoreFactory,
9    TerminationPolicy, TurnInjectionBridge, TurnInputInjectionBridge,
10};
11
12enum PluginSource {
13    Host(PluginHost),
14    Session(Arc<PluginSession>),
15}
16
17pub struct EmbeddedRuntimeBuilder {
18    session_id: Option<String>,
19    policy: Option<SessionPolicy>,
20    initial_state: Option<PersistedSessionState>,
21    plugin_source: PluginSource,
22    turn_injection_bridge: TurnInjectionBridge,
23    turn_input_injection_bridge: TurnInputInjectionBridge,
24    core: RuntimeCoreConfig,
25    session_store_factory: Option<Arc<dyn SessionStoreFactory>>,
26    store: Option<Arc<dyn RuntimePersistence>>,
27    background_task_host: Option<Arc<dyn BackgroundTaskHost>>,
28}
29
30impl Default for EmbeddedRuntimeBuilder {
31    fn default() -> Self {
32        Self {
33            session_id: None,
34            policy: None,
35            initial_state: None,
36            plugin_source: PluginSource::Host(PluginHost::empty()),
37            turn_injection_bridge: TurnInjectionBridge::new(),
38            turn_input_injection_bridge: TurnInputInjectionBridge::new(),
39            core: RuntimeCoreConfig::default(),
40            session_store_factory: None,
41            store: None,
42            background_task_host: None,
43        }
44    }
45}
46
47impl EmbeddedRuntimeBuilder {
48    pub fn new() -> Self {
49        Self::default()
50    }
51
52    pub fn session_id(&self) -> Option<&str> {
53        self.session_id.as_deref()
54    }
55
56    pub fn policy(&self) -> Option<&SessionPolicy> {
57        self.policy.as_ref()
58    }
59
60    pub fn with_session_id(mut self, session_id: impl Into<String>) -> Self {
61        self.session_id = Some(session_id.into());
62        self
63    }
64
65    pub fn with_policy(mut self, policy: SessionPolicy) -> Self {
66        self.policy = Some(policy);
67        self
68    }
69
70    pub fn with_initial_state(mut self, state: PersistedSessionState) -> Self {
71        self.initial_state = Some(state);
72        self
73    }
74
75    pub fn with_plugin_host(mut self, plugin_host: PluginHost) -> Self {
76        self.plugin_source = PluginSource::Host(if self.background_task_host.is_some() {
77            plugin_host.with_background_tasks()
78        } else {
79            plugin_host
80        });
81        self
82    }
83
84    pub fn with_plugin_session(mut self, plugin_session: Arc<PluginSession>) -> Self {
85        self.plugin_source = PluginSource::Session(plugin_session);
86        self
87    }
88
89    pub fn with_plugin_factories(mut self, factories: Vec<Arc<dyn PluginFactory>>) -> Self {
90        let host = PluginHost::new(factories);
91        self.plugin_source = PluginSource::Host(if self.background_task_host.is_some() {
92            host.with_background_tasks()
93        } else {
94            host
95        });
96        self
97    }
98
99    pub fn with_plugin_stack(self, stack: PluginStack) -> Self {
100        self.with_plugin_factories(stack.into_factories())
101    }
102
103    pub fn with_turn_injection_bridge(mut self, bridge: TurnInjectionBridge) -> Self {
104        self.turn_injection_bridge = bridge;
105        self
106    }
107
108    pub fn with_turn_input_injection_bridge(mut self, bridge: TurnInputInjectionBridge) -> Self {
109        self.turn_input_injection_bridge = bridge;
110        self
111    }
112
113    pub fn with_runtime_core(mut self, core: RuntimeCoreConfig) -> Self {
114        self.core = core;
115        self
116    }
117
118    pub fn with_attachment_store(
119        mut self,
120        attachment_store: Arc<dyn crate::AttachmentStore>,
121    ) -> Self {
122        self.core = self.core.with_attachment_store(attachment_store);
123        self
124    }
125
126    pub fn with_prompt_template(mut self, prompt_template: crate::PromptTemplate) -> Self {
127        self.core = self.core.with_prompt_template(prompt_template);
128        self
129    }
130
131    pub fn with_prompt_contribution(mut self, contribution: crate::PromptContribution) -> Self {
132        self.core = self.core.with_prompt_contribution(contribution);
133        self
134    }
135
136    pub fn with_replaced_prompt_slot(
137        mut self,
138        slot: crate::PromptSlot,
139        contributions: impl IntoIterator<Item = crate::PromptContribution>,
140    ) -> Self {
141        self.core = self.core.with_replaced_prompt_slot(slot, contributions);
142        self
143    }
144
145    pub fn with_cleared_prompt_slot(mut self, slot: crate::PromptSlot) -> Self {
146        self.core = self.core.with_cleared_prompt_slot(slot);
147        self
148    }
149
150    pub fn with_prompt_layer(mut self, prompt: crate::PromptLayer) -> Self {
151        self.core = self.core.with_prompt_layer(prompt);
152        self
153    }
154
155    pub fn with_trace_jsonl_path(mut self, trace_path: Option<PathBuf>) -> Self {
156        self.core = self.core.with_trace_jsonl_path(trace_path);
157        self
158    }
159
160    pub fn with_trace_sink(mut self, sink: Option<Arc<dyn lash_trace::TraceSink>>) -> Self {
161        self.core = self.core.with_trace_sink(sink);
162        self
163    }
164
165    pub fn with_trace_level(mut self, level: lash_trace::TraceLevel) -> Self {
166        self.core = self.core.with_trace_level(level);
167        self
168    }
169
170    pub fn with_trace_context(mut self, context: lash_trace::TraceContext) -> Self {
171        self.core = self.core.with_trace_context(context);
172        self
173    }
174
175    pub fn with_termination(mut self, termination: TerminationPolicy) -> Self {
176        self.core = self.core.with_termination(termination);
177        self
178    }
179
180    pub fn with_session_store_factory(
181        mut self,
182        session_store_factory: Arc<dyn SessionStoreFactory>,
183    ) -> Self {
184        self.session_store_factory = Some(session_store_factory);
185        self
186    }
187
188    pub fn with_store(mut self, store: Arc<dyn RuntimePersistence>) -> Self {
189        self.store = Some(store);
190        self
191    }
192
193    pub fn with_background_task_host(
194        mut self,
195        background_task_host: Arc<dyn BackgroundTaskHost>,
196    ) -> Self {
197        self.background_task_host = Some(background_task_host);
198        if let PluginSource::Host(host) = &mut self.plugin_source {
199            *host = host.clone().with_background_tasks();
200        }
201        self
202    }
203
204    fn resolve_state_from_defaults(&self) -> PersistedSessionState {
205        let mut state = self.initial_state.clone().unwrap_or_default();
206        if let Some(session_id) = &self.session_id {
207            state.session_id = session_id.clone();
208        }
209        if let Some(policy) = &self.policy {
210            state.policy = policy.clone();
211        }
212        state
213    }
214
215    async fn resolve_state(&self) -> Result<PersistedSessionState, SessionError> {
216        if let Some(state) = &self.initial_state {
217            return Ok({
218                let mut state = state.clone();
219                if let Some(session_id) = &self.session_id {
220                    state.session_id = session_id.clone();
221                }
222                if let Some(policy) = &self.policy {
223                    state.policy = policy.clone();
224                }
225                state
226            });
227        }
228        if let Some(store) = &self.store {
229            if let Some(mut state) = crate::store::load_persisted_session_state(store.as_ref())
230                .await
231                .map_err(|err| SessionError::Protocol(format!("failed to load store: {err}")))?
232            {
233                if let Some(session_id) = &self.session_id
234                    && &state.session_id != session_id
235                {
236                    return Err(SessionError::Protocol(format!(
237                        "store is bound to session `{}` but builder requested `{session_id}`",
238                        state.session_id
239                    )));
240                }
241                if let Some(policy) = &self.policy {
242                    state.policy = policy.clone();
243                }
244                return Ok(state);
245            }
246            let mut state = self.resolve_state_from_defaults();
247            if let Some(policy) = &self.policy {
248                state.policy = policy.clone();
249            }
250            return Ok(state);
251        }
252        Ok(self.resolve_state_from_defaults())
253    }
254
255    fn resolve_plugins(
256        &self,
257        state: &PersistedSessionState,
258    ) -> Result<Arc<PluginSession>, SessionError> {
259        match &self.plugin_source {
260            PluginSource::Session(session) => Ok(Arc::clone(session)),
261            PluginSource::Host(host) => host
262                .isolated_registry()
263                .build_session(
264                    state.session_id.clone(),
265                    state.policy.execution_mode.clone(),
266                    state.policy.standard_context_approach.clone(),
267                    None,
268                )
269                .map_err(|err| SessionError::Protocol(err.to_string())),
270        }
271    }
272
273    pub async fn build(self) -> Result<LashRuntime, SessionError> {
274        let state = self.resolve_state().await?;
275        let plugins = self.resolve_plugins(&state)?;
276        let embedded_host = EmbeddedRuntimeHost::new(self.core)
277            .with_session_store_factory_option(self.session_store_factory.clone());
278        match (self.store, self.background_task_host) {
279            (Some(store), Some(background_task_host)) => {
280                LashRuntime::from_persistent_background_state(
281                    state.policy.clone(),
282                    BackgroundRuntimeHost::new(embedded_host, background_task_host),
283                    PersistentRuntimeServices::new_with_bridges(
284                        plugins,
285                        self.turn_injection_bridge,
286                        self.turn_input_injection_bridge,
287                        store,
288                    ),
289                    state,
290                )
291                .await
292            }
293            (Some(store), None) => {
294                LashRuntime::from_persistent_embedded_state(
295                    state.policy.clone(),
296                    embedded_host,
297                    PersistentRuntimeServices::new_with_bridges(
298                        plugins,
299                        self.turn_injection_bridge,
300                        self.turn_input_injection_bridge,
301                        store,
302                    ),
303                    state,
304                )
305                .await
306            }
307            (None, Some(background_task_host)) => {
308                LashRuntime::from_background_state(
309                    state.policy.clone(),
310                    BackgroundRuntimeHost::new(embedded_host, background_task_host),
311                    RuntimeServices::new_with_bridges(
312                        plugins,
313                        self.turn_injection_bridge,
314                        self.turn_input_injection_bridge,
315                    ),
316                    state,
317                )
318                .await
319            }
320            (None, None) => {
321                LashRuntime::from_embedded_state(
322                    state.policy.clone(),
323                    embedded_host,
324                    RuntimeServices::new_with_bridges(
325                        plugins,
326                        self.turn_injection_bridge,
327                        self.turn_input_injection_bridge,
328                    ),
329                    state,
330                )
331                .await
332            }
333        }
334    }
335
336    pub async fn build_ephemeral(mut self) -> Result<LashRuntime, SessionError> {
337        self.store = None;
338        self.background_task_host = None;
339        if let PluginSource::Host(host) = &mut self.plugin_source {
340            *host = host.clone().with_background_tasks_available(false);
341        }
342        self.build().await
343    }
344
345    pub async fn build_persistent(
346        mut self,
347        store: Arc<dyn RuntimePersistence>,
348    ) -> Result<LashRuntime, SessionError> {
349        self.store = Some(store);
350        self.background_task_host = None;
351        if let PluginSource::Host(host) = &mut self.plugin_source {
352            *host = host.clone().with_background_tasks_available(false);
353        }
354        self.build().await
355    }
356
357    pub async fn build_background_persistent(
358        mut self,
359        store: Arc<dyn RuntimePersistence>,
360        background_task_host: Arc<dyn BackgroundTaskHost>,
361    ) -> Result<LashRuntime, SessionError> {
362        self.store = Some(store);
363        self = self.with_background_task_host(background_task_host);
364        self.build().await
365    }
366}
367
368impl LashRuntime {
369    pub fn builder() -> EmbeddedRuntimeBuilder {
370        EmbeddedRuntimeBuilder::new()
371    }
372}
373
374trait EmbeddedRuntimeHostExt {
375    fn with_session_store_factory_option(
376        self,
377        session_store_factory: Option<Arc<dyn SessionStoreFactory>>,
378    ) -> Self;
379}
380
381impl EmbeddedRuntimeHostExt for EmbeddedRuntimeHost {
382    fn with_session_store_factory_option(
383        mut self,
384        session_store_factory: Option<Arc<dyn SessionStoreFactory>>,
385    ) -> Self {
386        self.session_store_factory = session_store_factory;
387        self
388    }
389}