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 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 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(®istry) 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}