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    /// Ask every factory for its process-engine contributions and register them
77    /// on `runtime_host`, enforcing unique [`ProcessEngine::kind`](crate::ProcessEngine::kind)
78    /// across all engines (directly wired or plugin-contributed).
79    ///
80    /// This is the core-owned installation step that replaces facade-level
81    /// out-of-band wiring: engine construction that needs the fully-built plugin
82    /// host's extensions runs here, after the host is built. The trace context
83    /// handed to factories is the one already on `runtime_host`.
84    pub fn install_process_engine_contributions(
85        &self,
86        mut runtime_host: crate::runtime::RuntimeHostConfig,
87        process_lifecycle_available: bool,
88    ) -> Result<crate::runtime::RuntimeHostConfig, PluginError> {
89        let trace_context = runtime_host.tracing.trace_context.clone();
90        let ctx = super::ProcessEngineContributionContext::new(
91            &self.extensions,
92            &trace_context,
93            process_lifecycle_available,
94        );
95        for factory in self.factories() {
96            for engine in factory.process_engine_contributions(&ctx)? {
97                runtime_host.process_engines = runtime_host
98                    .process_engines
99                    .clone()
100                    .try_with_engine(engine)?;
101            }
102        }
103        Ok(runtime_host)
104    }
105
106    pub fn build_session(
107        &self,
108        session_id: impl Into<String>,
109        snapshot: Option<&PluginSessionSnapshot>,
110    ) -> Result<Arc<PluginSession>, PluginError> {
111        self.build_session_with_overlay(
112            session_id,
113            snapshot,
114            ToolCatalogContribution::default(),
115            None,
116        )
117    }
118
119    /// Variant of [`build_session`] that records the caller as the
120    /// parent of the new session. Plugin factories read
121    /// [`PluginSessionContext::is_root_session`] to gate root-only
122    /// behavior; anything that goes through the plain `build_session`
123    /// is treated as a root session by default.
124    pub fn build_session_with_parent(
125        &self,
126        session_id: impl Into<String>,
127        parent_session_id: Option<String>,
128        snapshot: Option<&PluginSessionSnapshot>,
129        authority: SessionAuthorityContext,
130    ) -> Result<Arc<PluginSession>, PluginError> {
131        self.build_session_with_parent_and_overlay(
132            session_id,
133            parent_session_id,
134            snapshot,
135            ToolCatalogContribution::default(),
136            None,
137            authority,
138        )
139    }
140
141    pub fn build_session_with_parent_and_overlay(
142        &self,
143        session_id: impl Into<String>,
144        parent_session_id: Option<String>,
145        snapshot: Option<&PluginSessionSnapshot>,
146        tool_catalog_overlay: ToolCatalogContribution,
147        tool_snapshot: Option<crate::ToolState>,
148        authority: SessionAuthorityContext,
149    ) -> Result<Arc<PluginSession>, PluginError> {
150        self.build_session_inner(BuildPluginSessionRequest {
151            session_id: session_id.into(),
152            parent_session_id,
153            snapshot,
154            tool_catalog_overlay,
155            tool_snapshot,
156            authority,
157        })
158    }
159
160    pub fn build_session_with_overlay(
161        &self,
162        session_id: impl Into<String>,
163        snapshot: Option<&PluginSessionSnapshot>,
164        tool_catalog_overlay: ToolCatalogContribution,
165        tool_snapshot: Option<crate::ToolState>,
166    ) -> Result<Arc<PluginSession>, PluginError> {
167        self.build_session_inner(BuildPluginSessionRequest {
168            session_id: session_id.into(),
169            parent_session_id: None,
170            snapshot,
171            tool_catalog_overlay,
172            tool_snapshot,
173            authority: SessionAuthorityContext::default(),
174        })
175    }
176
177    fn build_session_inner(
178        &self,
179        request: BuildPluginSessionRequest<'_>,
180    ) -> Result<Arc<PluginSession>, PluginError> {
181        let BuildPluginSessionRequest {
182            session_id,
183            parent_session_id,
184            snapshot,
185            tool_catalog_overlay,
186            tool_snapshot,
187            authority,
188        } = request;
189        let ctx = PluginSessionContext {
190            session_id,
191            tool_access: authority.tool_access.clone(),
192            subagent: authority.subagent.clone(),
193            plugin_options: authority.plugin_options.clone(),
194            extensions: self.extensions.clone(),
195            parent_session_id,
196        };
197        let session_id = ctx.session_id.clone();
198        let mut tool_snapshot = tool_snapshot;
199        if let Some(snapshot) = &mut tool_snapshot {
200            let hidden_tools = &authority.tool_access.hidden_tools;
201            if !hidden_tools.is_empty() {
202                snapshot.retain(|_, entry| !hidden_tools.contains(&entry.manifest().name));
203            }
204        }
205        let mut plugins = Vec::new();
206        let mut reg = PluginRegistrar::new();
207        for factory in self.factories() {
208            let plugin = factory.build(&ctx)?;
209            reg.registering_plugin_id = Some(plugin.id().to_string());
210            plugin.register(&mut reg)?;
211            reg.registering_plugin_id = None;
212            plugins.push(plugin);
213        }
214        let mut contributions = reg.contributions;
215        let protocol_session = contributions.protocol_session.take().ok_or_else(|| {
216            PluginError::Registration("missing protocol session capability".to_string())
217        })?;
218        let protocol_driver = contributions.protocol_driver.take().ok_or_else(|| {
219            PluginError::Registration("missing protocol driver capability".to_string())
220        })?;
221        contributions.protocol_session = Some(protocol_session);
222        contributions.protocol_driver = Some(protocol_driver);
223        contributions
224            .turn_context_transforms
225            .sort_by_key(|entry| std::cmp::Reverse(entry.0));
226        contributions
227            .context_compactors
228            .sort_by_key(|entry| std::cmp::Reverse(entry.0));
229        let triggers = crate::TriggerEventCatalog::from_events(contributions.triggers.clone())
230            .map_err(|message| {
231                PluginError::Registration(format!("invalid trigger event catalog: {message}"))
232            })?;
233        let registry = match tool_snapshot {
234            Some(snapshot) => Arc::new(
235                crate::ToolRegistry::from_tool_providers(contributions.tool_providers.clone())
236                    .map_err(|err| {
237                        PluginError::Registration(format!("failed to build tool registry: {err}"))
238                    })?
239                    .fork_with_state(snapshot)
240                    .map_err(|err| {
241                        PluginError::Session(format!(
242                            "tool state cannot be applied to this plugin host session: {err}"
243                        ))
244                    })?,
245            ),
246            None => Arc::new(
247                crate::ToolRegistry::from_tool_providers(contributions.tool_providers.clone())
248                    .map_err(|err| {
249                        PluginError::Registration(format!("failed to build tool registry: {err}"))
250                    })?,
251            ),
252        };
253        let tools = Arc::clone(&registry) as Arc<dyn ToolProvider>;
254
255        let session = Arc::new(PluginSession {
256            host: self.clone(),
257            session_id: ctx.session_id,
258            plugins,
259            tools,
260            tool_registry: registry,
261            tool_catalog_overlay,
262            tool_access: authority.tool_access,
263            subagent: authority.subagent,
264            extensions: self.extensions.clone(),
265            triggers,
266            contributions,
267        });
268        self.register_session(&session_id, &session)?;
269        let ready = SessionReadyContext {
270            session_id: session.session_id.clone(),
271            host: self.clone(),
272        };
273        for plugin in &session.plugins {
274            plugin.session_ready(ready.clone())?;
275        }
276        if let Some(snapshot) = snapshot {
277            session.restore(snapshot)?;
278        }
279        Ok(session)
280    }
281
282    fn register_session(
283        &self,
284        session_id: &str,
285        session: &Arc<PluginSession>,
286    ) -> Result<(), PluginError> {
287        let mut sessions = self.sessions.lock().map_err(|_| {
288            PluginError::Session("plugin host session registry poisoned".to_string())
289        })?;
290        if let Some(existing) = sessions.get(session_id).and_then(Weak::upgrade) {
291            if !Arc::ptr_eq(&existing, session) {
292                return Err(PluginError::Session(format!(
293                    "session `{session_id}` is already registered on this plugin host"
294                )));
295            }
296            return Ok(());
297        }
298        sessions.insert(session_id.to_string(), Arc::downgrade(session));
299        Ok(())
300    }
301
302    pub fn unregister_session(&self, session_id: &str) -> Result<(), PluginError> {
303        let mut sessions = self.sessions.lock().map_err(|_| {
304            PluginError::Session("plugin host session registry poisoned".to_string())
305        })?;
306        sessions.remove(session_id);
307        Ok(())
308    }
309
310    pub fn session(
311        &self,
312        session_id: &str,
313    ) -> Result<Arc<PluginSession>, PluginOperationInvokeError> {
314        let mut sessions = self
315            .sessions
316            .lock()
317            .map_err(|_| PluginOperationInvokeError::SessionRegistryPoisoned)?;
318        let Some(weak) = sessions.get(session_id).cloned() else {
319            return Err(PluginOperationInvokeError::UnknownSession(
320                session_id.to_string(),
321            ));
322        };
323        match weak.upgrade() {
324            Some(session) => Ok(session),
325            None => {
326                sessions.remove(session_id);
327                Err(PluginOperationInvokeError::UnknownSession(
328                    session_id.to_string(),
329                ))
330            }
331        }
332    }
333}