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

lash_core/plugin/
registry.rs
