tenant.openagents/omega

1use std::sync::Arc;
2
3use anyhow::{Result, bail};
4use extension::{ExtensionCapability, ExtensionManifest};
5use url::Url;
6
7pub struct CapabilityGranter {
8    granted_capabilities: Vec<ExtensionCapability>,
9    manifest: Arc<ExtensionManifest>,
10}
11
12impl CapabilityGranter {
13    pub fn new(
14        granted_capabilities: Vec<ExtensionCapability>,
15        manifest: Arc<ExtensionManifest>,
16    ) -> Self {
17        Self {
18            granted_capabilities,
19            manifest,
20        }
21    }
22
23    pub fn grant_exec(
24        &self,
25        desired_command: &str,
26        desired_args: &[impl AsRef<str> + std::fmt::Debug],
27    ) -> Result<()> {
28        self.manifest.allow_exec(desired_command, desired_args)?;
29
30        let is_allowed = self
31            .granted_capabilities
32            .iter()
33            .any(|capability| match capability {
34                ExtensionCapability::ProcessExec(capability) => {
35                    capability.allows(desired_command, desired_args)
36                }
37                _ => false,
38            });
39
40        if !is_allowed {
41            bail!(
42                "capability for process:exec {desired_command} {desired_args:?} is not granted by the extension host",
43            );
44        }
45
46        Ok(())
47    }
48
49    pub fn grant_download_file(&self, desired_url: &Url) -> Result<()> {
50        let is_allowed = self
51            .granted_capabilities
52            .iter()
53            .any(|capability| match capability {
54                ExtensionCapability::DownloadFile(capability) => capability.allows(desired_url),
55                _ => false,
56            });
57
58        if !is_allowed {
59            bail!(
60                "capability for download_file {desired_url} is not granted by the extension host",
61            );
62        }
63
64        Ok(())
65    }
66
67    pub fn grant_npm_install_package(&self, package_name: &str) -> Result<()> {
68        let is_allowed = self
69            .granted_capabilities
70            .iter()
71            .any(|capability| match capability {
72                ExtensionCapability::NpmInstallPackage(capability) => {
73                    capability.allows(package_name)
74                }
75                _ => false,
76            });
77
78        if !is_allowed {
79            bail!("capability for npm:install {package_name} is not granted by the extension host",);
80        }
81
82        Ok(())
83    }
84}
85
86#[cfg(test)]
87mod tests {
88    use std::collections::BTreeMap;
89
90    use extension::{ProcessExecCapability, SchemaVersion};
91
92    use super::*;
93
94    fn extension_manifest() -> ExtensionManifest {
95        ExtensionManifest {
96            id: "test".into(),
97            name: "Test".to_string(),
98            version: "1.0.0".into(),
99            schema_version: SchemaVersion::ZERO,
100            description: None,
101            repository: None,
102            authors: vec![],
103            lib: Default::default(),
104            themes: vec![],
105            icon_themes: vec![],
106            languages: vec![],
107            grammars: BTreeMap::default(),
108            language_servers: BTreeMap::default(),
109            context_servers: BTreeMap::default(),
110            slash_commands: BTreeMap::default(),
111            snippets: None,
112            capabilities: vec![],
113            debug_adapters: Default::default(),
114            debug_locators: Default::default(),
115            language_model_providers: BTreeMap::default(),
116        }
117    }
118
119    #[test]
120    fn test_grant_exec() {
121        let manifest = Arc::new(ExtensionManifest {
122            capabilities: vec![ExtensionCapability::ProcessExec(ProcessExecCapability {
123                command: "ls".to_string(),
124                args: vec!["-la".to_string()],
125            })],
126            ..extension_manifest()
127        });
128
129        // It returns an error when the extension host has no granted capabilities.
130        let granter = CapabilityGranter::new(Vec::new(), manifest.clone());
131        assert!(granter.grant_exec("ls", &["-la"]).is_err());
132
133        // It succeeds when the extension host has the exact capability.
134        let granter = CapabilityGranter::new(
135            vec![ExtensionCapability::ProcessExec(ProcessExecCapability {
136                command: "ls".to_string(),
137                args: vec!["-la".to_string()],
138            })],
139            manifest.clone(),
140        );
141        assert!(granter.grant_exec("ls", &["-la"]).is_ok());
142
143        // It succeeds when the extension host has a wildcard capability.
144        let granter = CapabilityGranter::new(
145            vec![ExtensionCapability::ProcessExec(ProcessExecCapability {
146                command: "*".to_string(),
147                args: vec!["**".to_string()],
148            })],
149            manifest,
150        );
151        assert!(granter.grant_exec("ls", &["-la"]).is_ok());
152    }
153}
154