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}