Skip to main content

lash_core/plugin/
registry.rs

1//! Plugin registration: `PluginSpec` (the declarative bundle of all a
2//! plugin's hooks), the `PluginFactory` / `SessionPlugin` traits
3//! plugin crates implement, and the two convenience factories
4//! (`StaticPluginFactory`, `PluginSpecFactory`) + the `SpecPlugin`
5//! glue that walks a spec and wires each field into the registrar.
6//!
7//! Split out of `plugin/mod.rs` for file size; outer path preserved by
8//! `pub use` in `plugin/mod.rs`.
9
10use std::sync::Arc;
11
12use super::{
13    AfterToolCallHook, AfterTurnHook, AssistantResponseHook, AssistantStreamHook,
14    BeforeToolCallHook, BeforeTurnHook, CheckpointHook, ContextCompactor, PluginAction,
15    PluginActionDef, PluginActionHandler, PluginError, PluginHost, PluginLifecycleEventHook,
16    PluginRegistrar, PluginSnapshotMeta, PromptContributor, SessionConfigMutator,
17    SessionToolAccess, SnapshotReader, SnapshotWriter, SubagentSessionContext,
18    ToolDiscoveryContributor, ToolResultProjector, ToolSurfaceContributor, TurnContextTransform,
19};
20use crate::{PluginOptions, ToolProvider};
21
22#[derive(Clone, Default)]
23pub struct PluginSpec {
24    pub tool_providers: Vec<Arc<dyn ToolProvider>>,
25    pub host_events: Vec<crate::HostEvent>,
26    pub prompt_contributors: Vec<PromptContributor>,
27    pub tool_surface_contributors: Vec<ToolSurfaceContributor>,
28    pub tool_discovery_contributors: Vec<ToolDiscoveryContributor>,
29    pub before_turn_hooks: Vec<BeforeTurnHook>,
30    pub before_tool_call_hooks: Vec<BeforeToolCallHook>,
31    pub after_tool_call_hooks: Vec<AfterToolCallHook>,
32    pub after_turn_hooks: Vec<AfterTurnHook>,
33    pub checkpoint_hooks: Vec<CheckpointHook>,
34    pub assistant_stream_hooks: Vec<AssistantStreamHook>,
35    pub assistant_response_hooks: Vec<AssistantResponseHook>,
36    pub tool_result_projector: Option<ToolResultProjector>,
37    pub runtime_event_hooks: Vec<PluginLifecycleEventHook>,
38    pub session_config_mutators: Vec<SessionConfigMutator>,
39    pub plugin_actions: Vec<(PluginActionDef, PluginActionHandler)>,
40    pub turn_context_transforms: Vec<(i32, Arc<dyn TurnContextTransform>)>,
41    pub context_compactors: Vec<(i32, Arc<dyn ContextCompactor>)>,
42}
43
44impl PluginSpec {
45    pub fn new() -> Self {
46        Self::default()
47    }
48
49    pub fn with_tool_provider(mut self, provider: Arc<dyn ToolProvider>) -> Self {
50        self.tool_providers.push(provider);
51        self
52    }
53
54    pub fn with_host_event(mut self, event: crate::HostEvent) -> Self {
55        self.host_events.push(event);
56        self
57    }
58
59    pub fn with_prompt_contributor(mut self, contributor: PromptContributor) -> Self {
60        self.prompt_contributors.push(contributor);
61        self
62    }
63
64    pub fn with_tool_surface_contributor(mut self, contributor: ToolSurfaceContributor) -> Self {
65        self.tool_surface_contributors.push(contributor);
66        self
67    }
68
69    pub fn with_tool_discovery_contributor(
70        mut self,
71        contributor: ToolDiscoveryContributor,
72    ) -> Self {
73        self.tool_discovery_contributors.push(contributor);
74        self
75    }
76
77    pub fn with_before_turn(mut self, hook: BeforeTurnHook) -> Self {
78        self.before_turn_hooks.push(hook);
79        self
80    }
81
82    pub fn with_before_tool_call(mut self, hook: BeforeToolCallHook) -> Self {
83        self.before_tool_call_hooks.push(hook);
84        self
85    }
86
87    pub fn with_after_tool_call(mut self, hook: AfterToolCallHook) -> Self {
88        self.after_tool_call_hooks.push(hook);
89        self
90    }
91
92    pub fn with_after_turn(mut self, hook: AfterTurnHook) -> Self {
93        self.after_turn_hooks.push(hook);
94        self
95    }
96
97    pub fn with_checkpoint(mut self, hook: CheckpointHook) -> Self {
98        self.checkpoint_hooks.push(hook);
99        self
100    }
101
102    pub fn with_assistant_stream(mut self, hook: AssistantStreamHook) -> Self {
103        self.assistant_stream_hooks.push(hook);
104        self
105    }
106
107    pub fn with_assistant_response(mut self, hook: AssistantResponseHook) -> Self {
108        self.assistant_response_hooks.push(hook);
109        self
110    }
111
112    pub fn with_tool_result_projector(mut self, projector: ToolResultProjector) -> Self {
113        self.tool_result_projector = Some(projector);
114        self
115    }
116
117    pub fn with_runtime_event(mut self, hook: PluginLifecycleEventHook) -> Self {
118        self.runtime_event_hooks.push(hook);
119        self
120    }
121
122    pub fn with_session_config_mutator(mut self, hook: SessionConfigMutator) -> Self {
123        self.session_config_mutators.push(hook);
124        self
125    }
126
127    pub fn with_plugin_action(
128        mut self,
129        def: PluginActionDef,
130        handler: PluginActionHandler,
131    ) -> Self {
132        self.plugin_actions.push((def, handler));
133        self
134    }
135
136    pub fn with_plugin_action_typed<Op, F, Fut>(self, handler: F) -> Self
137    where
138        Op: PluginAction,
139        F: Fn(super::PluginActionContext, Op::Args) -> Fut + Send + Sync + 'static,
140        Fut: std::future::Future<Output = Result<Op::Output, super::PluginActionFailure>>
141            + Send
142            + 'static,
143    {
144        self.with_plugin_action(
145            super::plugin_action_def::<Op>(),
146            Arc::new(move |ctx, args| {
147                let parsed = serde_json::from_value::<Op::Args>(args);
148                match parsed {
149                    Ok(args) => {
150                        let fut = handler(ctx, args);
151                        Box::pin(async move {
152                            match fut.await {
153                                Ok(output) => match serde_json::to_value(output) {
154                                    Ok(value) => crate::ToolResult::ok(value),
155                                    Err(err) => crate::ToolResult::err(serde_json::json!(format!(
156                                        "failed to serialize {} output: {err}",
157                                        Op::NAME
158                                    ))),
159                                },
160                                Err(err) => {
161                                    crate::ToolResult::err(serde_json::json!(err.to_string()))
162                                }
163                            }
164                        })
165                    }
166                    Err(err) => Box::pin(async move {
167                        crate::ToolResult::err(serde_json::json!(format!(
168                            "invalid {} args: {err}",
169                            Op::NAME
170                        )))
171                    }),
172                }
173            }),
174        )
175    }
176
177    pub fn with_plugin_action_sync<Op, F>(self, handler: F) -> Self
178    where
179        Op: PluginAction,
180        F: Fn(
181                super::PluginActionContext,
182                Op::Args,
183            ) -> Result<Op::Output, super::PluginActionFailure>
184            + Send
185            + Sync
186            + 'static,
187    {
188        self.with_plugin_action_typed::<Op, _, _>(move |ctx, args| {
189            let result = handler(ctx, args);
190            async move { result }
191        })
192    }
193
194    pub fn with_turn_context_transform(
195        mut self,
196        priority: i32,
197        transform: Arc<dyn TurnContextTransform>,
198    ) -> Self {
199        self.turn_context_transforms.push((priority, transform));
200        self
201    }
202
203    pub fn with_context_compactor(
204        mut self,
205        priority: i32,
206        compactor: Arc<dyn ContextCompactor>,
207    ) -> Self {
208        self.context_compactors.push((priority, compactor));
209        self
210    }
211}
212
213#[derive(Clone, Debug)]
214pub struct PluginSessionContext {
215    pub session_id: String,
216    pub tool_access: SessionToolAccess,
217    pub subagent: Option<SubagentSessionContext>,
218    pub plugin_options: PluginOptions,
219    pub lashlang_abilities: lashlang::LashlangAbilities,
220    pub lashlang_language_features: lashlang::LashlangLanguageFeatures,
221    /// Session id of the caller that created this one. `None` identifies
222    /// a root session; any subagent / compaction / forked-child session
223    /// carries the parent here so plugin factories can gate themselves
224    /// on root-only behavior (e.g. `update_plan`'s sticky plan dock).
225    pub parent_session_id: Option<String>,
226}
227
228impl PluginSessionContext {
229    /// Returns `true` when this context represents a root session, not a
230    /// subagent or internal child. Plugins that should only surface in
231    /// user-facing top-level turns check this in their `build`.
232    pub fn is_root_session(&self) -> bool {
233        self.parent_session_id.is_none()
234    }
235}
236
237#[derive(Clone)]
238pub struct SessionReadyContext {
239    pub session_id: String,
240    pub host: PluginHost,
241}
242
243pub trait SessionPlugin: Send + Sync {
244    fn id(&self) -> &'static str;
245
246    fn version(&self) -> &'static str {
247        "1"
248    }
249
250    fn register(&self, reg: &mut PluginRegistrar) -> Result<(), PluginError>;
251
252    fn snapshot(
253        &self,
254        _writer: &mut dyn SnapshotWriter,
255    ) -> Result<PluginSnapshotMeta, PluginError> {
256        Ok(PluginSnapshotMeta {
257            plugin_id: self.id().to_string(),
258            plugin_version: self.version().to_string(),
259            revision: self.snapshot_revision(),
260            state: None,
261        })
262    }
263
264    fn snapshot_revision(&self) -> u64 {
265        0
266    }
267
268    fn restore(
269        &self,
270        _meta: &PluginSnapshotMeta,
271        _reader: &dyn SnapshotReader,
272    ) -> Result<(), PluginError> {
273        Ok(())
274    }
275
276    fn session_ready(&self, _ctx: SessionReadyContext) -> Result<(), PluginError> {
277        Ok(())
278    }
279}
280
281/// Registers a plugin with the runtime and produces a per-session
282/// `SessionPlugin` instance for each new session.
283///
284/// # Cheap-build / stateful-factory contract
285///
286/// `build(ctx)` **must be cheap**. It runs on the hot path every time
287/// a new session is created (including subagents, forked children,
288/// and compaction children) and any latency here is paid per session.
289///
290/// Specifically, `build` must **not**:
291/// - perform any I/O (disk reads, HTTP calls, DB queries),
292/// - compile regexes, templates, or schemas,
293/// - open network connections or initialize connection pools,
294/// - load models, parse large config files, or allocate large buffers,
295/// - block the current thread for non-trivial work.
296///
297/// Expensive state belongs on the `PluginFactory` struct itself,
298/// wrapped in `Arc` so it can be cheaply cloned into per-session
299/// closures. The `PluginFactory` is constructed once by the embedder
300/// and held in the `RuntimeEnvironment`; its fields outlive every
301/// session. Hooks captured into a `PluginSpec` are closures that
302/// clone the `Arc`s off `self` and reference the shared state
303/// directly, so every session sees the same pool / cache / compiled
304/// artifact without rebuilding it.
305///
306/// The typical shape is:
307/// ```ignore
308/// pub struct MyFactory {
309///     pool: Arc<ConnectionPool>,          // expensive, built once
310///     compiled: Arc<Regex>,               // expensive, built once
311/// }
312///
313/// impl PluginFactory for MyFactory {
314///     fn id(&self) -> &'static str { "my_plugin" }
315///
316///     fn build(&self, _ctx: &PluginSessionContext)
317///         -> Result<Arc<dyn SessionPlugin>, PluginError>
318///     {
319///         // Cheap: clone Arcs, assemble spec, wrap in SpecPlugin.
320///         let pool = Arc::clone(&self.pool);
321///         let spec = PluginSpec::new().with_before_turn(Arc::new(move |_ctx| {
322///             let pool = Arc::clone(&pool);
323///             Box::pin(async move { /* use pool */ Ok(vec![]) })
324///         }));
325///         Ok(Arc::new(SpecPluginFromSpec::new("my_plugin", spec)))
326///     }
327/// }
328/// ```
329pub trait PluginFactory: Send + Sync {
330    fn id(&self) -> &'static str;
331
332    fn lashlang_abilities(&self) -> lashlang::LashlangAbilities {
333        lashlang::LashlangAbilities::default()
334    }
335
336    fn lashlang_language_features(&self) -> lashlang::LashlangLanguageFeatures {
337        lashlang::LashlangLanguageFeatures::default()
338    }
339
340    /// Host-owned Lashlang catalog entries that code may link against.
341    ///
342    /// This only affects the execution surface. It is intentionally not
343    /// rendered into prompts automatically; hosts remain responsible for
344    /// describing their resources through prompt contributions.
345    fn lashlang_resources(&self) -> lashlang::ResourceCatalog {
346        lashlang::ResourceCatalog::new()
347    }
348
349    /// Produce a session-scoped plugin. **Must be cheap** — see the
350    /// trait-level docs for the full contract.
351    fn build(&self, ctx: &PluginSessionContext) -> Result<Arc<dyn SessionPlugin>, PluginError>;
352}
353
354pub type PluginSpecBuilder =
355    Arc<dyn Fn(&PluginSessionContext) -> Result<PluginSpec, PluginError> + Send + Sync>;
356
357pub struct PluginSpecFactory {
358    id: &'static str,
359    builder: PluginSpecBuilder,
360}
361
362impl PluginSpecFactory {
363    pub fn new(id: &'static str, builder: PluginSpecBuilder) -> Self {
364        Self { id, builder }
365    }
366}
367
368pub struct StaticPluginFactory {
369    id: &'static str,
370    spec: PluginSpec,
371}
372
373impl StaticPluginFactory {
374    pub fn new(id: &'static str, spec: PluginSpec) -> Self {
375        Self { id, spec }
376    }
377}
378
379struct SpecPlugin {
380    id: &'static str,
381    spec: PluginSpec,
382}
383
384impl PluginFactory for PluginSpecFactory {
385    fn id(&self) -> &'static str {
386        self.id
387    }
388
389    fn build(&self, ctx: &PluginSessionContext) -> Result<Arc<dyn SessionPlugin>, PluginError> {
390        Ok(Arc::new(SpecPlugin {
391            id: self.id,
392            spec: (self.builder)(ctx)?,
393        }))
394    }
395}
396
397impl PluginFactory for StaticPluginFactory {
398    fn id(&self) -> &'static str {
399        self.id
400    }
401
402    fn build(&self, _ctx: &PluginSessionContext) -> Result<Arc<dyn SessionPlugin>, PluginError> {
403        Ok(Arc::new(SpecPlugin {
404            id: self.id,
405            spec: self.spec.clone(),
406        }))
407    }
408}
409
410impl SessionPlugin for SpecPlugin {
411    fn id(&self) -> &'static str {
412        self.id
413    }
414
415    fn register(&self, reg: &mut PluginRegistrar) -> Result<(), PluginError> {
416        for provider in &self.spec.tool_providers {
417            reg.tools().provider(Arc::clone(provider))?;
418        }
419        for event in &self.spec.host_events {
420            reg.host_events().declare(event.clone())?;
421        }
422        for contributor in &self.spec.prompt_contributors {
423            reg.prompt().contribute(Arc::clone(contributor));
424        }
425        for contributor in &self.spec.tool_surface_contributors {
426            reg.surface().contribute(Arc::clone(contributor));
427        }
428        for contributor in &self.spec.tool_discovery_contributors {
429            reg.discovery().contribute(Arc::clone(contributor));
430        }
431        for hook in &self.spec.before_turn_hooks {
432            reg.turn().before(Arc::clone(hook));
433        }
434        for hook in &self.spec.before_tool_call_hooks {
435            reg.tool_calls().before(Arc::clone(hook));
436        }
437        for hook in &self.spec.after_tool_call_hooks {
438            reg.tool_calls().after(Arc::clone(hook));
439        }
440        for hook in &self.spec.after_turn_hooks {
441            reg.turn().after(Arc::clone(hook));
442        }
443        for hook in &self.spec.checkpoint_hooks {
444            reg.turn().checkpoint(Arc::clone(hook));
445        }
446        for hook in &self.spec.assistant_stream_hooks {
447            reg.output().stream(Arc::clone(hook));
448        }
449        for hook in &self.spec.assistant_response_hooks {
450            reg.output().response(Arc::clone(hook));
451        }
452        if let Some(projector) = &self.spec.tool_result_projector {
453            reg.tool_results().projector(Arc::clone(projector))?;
454        }
455        for hook in &self.spec.runtime_event_hooks {
456            reg.session().on_event(Arc::clone(hook));
457        }
458        for hook in &self.spec.session_config_mutators {
459            reg.session().config_mutator(Arc::clone(hook));
460        }
461        for (def, handler) in &self.spec.plugin_actions {
462            reg.actions().op(def.clone(), Arc::clone(handler))?;
463        }
464        for (priority, transform) in &self.spec.turn_context_transforms {
465            reg.context().prepare_turn(*priority, Arc::clone(transform));
466        }
467        for (priority, compactor) in &self.spec.context_compactors {
468            reg.context().compact(*priority, Arc::clone(compactor));
469        }
470        Ok(())
471    }
472}