Skip to main content

lash_core/plugin/
runtime_impl.rs

1use std::collections::{BTreeMap, BTreeSet};
2use std::sync::Arc;
3use std::sync::{Mutex as StdMutex, Weak};
4
5use super::*;
6
7#[derive(Clone)]
8pub struct PluginHost {
9    factories: Arc<Vec<Arc<dyn PluginFactory>>>,
10    extensions: PluginExtensions,
11    sessions: Arc<StdMutex<BTreeMap<String, Weak<PluginSession>>>>,
12}
13
14struct BuildPluginSessionRequest<'a> {
15    session_id: String,
16    parent_session_id: Option<String>,
17    snapshot: Option<&'a PluginSessionSnapshot>,
18    tool_catalog_overlay: ToolCatalogContribution,
19    tool_snapshot: Option<crate::ToolState>,
20    authority: SessionAuthorityContext,
21}
22
23#[derive(Clone, Debug, Default)]
24pub struct SessionAuthorityContext {
25    pub tool_access: SessionToolAccess,
26    pub subagent: Option<SubagentSessionContext>,
27    pub plugin_options: PluginOptions,
28}
29
30impl PluginHost {
31    pub fn empty() -> Self {
32        Self::new(Vec::new())
33    }
34
35    pub fn new(factories: Vec<Arc<dyn PluginFactory>>) -> Self {
36        let override_ids: BTreeSet<&'static str> =
37            factories.iter().map(|factory| factory.id()).collect();
38        let mut all_factories = super::builtin_plugin_factories();
39        if !override_ids.is_empty() {
40            all_factories.retain(|factory| !override_ids.contains(factory.id()));
41        }
42        all_factories.extend(factories);
43        let extensions = PluginExtensions::from_contributions(
44            all_factories
45                .iter()
46                .flat_map(|factory| factory.extension_contributions()),
47        );
48        Self {
49            factories: Arc::new(all_factories),
50            extensions,
51            sessions: Arc::new(StdMutex::new(BTreeMap::new())),
52        }
53    }
54
55    pub fn with_extensions(mut self, extensions: PluginExtensions) -> Self {
56        self.extensions = extensions;
57        self
58    }
59
60    pub fn isolated_registry(&self) -> Self {
61        Self {
62            factories: Arc::clone(&self.factories),
63            extensions: self.extensions.clone(),
64            sessions: Arc::new(StdMutex::new(BTreeMap::new())),
65        }
66    }
67
68    pub fn extensions(&self) -> &PluginExtensions {
69        &self.extensions
70    }
71
72    pub fn factories(&self) -> &[Arc<dyn PluginFactory>] {
73        self.factories.as_ref().as_slice()
74    }
75
76    pub fn build_session(
77        &self,
78        session_id: impl Into<String>,
79        snapshot: Option<&PluginSessionSnapshot>,
80    ) -> Result<Arc<PluginSession>, PluginError> {
81        self.build_session_with_overlay(
82            session_id,
83            snapshot,
84            ToolCatalogContribution::default(),
85            None,
86        )
87    }
88
89    /// Variant of [`build_session`] that records the caller as the
90    /// parent of the new session. Plugin factories read
91    /// [`PluginSessionContext::is_root_session`] to gate root-only
92    /// behavior; anything that goes through the plain `build_session`
93    /// is treated as a root session by default.
94    pub fn build_session_with_parent(
95        &self,
96        session_id: impl Into<String>,
97        parent_session_id: Option<String>,
98        snapshot: Option<&PluginSessionSnapshot>,
99        authority: SessionAuthorityContext,
100    ) -> Result<Arc<PluginSession>, PluginError> {
101        self.build_session_with_parent_and_overlay(
102            session_id,
103            parent_session_id,
104            snapshot,
105            ToolCatalogContribution::default(),
106            None,
107            authority,
108        )
109    }
110
111    pub fn build_session_with_parent_and_overlay(
112        &self,
113        session_id: impl Into<String>,
114        parent_session_id: Option<String>,
115        snapshot: Option<&PluginSessionSnapshot>,
116        tool_catalog_overlay: ToolCatalogContribution,
117        tool_snapshot: Option<crate::ToolState>,
118        authority: SessionAuthorityContext,
119    ) -> Result<Arc<PluginSession>, PluginError> {
120        self.build_session_inner(BuildPluginSessionRequest {
121            session_id: session_id.into(),
122            parent_session_id,
123            snapshot,
124            tool_catalog_overlay,
125            tool_snapshot,
126            authority,
127        })
128    }
129
130    pub fn build_session_with_overlay(
131        &self,
132        session_id: impl Into<String>,
133        snapshot: Option<&PluginSessionSnapshot>,
134        tool_catalog_overlay: ToolCatalogContribution,
135        tool_snapshot: Option<crate::ToolState>,
136    ) -> Result<Arc<PluginSession>, PluginError> {
137        self.build_session_inner(BuildPluginSessionRequest {
138            session_id: session_id.into(),
139            parent_session_id: None,
140            snapshot,
141            tool_catalog_overlay,
142            tool_snapshot,
143            authority: SessionAuthorityContext::default(),
144        })
145    }
146
147    fn build_session_inner(
148        &self,
149        request: BuildPluginSessionRequest<'_>,
150    ) -> Result<Arc<PluginSession>, PluginError> {
151        let BuildPluginSessionRequest {
152            session_id,
153            parent_session_id,
154            snapshot,
155            tool_catalog_overlay,
156            tool_snapshot,
157            authority,
158        } = request;
159        let ctx = PluginSessionContext {
160            session_id,
161            tool_access: authority.tool_access.clone(),
162            subagent: authority.subagent.clone(),
163            plugin_options: authority.plugin_options.clone(),
164            extensions: self.extensions.clone(),
165            parent_session_id,
166        };
167        let session_id = ctx.session_id.clone();
168        let mut tool_snapshot = tool_snapshot;
169        if let Some(snapshot) = &mut tool_snapshot {
170            let hidden_tools = &authority.tool_access.hidden_tools;
171            if !hidden_tools.is_empty() {
172                snapshot.retain(|name, _| !hidden_tools.contains(name));
173            }
174        }
175        let mut plugins = Vec::new();
176        let mut reg = PluginRegistrar::new();
177        for factory in self.factories() {
178            let plugin = factory.build(&ctx)?;
179            reg.registering_plugin_id = Some(plugin.id().to_string());
180            plugin.register(&mut reg)?;
181            reg.registering_plugin_id = None;
182            plugins.push(plugin);
183        }
184        let mut contributions = reg.contributions;
185        let protocol_session = contributions.protocol_session.take().ok_or_else(|| {
186            PluginError::Registration("missing protocol session capability".to_string())
187        })?;
188        let protocol_driver = contributions.protocol_driver.take().ok_or_else(|| {
189            PluginError::Registration("missing protocol driver capability".to_string())
190        })?;
191        contributions.protocol_session = Some(protocol_session);
192        contributions.protocol_driver = Some(protocol_driver);
193        contributions
194            .turn_context_transforms
195            .sort_by_key(|entry| std::cmp::Reverse(entry.0));
196        contributions
197            .context_compactors
198            .sort_by_key(|entry| std::cmp::Reverse(entry.0));
199        let triggers = crate::TriggerEventCatalog::from_events(contributions.triggers.clone())
200            .map_err(|message| {
201                PluginError::Registration(format!("invalid trigger event catalog: {message}"))
202            })?;
203        let registry = match tool_snapshot {
204            Some(snapshot) => Arc::new(
205                crate::ToolRegistry::from_tool_providers(contributions.tool_providers.clone())
206                    .map_err(|err| {
207                        PluginError::Registration(format!("failed to build tool registry: {err}"))
208                    })?
209                    .fork_with_state(snapshot)
210                    .map_err(|err| {
211                        PluginError::Session(format!(
212                            "tool state cannot be applied to this plugin host session: {err}"
213                        ))
214                    })?,
215            ),
216            None => Arc::new(
217                crate::ToolRegistry::from_tool_providers(contributions.tool_providers.clone())
218                    .map_err(|err| {
219                        PluginError::Registration(format!("failed to build tool registry: {err}"))
220                    })?,
221            ),
222        };
223        let tools = Arc::clone(&registry) as Arc<dyn ToolProvider>;
224
225        let session = Arc::new(PluginSession {
226            host: self.clone(),
227            session_id: ctx.session_id,
228            plugins,
229            tools,
230            tool_registry: registry,
231            tool_catalog_overlay,
232            tool_access: authority.tool_access,
233            subagent: authority.subagent,
234            extensions: self.extensions.clone(),
235            triggers,
236            contributions,
237        });
238        self.register_session(&session_id, &session)?;
239        let ready = SessionReadyContext {
240            session_id: session.session_id.clone(),
241            host: self.clone(),
242        };
243        for plugin in &session.plugins {
244            plugin.session_ready(ready.clone())?;
245        }
246        if let Some(snapshot) = snapshot {
247            session.restore(snapshot)?;
248        }
249        Ok(session)
250    }
251
252    pub async fn invoke_plugin_action_sessionless(
253        &self,
254        name: &str,
255        args: serde_json::Value,
256    ) -> Result<ToolResult, PluginError> {
257        let session = self.build_session(
258            format!("__external__-{}", uuid::Uuid::new_v4().simple()),
259            None,
260        )?;
261        session
262            .invoke_plugin_action(
263                name,
264                args,
265                None,
266                false,
267                Arc::new(NoopSessionManager),
268                Arc::new(NoopSessionManager),
269                Arc::new(NoopSessionManager),
270                Arc::new(crate::UnavailableProcessService),
271            )
272            .await
273            .map_err(|err| PluginError::Invoke(err.to_string()))
274    }
275
276    fn register_session(
277        &self,
278        session_id: &str,
279        session: &Arc<PluginSession>,
280    ) -> Result<(), PluginError> {
281        let mut sessions = self.sessions.lock().map_err(|_| {
282            PluginError::Session("plugin host session registry poisoned".to_string())
283        })?;
284        if let Some(existing) = sessions.get(session_id).and_then(Weak::upgrade) {
285            if !Arc::ptr_eq(&existing, session) {
286                return Err(PluginError::Session(format!(
287                    "session `{session_id}` is already registered on this plugin host"
288                )));
289            }
290            return Ok(());
291        }
292        sessions.insert(session_id.to_string(), Arc::downgrade(session));
293        Ok(())
294    }
295
296    pub fn unregister_session(&self, session_id: &str) -> Result<(), PluginError> {
297        let mut sessions = self.sessions.lock().map_err(|_| {
298            PluginError::Session("plugin host session registry poisoned".to_string())
299        })?;
300        sessions.remove(session_id);
301        Ok(())
302    }
303
304    pub fn session(&self, session_id: &str) -> Result<Arc<PluginSession>, PluginActionInvokeError> {
305        let mut sessions = self
306            .sessions
307            .lock()
308            .map_err(|_| PluginActionInvokeError::SessionRegistryPoisoned)?;
309        let Some(weak) = sessions.get(session_id).cloned() else {
310            return Err(PluginActionInvokeError::UnknownSession(
311                session_id.to_string(),
312            ));
313        };
314        match weak.upgrade() {
315            Some(session) => Ok(session),
316            None => {
317                sessions.remove(session_id);
318                Err(PluginActionInvokeError::UnknownSession(
319                    session_id.to_string(),
320                ))
321            }
322        }
323    }
324
325    #[expect(
326        clippy::too_many_arguments,
327        reason = "host action invocation wires the runtime service bundle at the plugin boundary"
328    )]
329    pub async fn invoke_plugin_action_for_session(
330        &self,
331        session_id: &str,
332        name: &str,
333        args: serde_json::Value,
334        sessions: Arc<dyn SessionStateService>,
335        session_lifecycle: Arc<dyn SessionLifecycleService>,
336        session_graph: Arc<dyn SessionGraphService>,
337        processes: Arc<dyn crate::ProcessService>,
338    ) -> Result<ToolResult, PluginActionInvokeError> {
339        let session = self.session(session_id)?;
340        session
341            .invoke_plugin_action(
342                name,
343                args,
344                Some(session_id.to_string()),
345                false,
346                sessions,
347                session_lifecycle,
348                session_graph,
349                processes,
350            )
351            .await
352    }
353}