lash_core/runtime/process/
engine.rs1use std::collections::BTreeMap;
2use std::future::Future;
3use std::pin::Pin;
4use std::sync::Arc;
5
6use tokio_util::sync::CancellationToken;
7
8use super::events::ProcessAwaitOutput;
9use super::model::{
10 ProcessExecutionContext, ProcessExecutionEnvSpec, ProcessIdentity, ProcessInput,
11 ProcessRegistration,
12};
13use super::registry::ProcessRegistry;
14
15#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
17pub struct SegmentHandover {
18 pub reason: crate::BoundaryReason,
19 #[serde(default, skip_serializing_if = "Option::is_none")]
20 pub program_hash: Option<String>,
21 pub engine_state: Vec<u8>,
22}
23
24#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
29pub struct PersistedSegmentHandover {
30 pub segment_ordinal: u64,
31 pub program_hash: String,
32 pub handover: SegmentHandover,
33}
34
35#[derive(Clone, Debug, PartialEq)]
37pub enum ProcessRunOutcome {
38 Terminal(Box<ProcessAwaitOutput>),
39 SegmentBoundary(SegmentHandover),
40}
41
42impl From<ProcessAwaitOutput> for ProcessRunOutcome {
43 fn from(output: ProcessAwaitOutput) -> Self {
44 Self::Terminal(Box::new(output))
45 }
46}
47
48pub type ProcessEngineShutdownFuture<'run> = Pin<Box<dyn Future<Output = ()> + Send + 'run>>;
49
50pub struct ProcessEngineRunGuard<'run> {
51 shutdown: Option<Box<dyn FnOnce() -> ProcessEngineShutdownFuture<'run> + Send + 'run>>,
52}
53
54impl<'run> ProcessEngineRunGuard<'run> {
55 pub(crate) fn new(
56 shutdown: impl FnOnce() -> ProcessEngineShutdownFuture<'run> + Send + 'run,
57 ) -> Self {
58 Self {
59 shutdown: Some(Box::new(shutdown)),
60 }
61 }
62
63 pub async fn shutdown(mut self) {
64 if let Some(shutdown) = self.shutdown.take() {
65 shutdown().await;
66 }
67 }
68}
69
70pub struct ProcessEngineRuntimeContext<'run> {
71 context: crate::RuntimeExecutionContext<'run>,
72 guard: ProcessEngineRunGuard<'run>,
73}
74
75impl<'run> ProcessEngineRuntimeContext<'run> {
76 pub(crate) fn new(
77 context: crate::RuntimeExecutionContext<'run>,
78 guard: ProcessEngineRunGuard<'run>,
79 ) -> Self {
80 Self { context, guard }
81 }
82
83 pub fn context(&self) -> &crate::RuntimeExecutionContext<'run> {
84 &self.context
85 }
86
87 pub fn context_mut(&mut self) -> &mut crate::RuntimeExecutionContext<'run> {
88 &mut self.context
89 }
90
91 pub fn into_parts(
92 self,
93 ) -> (
94 crate::RuntimeExecutionContext<'run>,
95 ProcessEngineRunGuard<'run>,
96 ) {
97 (self.context, self.guard)
98 }
99
100 pub async fn shutdown(self) {
101 self.guard.shutdown().await;
102 }
103}
104
105type RuntimeContextBuilder<'run> = Box<
106 dyn FnOnce(
107 Arc<crate::ToolCatalog>,
108 ) -> Result<ProcessEngineRuntimeContext<'run>, crate::PluginError>
109 + Send
110 + 'run,
111>;
112
113pub struct ProcessEngineRunContext<'run> {
114 registration: ProcessRegistration,
115 execution_context: ProcessExecutionContext,
116 registry: Arc<dyn ProcessRegistry>,
117 session_id: String,
118 plugins: Arc<crate::PluginSession>,
119 store: Option<Arc<dyn crate::RuntimePersistence>>,
120 session_store_factory: Option<Arc<dyn crate::SessionStoreFactory>>,
121 queued_work_driver: Option<crate::QueuedWorkDriver>,
122 process_registry_available: bool,
123 cancellation: CancellationToken,
124 turn_phase_probe: Option<Arc<dyn crate::runtime::RuntimeTurnPhaseProbe>>,
125 scoped_effect_controller: crate::ScopedEffectController<'run>,
126 handover: Option<SegmentHandover>,
127 runtime_context_builder: Option<RuntimeContextBuilder<'run>>,
128}
129
130impl<'run> ProcessEngineRunContext<'run> {
131 #[allow(clippy::too_many_arguments)]
132 pub(crate) fn new(
133 registration: ProcessRegistration,
134 execution_context: ProcessExecutionContext,
135 registry: Arc<dyn ProcessRegistry>,
136 session_id: String,
137 plugins: Arc<crate::PluginSession>,
138 store: Option<Arc<dyn crate::RuntimePersistence>>,
139 session_store_factory: Option<Arc<dyn crate::SessionStoreFactory>>,
140 queued_work_driver: Option<crate::QueuedWorkDriver>,
141 process_registry_available: bool,
142 cancellation: CancellationToken,
143 turn_phase_probe: Option<Arc<dyn crate::runtime::RuntimeTurnPhaseProbe>>,
144 scoped_effect_controller: crate::ScopedEffectController<'run>,
145 handover: Option<SegmentHandover>,
146 runtime_context_builder: RuntimeContextBuilder<'run>,
147 ) -> Self {
148 Self {
149 registration,
150 execution_context,
151 registry,
152 session_id,
153 plugins,
154 store,
155 session_store_factory,
156 queued_work_driver,
157 process_registry_available,
158 cancellation,
159 turn_phase_probe,
160 scoped_effect_controller,
161 handover,
162 runtime_context_builder: Some(runtime_context_builder),
163 }
164 }
165
166 pub fn registration(&self) -> &ProcessRegistration {
167 &self.registration
168 }
169
170 pub fn execution_context(&self) -> &ProcessExecutionContext {
171 &self.execution_context
172 }
173
174 pub fn registry(&self) -> Arc<dyn ProcessRegistry> {
175 Arc::clone(&self.registry)
176 }
177
178 pub fn session_id(&self) -> &str {
179 &self.session_id
180 }
181
182 pub fn plugins(&self) -> Arc<crate::PluginSession> {
183 Arc::clone(&self.plugins)
184 }
185
186 pub fn store(&self) -> Option<Arc<dyn crate::RuntimePersistence>> {
187 self.store.clone()
188 }
189
190 pub fn session_store_factory(&self) -> Option<Arc<dyn crate::SessionStoreFactory>> {
191 self.session_store_factory.clone()
192 }
193
194 pub fn queued_work_driver(&self) -> Option<crate::QueuedWorkDriver> {
195 self.queued_work_driver.clone()
196 }
197
198 pub fn process_registry_available(&self) -> bool {
199 self.process_registry_available
200 }
201
202 pub fn cancellation_token(&self) -> CancellationToken {
203 self.cancellation.clone()
204 }
205
206 pub fn effect_controller(&self) -> &dyn crate::RuntimeEffectController {
207 self.scoped_effect_controller.controller()
208 }
209
210 pub fn scoped_effect_controller(&self) -> crate::ScopedEffectController<'run> {
211 self.scoped_effect_controller.clone()
212 }
213
214 pub fn take_handover(&mut self) -> Option<SegmentHandover> {
215 self.handover.take()
216 }
217
218 #[doc(hidden)]
219 pub fn named_phase(&self, phase: &'static str) -> crate::runtime::RuntimeNamedPhase {
220 crate::runtime::RuntimeNamedPhase::begin(self.turn_phase_probe.clone(), phase)
221 }
222
223 #[doc(hidden)]
224 pub fn turn_phase_probe(&self) -> Option<Arc<dyn crate::runtime::RuntimeTurnPhaseProbe>> {
225 self.turn_phase_probe.clone()
226 }
227
228 pub fn resolved_tool_catalog(&self) -> Result<Arc<crate::ToolCatalog>, crate::PluginError> {
229 self.plugins.resolved_tool_catalog(&self.session_id)
230 }
231
232 pub fn into_runtime_context(
233 mut self,
234 tool_catalog: Arc<crate::ToolCatalog>,
235 ) -> Result<ProcessEngineRuntimeContext<'run>, crate::PluginError> {
236 let builder = self.runtime_context_builder.take().ok_or_else(|| {
237 crate::PluginError::Session("process engine runtime context was already built".into())
238 })?;
239 builder(tool_catalog)
240 }
241}
242
243pub struct ProcessEngineValidationContext<'a> {
244 plugin_host: &'a crate::PluginHost,
245 tool_catalog: Arc<crate::ToolCatalog>,
246 process_registry_available: bool,
247}
248
249impl<'a> ProcessEngineValidationContext<'a> {
250 pub(crate) fn new(
251 plugin_host: &'a crate::PluginHost,
252 tool_catalog: Arc<crate::ToolCatalog>,
253 process_registry_available: bool,
254 ) -> Self {
255 Self {
256 plugin_host,
257 tool_catalog,
258 process_registry_available,
259 }
260 }
261
262 pub fn plugin_host(&self) -> &crate::PluginHost {
263 self.plugin_host
264 }
265
266 pub fn tool_catalog(&self) -> &crate::ToolCatalog {
267 self.tool_catalog.as_ref()
268 }
269
270 pub fn process_registry_available(&self) -> bool {
271 self.process_registry_available
272 }
273}
274
275#[async_trait::async_trait]
276pub trait ProcessEngine: Send + Sync {
283 fn kind(&self) -> &'static str;
284
285 async fn validate_start(
286 &self,
287 _context: ProcessEngineValidationContext<'_>,
288 _payload: &serde_json::Value,
289 _env_spec: Option<&ProcessExecutionEnvSpec>,
290 ) -> Result<(), crate::PluginError> {
291 Ok(())
292 }
293
294 async fn run(
295 &self,
296 context: ProcessEngineRunContext<'_>,
297 payload: serde_json::Value,
298 ) -> ProcessRunOutcome;
299
300 fn identity(&self, payload: &serde_json::Value) -> ProcessIdentity {
301 let _ = payload;
302 ProcessIdentity::new(self.kind())
303 }
304
305 fn durability_tier(&self) -> crate::DurabilityTier {
311 crate::DurabilityTier::Inline
312 }
313}
314
315#[derive(Clone, Default)]
316pub struct ProcessEngineRegistry {
317 engines: Arc<BTreeMap<String, Arc<dyn ProcessEngine>>>,
318}
319
320impl ProcessEngineRegistry {
321 pub fn new() -> Self {
322 Self::default()
323 }
324
325 pub fn with_engine(self, engine: Arc<dyn ProcessEngine>) -> Self {
326 let mut engines = (*self.engines).clone();
327 engines.insert(engine.kind().to_string(), engine);
328 Self {
329 engines: Arc::new(engines),
330 }
331 }
332
333 pub fn try_with_engine(
338 self,
339 engine: Arc<dyn ProcessEngine>,
340 ) -> Result<Self, crate::PluginError> {
341 if self.engines.contains_key(engine.kind()) {
342 return Err(crate::PluginError::Registration(format!(
343 "duplicate process engine kind `{}`; each engine kind may be registered once",
344 engine.kind()
345 )));
346 }
347 Ok(self.with_engine(engine))
348 }
349
350 pub fn get(&self, kind: &str) -> Option<Arc<dyn ProcessEngine>> {
351 self.engines.get(kind).cloned()
352 }
353
354 pub fn engines(&self) -> impl Iterator<Item = &Arc<dyn ProcessEngine>> {
357 self.engines.values()
358 }
359
360 pub fn require(&self, kind: &str) -> Result<Arc<dyn ProcessEngine>, crate::PluginError> {
361 self.get(kind).ok_or_else(|| {
362 crate::PluginError::Session(format!("process engine `{kind}` is not configured"))
363 })
364 }
365
366 pub fn validate_input(&self, input: &ProcessInput) -> Result<(), crate::PluginError> {
367 if let ProcessInput::Engine { kind, .. } = input {
368 self.require(kind).map(|_| ())
369 } else {
370 Ok(())
371 }
372 }
373}