Skip to main content

lash_core/runtime/process/
engine.rs

1use 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/// Opaque engine-owned state carried between in-process execution segments.
16#[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/// The single bounded continuation durably retained for a process incarnation.
25///
26/// This is registry-internal execution state: it is deliberately not a process
27/// event and therefore never appears in change feeds or provenance.
28#[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/// Result of one process invocation. A segment boundary is never terminal.
36#[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]
276/// Deployment extension point for non-kernel process runtimes.
277///
278/// Core built-ins (`ToolCall`, `SessionTurn`, and `External`) are intentionally
279/// not registered here; they are kernel primitives with direct orchestration
280/// support. Implement `ProcessEngine` for process kinds stored as
281/// [`ProcessInput::Engine`](super::model::ProcessInput::Engine).
282pub 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    /// Durability tier this engine provides. Mirrors
306    /// [`EffectHost::durability_tier`](crate::runtime::AwaitEventResolver::durability_tier):
307    /// an engine that only persists inline reports [`DurabilityTier::Inline`](crate::DurabilityTier::Inline),
308    /// and the store-peer coherence validator rejects wiring it behind a durable
309    /// session store. Defaults to `Inline`.
310    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    /// Register an engine, rejecting a duplicate
334    /// [`ProcessEngine::kind`]. This is the single enforcement point for unique
335    /// engine kinds across everything registered on a runtime host, whether the
336    /// engine was wired directly or contributed through the plugin contract.
337    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    /// Iterate over every registered engine, used by the store-peer coherence
355    /// validator to sweep engine durability tiers.
356    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}