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

lash_core/plugin/
registry.rs
