af_agent/
plugin.rs

1use std::collections::{BTreeMap, BTreeSet};
2use std::future::Future;
3use std::pin::Pin;
4use std::sync::Arc;
5
6use async_trait::async_trait;
7use semver::{Version, VersionReq};
8use serde::{Deserialize, Serialize};
9use serde_json::Value;
10
11use crate::{
12    validate_json_schema, ContextContributor, PromptAuthority, PromptSection, Tool, ToolRegistry,
13};
14
15#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
16pub struct PluginManifest {
17    pub id: String,
18    pub version: String,
19    #[serde(default)]
20    pub dependencies: BTreeMap<String, String>,
21    #[serde(default)]
22    pub config_schema: Value,
23    #[serde(default)]
24    pub permissions: BTreeSet<PluginPermission>,
25}
26
27impl PluginManifest {
28    pub fn validate(&self) -> Result<(), PluginError> {
29        if self.id.trim().is_empty()
30            || !self.id.chars().all(|ch| {
31                ch.is_ascii_lowercase() || ch.is_ascii_digit() || matches!(ch, '.' | '-' | '_')
32            })
33        {
34            return Err(PluginError::InvalidManifest(format!(
35                "invalid plugin id {}",
36                self.id
37            )));
38        }
39        Version::parse(&self.version)
40            .map_err(|error| PluginError::InvalidManifest(format!("{}: {error}", self.id)))?;
41        Ok(())
42    }
43}
44
45#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
46#[serde(rename_all = "snake_case")]
47pub enum PluginPermission {
48    Prompt,
49    Context,
50    Tool,
51    Hook,
52}
53
54#[derive(Debug, Clone, PartialEq, Eq)]
55pub struct PluginMountContext {
56    pub tenant_id: String,
57    pub subject_id: String,
58    pub session_id: String,
59    pub profile_revision_id: String,
60}
61
62#[derive(Debug, Clone, PartialEq)]
63pub enum HookDecision {
64    Continue,
65    Deny { reason: String },
66    WaitForInput { kind: String, payload: Value },
67}
68
69pub type ToolExecutionFuture<'a> = Pin<Box<dyn Future<Output = Result<Value, String>> + Send + 'a>>;
70
71#[async_trait]
72pub trait Hook: Send + Sync {
73    async fn before_tool(
74        &self,
75        _context: &crate::ToolExecutionContext,
76        _tool: &str,
77        _arguments: &Value,
78    ) -> HookDecision {
79        HookDecision::Continue
80    }
81    fn around_tool<'a>(
82        &'a self,
83        _context: &'a crate::ToolExecutionContext,
84        _tool: &'a str,
85        _arguments: &'a Value,
86        next: ToolExecutionFuture<'a>,
87    ) -> ToolExecutionFuture<'a> {
88        next
89    }
90    async fn after_tool(
91        &self,
92        _context: &crate::ToolExecutionContext,
93        _tool: &str,
94        _arguments: &Value,
95        _result: &Value,
96    ) -> Result<(), PluginError> {
97        Ok(())
98    }
99}
100
101#[derive(Default, Clone)]
102struct PluginContributions {
103    prompts: BTreeMap<String, PromptSection>,
104    contexts: Vec<Arc<dyn ContextContributor>>,
105    tools: ToolRegistry,
106    hooks: Vec<Arc<dyn Hook>>,
107}
108
109impl PluginContributions {
110    fn merge(&mut self, additions: Self) -> Result<(), PluginError> {
111        let mut merged = self.clone();
112        for (id, prompt) in additions.prompts {
113            insert_once(&mut merged.prompts, id, prompt)?;
114        }
115        merged
116            .tools
117            .extend(&additions.tools)
118            .map_err(PluginError::Conflict)?;
119        merged.contexts.extend(additions.contexts);
120        merged.hooks.extend(additions.hooks);
121        *self = merged;
122        Ok(())
123    }
124}
125
126pub struct AgentRegistrar<'a> {
127    owner: &'a PluginManifest,
128    contributions: &'a mut PluginContributions,
129}
130
131impl AgentRegistrar<'_> {
132    pub fn prompt(
133        &mut self,
134        id: impl Into<String>,
135        prompt: impl Into<String>,
136    ) -> Result<(), PluginError> {
137        require(self.owner, PluginPermission::Prompt)?;
138        let id = format!("plugin.{}.{}", self.owner.id, id.into());
139        insert_once(
140            &mut self.contributions.prompts,
141            id.clone(),
142            PromptSection {
143                id,
144                order: 400,
145                authority: PromptAuthority::Plugin,
146                source: self.owner.id.clone(),
147                version: self.owner.version.clone(),
148                content: prompt.into(),
149            },
150        )
151    }
152    pub fn tool(&mut self, tool: Arc<dyn Tool>) -> Result<(), PluginError> {
153        require(self.owner, PluginPermission::Tool)?;
154        if self.contributions.tools.contains(tool.name()) {
155            return Err(PluginError::Conflict(format!("tool {}", tool.name())));
156        }
157        self.contributions
158            .tools
159            .register(tool)
160            .map(|_| ())
161            .map_err(PluginError::Conflict)
162    }
163    pub fn context(&mut self, contributor: Arc<dyn ContextContributor>) -> Result<(), PluginError> {
164        require(self.owner, PluginPermission::Context)?;
165        self.contributions.contexts.push(contributor);
166        Ok(())
167    }
168    pub fn hook(&mut self, hook: Arc<dyn Hook>) -> Result<(), PluginError> {
169        require(self.owner, PluginPermission::Hook)?;
170        self.contributions.hooks.push(hook);
171        Ok(())
172    }
173}
174
175#[async_trait]
176pub trait PluginLease: Send + Sync {
177    /// Cleanup must be idempotent because a failed release is retried.
178    async fn unmount(&mut self) -> Result<(), PluginError>;
179}
180
181#[async_trait]
182pub trait AgentPlugin: Send + Sync {
183    fn manifest(&self) -> &PluginManifest;
184
185    /// Creates the cleanup lease before activation performs any side effects.
186    /// The lease must be safe to unmount before, during, or after activation.
187    fn lease(&self, context: &PluginMountContext, config: &Value) -> Box<dyn PluginLease>;
188
189    async fn activate(
190        &self,
191        context: &PluginMountContext,
192        config: &Value,
193        registrar: &mut AgentRegistrar<'_>,
194    ) -> Result<(), PluginError>;
195}
196
197pub struct MountedPlugins {
198    contributions: PluginContributions,
199    leases: Arc<LeaseSet>,
200}
201
202impl std::fmt::Debug for MountedPlugins {
203    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
204        formatter
205            .debug_struct("MountedPlugins")
206            .field(
207                "pending_leases",
208                &self.leases.state.lock().unwrap().leases.len(),
209            )
210            .finish_non_exhaustive()
211    }
212}
213
214impl MountedPlugins {
215    fn from_parts(contributions: PluginContributions, leases: Vec<Box<dyn PluginLease>>) -> Self {
216        Self {
217            contributions,
218            leases: Arc::new(LeaseSet {
219                state: std::sync::Mutex::new(LeaseState {
220                    leases,
221                    holders: 1,
222                    closing: false,
223                }),
224                notify: tokio::sync::Notify::new(),
225            }),
226        }
227    }
228
229    pub fn prompts(&self) -> impl Iterator<Item = &PromptSection> {
230        self.contributions.prompts.values()
231    }
232    pub fn contexts(&self) -> &[Arc<dyn ContextContributor>] {
233        &self.contributions.contexts
234    }
235    pub fn tools(&self) -> &ToolRegistry {
236        &self.contributions.tools
237    }
238    pub fn hooks(&self) -> &[Arc<dyn Hook>] {
239        &self.contributions.hooks
240    }
241
242    pub async fn unmount(&mut self) -> Result<(), PluginError> {
243        let leases = Arc::clone(&self.leases);
244        {
245            leases.state.lock().unwrap().closing = true;
246        }
247        loop {
248            let notified = leases.notify.notified();
249            let pending = {
250                let mut state = leases.state.lock().unwrap();
251                (state.holders == 1).then(|| std::mem::take(&mut state.leases))
252            };
253            if let Some(mut pending) = pending {
254                let mut failed = Vec::new();
255                let mut errors = Vec::new();
256                while let Some(mut lease) = pending.pop() {
257                    if let Err(error) = lease.unmount().await {
258                        errors.push(error.to_string());
259                        failed.push(lease);
260                    }
261                }
262                leases.state.lock().unwrap().leases = failed;
263                return if errors.is_empty() {
264                    Ok(())
265                } else {
266                    Err(PluginError::Mount(format!(
267                        "plugin cleanup failed: {}",
268                        errors.join("; ")
269                    )))
270                };
271            }
272            notified.await;
273        }
274    }
275}
276
277impl Clone for MountedPlugins {
278    fn clone(&self) -> Self {
279        self.leases.state.lock().unwrap().holders += 1;
280        Self {
281            contributions: self.contributions.clone(),
282            leases: Arc::clone(&self.leases),
283        }
284    }
285}
286
287impl Drop for MountedPlugins {
288    fn drop(&mut self) {
289        let mut state = self.leases.state.lock().unwrap();
290        state.holders -= 1;
291        if state.closing && state.holders <= 1 {
292            self.leases.notify.notify_waiters();
293        }
294    }
295}
296
297struct LeaseSet {
298    state: std::sync::Mutex<LeaseState>,
299    notify: tokio::sync::Notify,
300}
301
302struct LeaseState {
303    leases: Vec<Box<dyn PluginLease>>,
304    holders: usize,
305    closing: bool,
306}
307
308#[derive(Default)]
309pub struct PluginCatalog {
310    plugins: BTreeMap<(String, String), Arc<dyn AgentPlugin>>,
311}
312
313impl std::fmt::Debug for PluginCatalog {
314    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
315        formatter
316            .debug_struct("PluginCatalog")
317            .field("count", &self.plugins.len())
318            .finish()
319    }
320}
321
322impl PluginCatalog {
323    pub fn register(&mut self, plugin: Arc<dyn AgentPlugin>) -> Result<(), PluginError> {
324        plugin.manifest().validate()?;
325        let id = plugin.manifest().id.clone();
326        let version = plugin.manifest().version.clone();
327        if self.plugins.contains_key(&(id.clone(), version.clone())) {
328            return Err(PluginError::Conflict(format!("{id}@{version}")));
329        }
330        self.plugins.insert((id, version), plugin);
331        Ok(())
332    }
333
334    pub async fn mount(
335        &self,
336        context: &PluginMountContext,
337        references: &BTreeMap<String, (String, Value)>,
338    ) -> Result<MountedPlugins, PluginError> {
339        let selected = self.select(references)?;
340        let configs = references
341            .iter()
342            .map(|(id, (_, config))| (id.clone(), config.clone()))
343            .collect();
344        let borrowed = selected
345            .iter()
346            .map(|plugin| plugin.as_ref())
347            .collect::<Vec<_>>();
348        mount_plugins(context, &borrowed, &configs).await
349    }
350
351    pub fn validate_references(
352        &self,
353        references: &BTreeMap<String, (String, Value)>,
354    ) -> Result<(), PluginError> {
355        self.select(references).map(|_| ())
356    }
357
358    fn select(
359        &self,
360        references: &BTreeMap<String, (String, Value)>,
361    ) -> Result<Vec<Arc<dyn AgentPlugin>>, PluginError> {
362        let mut selected = Vec::with_capacity(references.len());
363        let mut configs = BTreeMap::new();
364        for (id, (version, config)) in references {
365            let plugin = if let Some(plugin) = self.plugins.get(&(id.clone(), version.clone())) {
366                plugin
367            } else if let Some(((_, actual), _)) = self
368                .plugins
369                .iter()
370                .find(|((plugin_id, _), _)| plugin_id == id)
371            {
372                return Err(PluginError::DependencyVersion {
373                    plugin: id.clone(),
374                    dependency: id.clone(),
375                    required: version.clone(),
376                    actual: actual.clone(),
377                });
378            } else {
379                return Err(PluginError::MissingDependency(format!("{id}@{version}")));
380            };
381            selected.push(Arc::clone(plugin));
382            configs.insert(id.clone(), config.clone());
383        }
384        let borrowed = selected
385            .iter()
386            .map(|plugin| plugin.as_ref())
387            .collect::<Vec<_>>();
388        resolve_plugin_order(&borrowed)?;
389        for plugin in &selected {
390            let manifest = plugin.manifest();
391            validate_json_schema(&manifest.config_schema, &configs[&manifest.id])
392                .map_err(|error| PluginError::InvalidConfig(manifest.id.clone(), error))?;
393        }
394        Ok(selected)
395    }
396}
397
398/// Mount a complete Session plugin set atomically. Any validation or activation
399/// failure tears down current and already-mounted plugins in reverse order.
400pub async fn mount_plugins(
401    context: &PluginMountContext,
402    plugins: &[&dyn AgentPlugin],
403    configs: &BTreeMap<String, Value>,
404) -> Result<MountedPlugins, PluginError> {
405    let ordered = resolve_plugin_order(plugins)?;
406    let mut contributions = PluginContributions::default();
407    let mut leases = Vec::with_capacity(ordered.len());
408    for plugin in ordered {
409        let manifest = plugin.manifest();
410        let config = configs
411            .get(&manifest.id)
412            .cloned()
413            .unwrap_or_else(|| Value::Object(Default::default()));
414        if let Err(error) = validate_json_schema(&manifest.config_schema, &config) {
415            return Err(rollback_error(
416                PluginError::InvalidConfig(manifest.id.clone(), error),
417                leases,
418            )
419            .await);
420        }
421        let mut additions = PluginContributions::default();
422        let mut registrar = AgentRegistrar {
423            owner: manifest,
424            contributions: &mut additions,
425        };
426        leases.push(plugin.lease(context, &config));
427        if let Err(error) = plugin.activate(context, &config, &mut registrar).await {
428            return Err(rollback_error(error, leases).await);
429        }
430        if let Err(error) = contributions.merge(additions) {
431            return Err(rollback_error(error, leases).await);
432        }
433    }
434    Ok(MountedPlugins::from_parts(contributions, leases))
435}
436
437async fn rollback_error(cause: PluginError, leases: Vec<Box<dyn PluginLease>>) -> PluginError {
438    let mut cleanup = MountedPlugins::from_parts(PluginContributions::default(), leases);
439    match cleanup.unmount().await {
440        Ok(()) => cause,
441        Err(error) => PluginError::Rollback {
442            cause: Box::new(cause),
443            cleanup_error: error.to_string(),
444            cleanup: Box::new(cleanup),
445        },
446    }
447}
448
449pub fn resolve_plugin_order<'a>(
450    plugins: &'a [&'a dyn AgentPlugin],
451) -> Result<Vec<&'a dyn AgentPlugin>, PluginError> {
452    let by_id = plugins
453        .iter()
454        .map(|plugin| (plugin.manifest().id.as_str(), *plugin))
455        .collect::<BTreeMap<_, _>>();
456    if by_id.len() != plugins.len() {
457        return Err(PluginError::Conflict("duplicate plugin id".into()));
458    }
459    let mut visiting = BTreeSet::new();
460    let mut visited = BTreeSet::new();
461    let mut ordered = Vec::with_capacity(plugins.len());
462    fn visit<'a>(
463        id: &'a str,
464        by_id: &BTreeMap<&'a str, &'a dyn AgentPlugin>,
465        visiting: &mut BTreeSet<&'a str>,
466        visited: &mut BTreeSet<&'a str>,
467        ordered: &mut Vec<&'a dyn AgentPlugin>,
468    ) -> Result<(), PluginError> {
469        if visited.contains(id) {
470            return Ok(());
471        }
472        if !visiting.insert(id) {
473            return Err(PluginError::DependencyCycle(id.into()));
474        }
475        let plugin = by_id
476            .get(id)
477            .ok_or_else(|| PluginError::MissingDependency(id.into()))?;
478        plugin.manifest().validate()?;
479        for (dependency, requirement) in &plugin.manifest().dependencies {
480            if !by_id.contains_key(dependency.as_str()) {
481                return Err(PluginError::MissingDependency(format!(
482                    "{id} -> {dependency}"
483                )));
484            }
485            let required = VersionReq::parse(requirement).map_err(|error| {
486                PluginError::InvalidManifest(format!("{id} dependency {dependency}: {error}"))
487            })?;
488            let actual = Version::parse(&by_id[dependency.as_str()].manifest().version)
489                .map_err(|error| PluginError::InvalidManifest(error.to_string()))?;
490            if !required.matches(&actual) {
491                return Err(PluginError::DependencyVersion {
492                    plugin: id.into(),
493                    dependency: dependency.clone(),
494                    required: requirement.clone(),
495                    actual: actual.to_string(),
496                });
497            }
498            visit(dependency, by_id, visiting, visited, ordered)?;
499        }
500        visiting.remove(id);
501        visited.insert(id);
502        ordered.push(*plugin);
503        Ok(())
504    }
505    for id in by_id.keys() {
506        visit(id, &by_id, &mut visiting, &mut visited, &mut ordered)?;
507    }
508    Ok(ordered)
509}
510
511fn require(manifest: &PluginManifest, permission: PluginPermission) -> Result<(), PluginError> {
512    manifest
513        .permissions
514        .contains(&permission)
515        .then_some(())
516        .ok_or_else(|| PluginError::Permission(format!("{} lacks {permission:?}", manifest.id)))
517}
518
519fn insert_once<T>(map: &mut BTreeMap<String, T>, key: String, value: T) -> Result<(), PluginError> {
520    if map.contains_key(&key) {
521        return Err(PluginError::Conflict(key));
522    }
523    map.insert(key, value);
524    Ok(())
525}
526
527#[derive(Debug, thiserror::Error)]
528pub enum PluginError {
529    #[error("invalid plugin manifest: {0}")]
530    InvalidManifest(String),
531    #[error("missing plugin dependency: {0}")]
532    MissingDependency(String),
533    #[error("plugin dependency cycle at {0}")]
534    DependencyCycle(String),
535    #[error("plugin permission denied: {0}")]
536    Permission(String),
537    #[error("plugin registration conflict: {0}")]
538    Conflict(String),
539    #[error("plugin mount failed: {0}")]
540    Mount(String),
541    #[error("{cause}; plugin rollback failed: {cleanup_error}")]
542    Rollback {
543        cause: Box<PluginError>,
544        cleanup_error: String,
545        cleanup: Box<MountedPlugins>,
546    },
547    #[error("plugin {0} configuration is invalid: {1}")]
548    InvalidConfig(String, String),
549    #[error("plugin {plugin} requires {dependency} {required}, found {actual}")]
550    DependencyVersion {
551        plugin: String,
552        dependency: String,
553        required: String,
554        actual: String,
555    },
556}
557
558impl PluginError {
559    pub async fn retry_cleanup(&mut self) -> Result<(), PluginError> {
560        let Self::Rollback {
561            cleanup,
562            cleanup_error,
563            ..
564        } = self
565        else {
566            return Ok(());
567        };
568        match cleanup.unmount().await {
569            Ok(()) => Ok(()),
570            Err(error) => {
571                *cleanup_error = error.to_string();
572                Err(error)
573            }
574        }
575    }
576}
577
578#[cfg(test)]
579mod tests {
580    use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
581
582    use super::*;
583    use crate::ToolExecutionContext;
584
585    struct TestTool;
586    #[async_trait]
587    impl Tool for TestTool {
588        fn name(&self) -> &str {
589            "test.tool"
590        }
591        fn description(&self) -> &str {
592            "test"
593        }
594        fn parameters(&self) -> Value {
595            serde_json::json!({"type":"object"})
596        }
597        fn output_schema(&self) -> Value {
598            serde_json::json!({"type":"object"})
599        }
600        async fn call(&self, arguments: Value) -> Result<Value, String> {
601            Ok(arguments)
602        }
603    }
604
605    struct TestHook;
606    #[async_trait]
607    impl Hook for TestHook {}
608
609    struct Lease(Arc<AtomicUsize>);
610    #[async_trait]
611    impl PluginLease for Lease {
612        async fn unmount(&mut self) -> Result<(), PluginError> {
613            self.0.fetch_add(1, Ordering::SeqCst);
614            Ok(())
615        }
616    }
617
618    struct TestPlugin {
619        manifest: PluginManifest,
620        fail: bool,
621        unmounted: Arc<AtomicUsize>,
622    }
623
624    struct RegisteringPlugin {
625        manifest: PluginManifest,
626        unmounted: Arc<AtomicUsize>,
627    }
628
629    #[async_trait]
630    impl AgentPlugin for RegisteringPlugin {
631        fn manifest(&self) -> &PluginManifest {
632            &self.manifest
633        }
634        fn lease(&self, _: &PluginMountContext, _: &Value) -> Box<dyn PluginLease> {
635            Box::new(Lease(Arc::clone(&self.unmounted)))
636        }
637        async fn activate(
638            &self,
639            _: &PluginMountContext,
640            _: &Value,
641            registrar: &mut AgentRegistrar<'_>,
642        ) -> Result<(), PluginError> {
643            registrar.tool(Arc::new(TestTool))?;
644            Ok(())
645        }
646    }
647
648    struct RetryLease {
649        id: &'static str,
650        fail_once: Arc<AtomicBool>,
651        calls: Arc<std::sync::Mutex<Vec<&'static str>>>,
652    }
653
654    #[async_trait]
655    impl PluginLease for RetryLease {
656        async fn unmount(&mut self) -> Result<(), PluginError> {
657            self.calls.lock().unwrap().push(self.id);
658            if self.fail_once.swap(false, Ordering::SeqCst) {
659                Err(PluginError::Mount(format!("{} failed", self.id)))
660            } else {
661                Ok(())
662            }
663        }
664    }
665
666    struct CleanupPlugin {
667        manifest: PluginManifest,
668        lease_id: &'static str,
669        fail_activation: bool,
670        fail_once: Arc<AtomicBool>,
671        calls: Arc<std::sync::Mutex<Vec<&'static str>>>,
672    }
673
674    #[async_trait]
675    impl AgentPlugin for CleanupPlugin {
676        fn manifest(&self) -> &PluginManifest {
677            &self.manifest
678        }
679        fn lease(&self, _: &PluginMountContext, _: &Value) -> Box<dyn PluginLease> {
680            Box::new(RetryLease {
681                id: self.lease_id,
682                fail_once: Arc::clone(&self.fail_once),
683                calls: Arc::clone(&self.calls),
684            })
685        }
686        async fn activate(
687            &self,
688            _: &PluginMountContext,
689            _: &Value,
690            _: &mut AgentRegistrar<'_>,
691        ) -> Result<(), PluginError> {
692            if self.fail_activation {
693                Err(PluginError::Mount(self.manifest.id.clone()))
694            } else {
695                Ok(())
696            }
697        }
698    }
699    #[async_trait]
700    impl AgentPlugin for TestPlugin {
701        fn manifest(&self) -> &PluginManifest {
702            &self.manifest
703        }
704        fn lease(&self, _: &PluginMountContext, _: &Value) -> Box<dyn PluginLease> {
705            Box::new(Lease(Arc::clone(&self.unmounted)))
706        }
707        async fn activate(
708            &self,
709            _context: &PluginMountContext,
710            _config: &Value,
711            _registrar: &mut AgentRegistrar<'_>,
712        ) -> Result<(), PluginError> {
713            if self.fail {
714                Err(PluginError::Mount(self.manifest.id.clone()))
715            } else {
716                Ok(())
717            }
718        }
719    }
720
721    fn plugin(
722        id: &str,
723        dependencies: BTreeMap<String, String>,
724        fail: bool,
725        unmounted: Arc<AtomicUsize>,
726    ) -> TestPlugin {
727        TestPlugin {
728            manifest: PluginManifest {
729                id: id.into(),
730                version: "1.0.0".into(),
731                dependencies,
732                config_schema: serde_json::json!({"type":"object"}),
733                permissions: BTreeSet::new(),
734            },
735            fail,
736            unmounted,
737        }
738    }
739
740    #[tokio::test]
741    async fn activation_failure_rolls_back_current_and_prior_leases() {
742        let unmounted = Arc::new(AtomicUsize::new(0));
743        let first = plugin("first", BTreeMap::new(), false, Arc::clone(&unmounted));
744        let second = plugin(
745            "second",
746            BTreeMap::from([("first".into(), "^1".into())]),
747            true,
748            Arc::clone(&unmounted),
749        );
750        let context = PluginMountContext {
751            tenant_id: "t".into(),
752            subject_id: "s".into(),
753            session_id: "session".into(),
754            profile_revision_id: "profile".into(),
755        };
756        assert!(
757            mount_plugins(&context, &[&second, &first], &BTreeMap::new())
758                .await
759                .is_err()
760        );
761        assert_eq!(unmounted.load(Ordering::SeqCst), 2);
762    }
763
764    #[tokio::test]
765    async fn activation_failure_cleans_current_plugin_and_retries_failed_cleanup() {
766        let calls = Arc::new(std::sync::Mutex::new(Vec::new()));
767        let first = CleanupPlugin {
768            manifest: plugin(
769                "first",
770                BTreeMap::new(),
771                false,
772                Arc::new(AtomicUsize::new(0)),
773            )
774            .manifest,
775            lease_id: "first",
776            fail_activation: false,
777            fail_once: Arc::new(AtomicBool::new(false)),
778            calls: Arc::clone(&calls),
779        };
780        let second = CleanupPlugin {
781            manifest: plugin(
782                "second",
783                BTreeMap::from([("first".into(), "^1".into())]),
784                false,
785                Arc::new(AtomicUsize::new(0)),
786            )
787            .manifest,
788            lease_id: "second",
789            fail_activation: true,
790            fail_once: Arc::new(AtomicBool::new(true)),
791            calls: Arc::clone(&calls),
792        };
793        let context = PluginMountContext {
794            tenant_id: "t".into(),
795            subject_id: "s".into(),
796            session_id: "rollback".into(),
797            profile_revision_id: "profile".into(),
798        };
799
800        let mut error = mount_plugins(&context, &[&second, &first], &BTreeMap::new())
801            .await
802            .unwrap_err();
803        assert!(matches!(error, PluginError::Rollback { .. }));
804        assert_eq!(*calls.lock().unwrap(), vec!["second", "first"]);
805        error.retry_cleanup().await.unwrap();
806        assert_eq!(*calls.lock().unwrap(), vec!["second", "first", "second"]);
807    }
808
809    #[tokio::test]
810    async fn registration_conflict_rolls_back_current_and_prior_leases() {
811        let unmounted = Arc::new(AtomicUsize::new(0));
812        let manifest = |id: &str, dependencies| PluginManifest {
813            id: id.into(),
814            version: "1.0.0".into(),
815            dependencies,
816            config_schema: serde_json::json!({"type":"object"}),
817            permissions: BTreeSet::from([PluginPermission::Tool]),
818        };
819        let first = RegisteringPlugin {
820            manifest: manifest("first", BTreeMap::new()),
821            unmounted: Arc::clone(&unmounted),
822        };
823        let second = RegisteringPlugin {
824            manifest: manifest("second", BTreeMap::from([("first".into(), "^1".into())])),
825            unmounted: Arc::clone(&unmounted),
826        };
827        let context = PluginMountContext {
828            tenant_id: "t".into(),
829            subject_id: "s".into(),
830            session_id: "session".into(),
831            profile_revision_id: "profile".into(),
832        };
833
834        assert!(matches!(
835            mount_plugins(&context, &[&second, &first], &BTreeMap::new()).await,
836            Err(PluginError::Conflict(_))
837        ));
838        assert_eq!(unmounted.load(Ordering::SeqCst), 2);
839    }
840
841    #[tokio::test]
842    async fn session_lease_unmounts_after_last_runtime_clone() {
843        let unmounted = Arc::new(AtomicUsize::new(0));
844        let plugin = plugin("session", BTreeMap::new(), false, Arc::clone(&unmounted));
845        let context = PluginMountContext {
846            tenant_id: "t".into(),
847            subject_id: "s".into(),
848            session_id: "session".into(),
849            profile_revision_id: "profile".into(),
850        };
851        let mounted = mount_plugins(&context, &[&plugin], &BTreeMap::new())
852            .await
853            .unwrap();
854        let runtime_clone = mounted.clone();
855        let teardown = tokio::spawn(async move {
856            let mut mounted = mounted;
857            mounted.unmount().await
858        });
859        tokio::task::yield_now().await;
860        assert_eq!(unmounted.load(Ordering::SeqCst), 0);
861        drop(runtime_clone);
862        teardown.await.unwrap().unwrap();
863        assert_eq!(unmounted.load(Ordering::SeqCst), 1);
864    }
865
866    #[tokio::test]
867    async fn cleanup_attempts_every_lease_and_retries_only_failures() {
868        let calls = Arc::new(std::sync::Mutex::new(Vec::new()));
869        let first = CleanupPlugin {
870            manifest: plugin(
871                "first-cleanup",
872                BTreeMap::new(),
873                false,
874                Arc::new(AtomicUsize::new(0)),
875            )
876            .manifest,
877            lease_id: "first",
878            fail_activation: false,
879            fail_once: Arc::new(AtomicBool::new(false)),
880            calls: Arc::clone(&calls),
881        };
882        let second = CleanupPlugin {
883            manifest: plugin(
884                "second-cleanup",
885                BTreeMap::from([("first-cleanup".into(), "^1".into())]),
886                false,
887                Arc::new(AtomicUsize::new(0)),
888            )
889            .manifest,
890            lease_id: "second",
891            fail_activation: false,
892            fail_once: Arc::new(AtomicBool::new(true)),
893            calls: Arc::clone(&calls),
894        };
895        let context = PluginMountContext {
896            tenant_id: "t".into(),
897            subject_id: "s".into(),
898            session_id: "cleanup".into(),
899            profile_revision_id: "profile".into(),
900        };
901        let mut mounted = mount_plugins(&context, &[&second, &first], &BTreeMap::new())
902            .await
903            .unwrap();
904        assert!(mounted.unmount().await.is_err());
905        assert_eq!(*calls.lock().unwrap(), vec!["second", "first"]);
906        mounted.unmount().await.unwrap();
907        assert_eq!(*calls.lock().unwrap(), vec!["second", "first", "second"]);
908    }
909
910    #[test]
911    fn dependency_versions_and_cycles_fail_before_mount() {
912        let counter = Arc::new(AtomicUsize::new(0));
913        let base = plugin("base", BTreeMap::new(), false, Arc::clone(&counter));
914        let incompatible = plugin(
915            "consumer",
916            BTreeMap::from([("base".into(), "^2".into())]),
917            false,
918            Arc::clone(&counter),
919        );
920        assert!(matches!(
921            resolve_plugin_order(&[&base, &incompatible]),
922            Err(PluginError::DependencyVersion { .. })
923        ));
924        let left = plugin(
925            "left",
926            BTreeMap::from([("right".into(), "^1".into())]),
927            false,
928            Arc::clone(&counter),
929        );
930        let right = plugin(
931            "right",
932            BTreeMap::from([("left".into(), "^1".into())]),
933            false,
934            counter,
935        );
936        assert!(matches!(
937            resolve_plugin_order(&[&left, &right]),
938            Err(PluginError::DependencyCycle(_))
939        ));
940    }
941
942    #[tokio::test]
943    async fn registrar_and_catalog_enforce_manifest_permissions_atomically() {
944        let manifest = PluginManifest {
945            id: "test.plugin".into(),
946            version: "1.0.0".into(),
947            dependencies: BTreeMap::new(),
948            config_schema: serde_json::json!({"type":"object","required":["enabled"]}),
949            permissions: BTreeSet::from([
950                PluginPermission::Prompt,
951                PluginPermission::Tool,
952                PluginPermission::Hook,
953            ]),
954        };
955        manifest.validate().unwrap();
956        for invalid in [
957            PluginManifest {
958                id: "Bad Plugin".into(),
959                ..manifest.clone()
960            },
961            PluginManifest {
962                version: "latest".into(),
963                ..manifest.clone()
964            },
965        ] {
966            assert!(matches!(
967                invalid.validate(),
968                Err(PluginError::InvalidManifest(_))
969            ));
970        }
971
972        let mut contributions = PluginContributions::default();
973        {
974            let mut registrar = AgentRegistrar {
975                owner: &manifest,
976                contributions: &mut contributions,
977            };
978            registrar.prompt("system", "prompt").unwrap();
979            assert!(matches!(
980                registrar.prompt("system", "again"),
981                Err(PluginError::Conflict(_))
982            ));
983            registrar.tool(Arc::new(TestTool)).unwrap();
984            assert!(matches!(
985                registrar.tool(Arc::new(TestTool)),
986                Err(PluginError::Conflict(_))
987            ));
988            registrar.hook(Arc::new(TestHook)).unwrap();
989        }
990        let prompt = contributions.prompts.values().next().unwrap();
991        assert_eq!(prompt.id, "plugin.test.plugin.system");
992        assert_eq!(prompt.source, "test.plugin");
993        assert_eq!(prompt.version, "1.0.0");
994
995        let denied = PluginManifest {
996            permissions: BTreeSet::new(),
997            ..manifest.clone()
998        };
999        let mut denied_contributions = PluginContributions::default();
1000        let mut denied_registrar = AgentRegistrar {
1001            owner: &denied,
1002            contributions: &mut denied_contributions,
1003        };
1004        assert!(matches!(
1005            denied_registrar.prompt("x", "x"),
1006            Err(PluginError::Permission(_))
1007        ));
1008        let context = ToolExecutionContext {
1009            request: crate::RequestContext {
1010                tenant_id: "tenant".into(),
1011                subject_id: "subject".into(),
1012                roles: Default::default(),
1013                locale: "en".into(),
1014                request_id: "request".into(),
1015                entitlements: Default::default(),
1016            },
1017            session_id: "session".into(),
1018            run_id: "run".into(),
1019            step: 1,
1020            call_id: "call".into(),
1021            source_event_seq: 1,
1022            interaction_resolution: None,
1023            cancellation: crate::CancellationToken::default(),
1024            deadline: std::time::Instant::now() + std::time::Duration::from_secs(15),
1025        };
1026        assert_eq!(
1027            TestHook
1028                .before_tool(&context, "test.tool", &serde_json::json!({}))
1029                .await,
1030            HookDecision::Continue
1031        );
1032        TestHook
1033            .after_tool(
1034                &context,
1035                "test.tool",
1036                &serde_json::json!({}),
1037                &serde_json::json!({}),
1038            )
1039            .await
1040            .unwrap();
1041        let unmounted = Arc::new(AtomicUsize::new(0));
1042        let plugin = Arc::new(TestPlugin {
1043            manifest: manifest.clone(),
1044            fail: false,
1045            unmounted: Arc::clone(&unmounted),
1046        });
1047        let mut catalog = PluginCatalog::default();
1048        catalog.register(plugin.clone()).unwrap();
1049        assert!(matches!(
1050            catalog.register(plugin),
1051            Err(PluginError::Conflict(_))
1052        ));
1053        let mount_context = PluginMountContext {
1054            tenant_id: "t".into(),
1055            subject_id: "s".into(),
1056            session_id: "session".into(),
1057            profile_revision_id: "profile".into(),
1058        };
1059        assert!(matches!(
1060            catalog
1061                .mount(
1062                    &mount_context,
1063                    &BTreeMap::from([("missing".into(), ("1.0.0".into(), serde_json::json!({})))])
1064                )
1065                .await,
1066            Err(PluginError::MissingDependency(_))
1067        ));
1068        assert!(matches!(
1069            catalog
1070                .mount(
1071                    &mount_context,
1072                    &BTreeMap::from([(
1073                        "test.plugin".into(),
1074                        ("2.0.0".into(), serde_json::json!({}))
1075                    )])
1076                )
1077                .await,
1078            Err(PluginError::DependencyVersion { .. })
1079        ));
1080        assert!(matches!(
1081            catalog
1082                .mount(
1083                    &mount_context,
1084                    &BTreeMap::from([(
1085                        "test.plugin".into(),
1086                        ("1.0.0".into(), serde_json::json!({}))
1087                    )])
1088                )
1089                .await,
1090            Err(PluginError::InvalidConfig(_, _))
1091        ));
1092        let mut mounted = catalog
1093            .mount(
1094                &mount_context,
1095                &BTreeMap::from([(
1096                    "test.plugin".into(),
1097                    ("1.0.0".into(), serde_json::json!({"enabled":true})),
1098                )]),
1099            )
1100            .await
1101            .unwrap();
1102        mounted.unmount().await.unwrap();
1103        assert_eq!(unmounted.load(Ordering::SeqCst), 1);
1104    }
1105}
af_agent/plugin.rs

af_agent/
plugin.rs
