universal_bot_core/
plugin.rs

1//! Plugin system for extending bot functionality
2//!
3//! This module provides a plugin architecture that allows extending
4//! the bot's capabilities without modifying core code.
5
6use std::collections::HashMap;
7use std::fmt;
8
9use anyhow::Result;
10use async_trait::async_trait;
11use serde::{Deserialize, Serialize};
12use tracing::{debug, info, instrument, warn};
13
14use crate::{
15    error::Error,
16    message::{Message, Response},
17};
18
19/// Plugin trait for extending bot functionality
20#[async_trait]
21pub trait Plugin: Send + Sync {
22    /// Get the plugin name
23    fn name(&self) -> &str;
24
25    /// Get the plugin version
26    fn version(&self) -> &str;
27
28    /// Get plugin description
29    fn description(&self) -> &str {
30        "No description provided"
31    }
32
33    /// Get plugin capabilities
34    fn capabilities(&self) -> Vec<Capability>;
35
36    /// Initialize the plugin
37    async fn initialize(&mut self, _config: PluginConfig) -> Result<()> {
38        Ok(())
39    }
40
41    /// Process a request
42    async fn process(&self, request: PluginRequest) -> Result<PluginResponse>;
43
44    /// Shutdown the plugin
45    async fn shutdown(&mut self) -> Result<()> {
46        Ok(())
47    }
48
49    /// Check if the plugin can handle a message
50    fn can_handle(&self, _message: &Message) -> bool {
51        true
52    }
53
54    /// Get plugin metadata
55    fn metadata(&self) -> PluginMetadata {
56        PluginMetadata {
57            name: self.name().to_string(),
58            version: self.version().to_string(),
59            description: self.description().to_string(),
60            author: None,
61            homepage: None,
62            license: None,
63        }
64    }
65}
66
67/// Plugin capability
68#[derive(Debug, Clone, Serialize, Deserialize)]
69pub struct Capability {
70    /// Capability name
71    pub name: String,
72    /// Capability type
73    pub capability_type: CapabilityType,
74    /// Description
75    pub description: String,
76    /// Required permissions
77    pub required_permissions: Vec<Permission>,
78}
79
80/// Type of capability
81#[derive(Debug, Clone, Serialize, Deserialize)]
82#[serde(rename_all = "snake_case")]
83pub enum CapabilityType {
84    /// Message processing
85    MessageProcessor,
86    /// Command handler
87    CommandHandler,
88    /// Event listener
89    EventListener,
90    /// Tool provider
91    ToolProvider,
92    /// Middleware
93    Middleware,
94    /// Custom capability
95    Custom(String),
96}
97
98/// Plugin permission
99#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
100#[serde(rename_all = "snake_case")]
101pub enum Permission {
102    /// Read messages
103    ReadMessages,
104    /// Write messages
105    WriteMessages,
106    /// Access context
107    AccessContext,
108    /// Modify context
109    ModifyContext,
110    /// Make network requests
111    NetworkAccess,
112    /// Access filesystem
113    FileSystemAccess,
114    /// Execute commands
115    ExecuteCommands,
116    /// Access database
117    DatabaseAccess,
118    /// All permissions
119    All,
120}
121
122/// Plugin configuration
123#[derive(Debug, Clone, Default, Serialize, Deserialize)]
124pub struct PluginConfig {
125    /// Plugin-specific settings
126    pub settings: HashMap<String, serde_json::Value>,
127    /// Enabled features
128    pub enabled_features: Vec<String>,
129    /// Granted permissions
130    pub permissions: Vec<Permission>,
131    /// Resource limits
132    pub resource_limits: ResourceLimits,
133}
134
135/// Resource limits for plugins
136#[derive(Debug, Clone, Serialize, Deserialize)]
137pub struct ResourceLimits {
138    /// Maximum memory in bytes
139    pub max_memory: Option<usize>,
140    /// Maximum CPU percentage
141    pub max_cpu: Option<f32>,
142    /// Maximum execution time
143    pub max_execution_time: Option<std::time::Duration>,
144    /// Maximum concurrent operations
145    pub max_concurrent_ops: Option<usize>,
146}
147
148impl Default for ResourceLimits {
149    fn default() -> Self {
150        Self {
151            max_memory: Some(100 * 1024 * 1024), // 100MB
152            max_cpu: Some(50.0),                 // 50%
153            max_execution_time: Some(std::time::Duration::from_secs(30)),
154            max_concurrent_ops: Some(10),
155        }
156    }
157}
158
159/// Plugin request
160#[derive(Debug, Clone, Serialize, Deserialize)]
161pub struct PluginRequest {
162    /// Request ID
163    pub id: String,
164    /// Request type
165    pub request_type: RequestType,
166    /// Request data
167    pub data: serde_json::Value,
168    /// Request metadata
169    pub metadata: HashMap<String, serde_json::Value>,
170}
171
172/// Request type
173#[derive(Debug, Clone, Serialize, Deserialize)]
174#[serde(rename_all = "snake_case")]
175pub enum RequestType {
176    /// Process a message
177    ProcessMessage,
178    /// Execute a command
179    ExecuteCommand,
180    /// Handle an event
181    HandleEvent,
182    /// Invoke a tool
183    InvokeTool,
184    /// Custom request
185    Custom(String),
186}
187
188/// Plugin response
189#[derive(Debug, Clone, Serialize, Deserialize)]
190pub struct PluginResponse {
191    /// Response ID
192    pub id: String,
193    /// Success status
194    pub success: bool,
195    /// Response data
196    pub data: serde_json::Value,
197    /// Error message if failed
198    pub error: Option<String>,
199    /// Response metadata
200    pub metadata: HashMap<String, serde_json::Value>,
201}
202
203impl PluginResponse {
204    /// Create a successful response
205    pub fn success(id: impl Into<String>, data: serde_json::Value) -> Self {
206        Self {
207            id: id.into(),
208            success: true,
209            data,
210            error: None,
211            metadata: HashMap::new(),
212        }
213    }
214
215    /// Create an error response
216    pub fn error(id: impl Into<String>, error: impl fmt::Display) -> Self {
217        Self {
218            id: id.into(),
219            success: false,
220            data: serde_json::Value::Null,
221            error: Some(error.to_string()),
222            metadata: HashMap::new(),
223        }
224    }
225}
226
227/// Plugin metadata
228#[derive(Debug, Clone, Serialize, Deserialize)]
229pub struct PluginMetadata {
230    /// Plugin name
231    pub name: String,
232    /// Plugin version
233    pub version: String,
234    /// Plugin description
235    pub description: String,
236    /// Plugin author
237    pub author: Option<String>,
238    /// Plugin homepage
239    pub homepage: Option<String>,
240    /// Plugin license
241    pub license: Option<String>,
242}
243
244/// Plugin registry for managing plugins
245pub struct PluginRegistry {
246    plugins: HashMap<String, Box<dyn Plugin>>,
247    hooks: HashMap<HookType, Vec<String>>,
248    permissions: HashMap<String, Vec<Permission>>,
249}
250
251impl PluginRegistry {
252    /// Create a new plugin registry
253    #[must_use]
254    pub fn new() -> Self {
255        Self {
256            plugins: HashMap::new(),
257            hooks: HashMap::new(),
258            permissions: HashMap::new(),
259        }
260    }
261
262    /// Register a plugin
263    ///
264    /// # Errors
265    ///
266    /// Returns an error if a plugin with the same name already exists.
267    #[instrument(skip(self, plugin))]
268    pub fn register(&mut self, mut plugin: Box<dyn Plugin>) -> Result<()> {
269        let name = plugin.name().to_string();
270
271        if self.plugins.contains_key(&name) {
272            return Err(Error::Plugin(format!("Plugin '{name}' already registered")).into());
273        }
274
275        info!("Registering plugin: {} v{}", name, plugin.version());
276
277        // Initialize plugin with default config
278        let config = PluginConfig::default();
279        futures::executor::block_on(plugin.initialize(config))?;
280
281        // Register capabilities
282        for capability in plugin.capabilities() {
283            self.register_hook(&name, &capability);
284        }
285
286        self.plugins.insert(name.clone(), plugin);
287        self.permissions.insert(name, vec![Permission::All]);
288
289        Ok(())
290    }
291
292    /// Unregister a plugin
293    #[instrument(skip(self))]
294    pub async fn unregister(&mut self, name: &str) -> Result<()> {
295        if let Some(mut plugin) = self.plugins.remove(name) {
296            info!("Unregistering plugin: {}", name);
297            plugin.shutdown().await?;
298
299            // Remove from hooks
300            for hooks in self.hooks.values_mut() {
301                hooks.retain(|n| n != name);
302            }
303
304            self.permissions.remove(name);
305            Ok(())
306        } else {
307            Err(Error::NotFound(format!("Plugin '{name}' not found")).into())
308        }
309    }
310
311    /// Get a plugin by name
312    pub fn get(&self, name: &str) -> Option<&dyn Plugin> {
313        self.plugins.get(name).map(std::convert::AsRef::as_ref)
314    }
315
316    /// List all registered plugins
317    pub fn list(&self) -> Vec<PluginMetadata> {
318        self.plugins.values().map(|p| p.metadata()).collect()
319    }
320
321    /// Apply pre-processing plugins
322    #[instrument(skip(self, message))]
323    pub async fn apply_pre_processing(&self, mut message: Message) -> Result<Message> {
324        for plugin in self.plugins.values() {
325            if plugin.can_handle(&message) {
326                let request = PluginRequest {
327                    id: uuid::Uuid::new_v4().to_string(),
328                    request_type: RequestType::ProcessMessage,
329                    data: serde_json::to_value(&message)?,
330                    metadata: HashMap::new(),
331                };
332
333                match plugin.process(request).await {
334                    Ok(response) if response.success => {
335                        if let Ok(processed) = serde_json::from_value(response.data) {
336                            message = processed;
337                        }
338                    }
339                    Ok(response) => {
340                        warn!(
341                            "Plugin {} failed to process message: {:?}",
342                            plugin.name(),
343                            response.error
344                        );
345                    }
346                    Err(e) => {
347                        warn!("Plugin {} error: {}", plugin.name(), e);
348                    }
349                }
350            }
351        }
352
353        Ok(message)
354    }
355
356    /// Apply post-processing plugins
357    #[instrument(skip(self, response))]
358    pub async fn apply_post_processing(&self, mut response: Response) -> Result<Response> {
359        for plugin in self.plugins.values() {
360            let request = PluginRequest {
361                id: uuid::Uuid::new_v4().to_string(),
362                request_type: RequestType::Custom("post_process".to_string()),
363                data: serde_json::to_value(&response)?,
364                metadata: HashMap::new(),
365            };
366
367            match plugin.process(request).await {
368                Ok(plugin_response) if plugin_response.success => {
369                    if let Ok(processed) = serde_json::from_value(plugin_response.data) {
370                        response = processed;
371                    }
372                }
373                Ok(plugin_response) => {
374                    debug!(
375                        "Plugin {} post-processing failed: {:?}",
376                        plugin.name(),
377                        plugin_response.error
378                    );
379                }
380                Err(e) => {
381                    debug!("Plugin {} post-processing error: {}", plugin.name(), e);
382                }
383            }
384        }
385
386        Ok(response)
387    }
388
389    /// Check if a plugin has permission
390    pub fn has_permission(&self, plugin_name: &str, permission: &Permission) -> bool {
391        self.permissions
392            .get(plugin_name)
393            .is_some_and(|perms| perms.contains(permission) || perms.contains(&Permission::All))
394    }
395
396    // Private helper methods
397
398    fn register_hook(&mut self, plugin_name: &str, capability: &Capability) {
399        let hook_type = match &capability.capability_type {
400            CapabilityType::MessageProcessor => HookType::MessageProcessor,
401            CapabilityType::CommandHandler => HookType::CommandHandler,
402            CapabilityType::EventListener => HookType::EventListener,
403            CapabilityType::ToolProvider => HookType::ToolProvider,
404            CapabilityType::Middleware => HookType::Middleware,
405            CapabilityType::Custom(name) => HookType::Custom(name.clone()),
406        };
407
408        self.hooks
409            .entry(hook_type)
410            .or_default()
411            .push(plugin_name.to_string());
412    }
413}
414
415impl Default for PluginRegistry {
416    fn default() -> Self {
417        Self::new()
418    }
419}
420
421/// Hook type for plugin registration
422#[derive(Debug, Clone, PartialEq, Eq, Hash)]
423pub enum HookType {
424    /// Message processor hook
425    MessageProcessor,
426    /// Command handler hook
427    CommandHandler,
428    /// Event listener hook
429    EventListener,
430    /// Tool provider hook
431    ToolProvider,
432    /// Middleware hook
433    Middleware,
434    /// Custom hook
435    Custom(String),
436}
437
438/// Example plugin implementation
439pub struct EchoPlugin {
440    name: String,
441    version: String,
442}
443
444impl EchoPlugin {
445    /// Create a new echo plugin
446    #[must_use]
447    pub fn new() -> Self {
448        Self {
449            name: "echo".to_string(),
450            version: "1.0.0".to_string(),
451        }
452    }
453}
454
455impl Default for EchoPlugin {
456    fn default() -> Self {
457        Self::new()
458    }
459}
460
461#[async_trait]
462impl Plugin for EchoPlugin {
463    fn name(&self) -> &str {
464        &self.name
465    }
466
467    fn version(&self) -> &str {
468        &self.version
469    }
470
471    fn description(&self) -> &'static str {
472        "Simple echo plugin for testing"
473    }
474
475    fn capabilities(&self) -> Vec<Capability> {
476        vec![Capability {
477            name: "echo".to_string(),
478            capability_type: CapabilityType::MessageProcessor,
479            description: "Echoes messages back".to_string(),
480            required_permissions: vec![Permission::ReadMessages, Permission::WriteMessages],
481        }]
482    }
483
484    async fn process(&self, request: PluginRequest) -> Result<PluginResponse> {
485        match request.request_type {
486            RequestType::ProcessMessage => {
487                if let Ok(message) = serde_json::from_value::<Message>(request.data) {
488                    let echo_message = Message::text(format!("Echo: {}", message.content));
489                    Ok(PluginResponse::success(
490                        request.id,
491                        serde_json::to_value(echo_message)?,
492                    ))
493                } else {
494                    Ok(PluginResponse::error(request.id, "Invalid message data"))
495                }
496            }
497            _ => Ok(PluginResponse::error(
498                request.id,
499                "Unsupported request type",
500            )),
501        }
502    }
503}
504
505#[cfg(test)]
506mod tests {
507    use super::*;
508
509    #[test]
510    fn test_plugin_registry() {
511        let mut registry = PluginRegistry::new();
512        let plugin = Box::new(EchoPlugin::new());
513
514        assert!(registry.register(plugin).is_ok());
515        assert!(registry.get("echo").is_some());
516
517        let plugins = registry.list();
518        assert_eq!(plugins.len(), 1);
519        assert_eq!(plugins[0].name, "echo");
520    }
521
522    #[tokio::test]
523    async fn test_plugin_unregister() {
524        let mut registry = PluginRegistry::new();
525        let plugin = Box::new(EchoPlugin::new());
526
527        registry.register(plugin).unwrap();
528        assert!(registry.unregister("echo").await.is_ok());
529        assert!(registry.get("echo").is_none());
530    }
531
532    #[test]
533    fn test_plugin_permissions() {
534        let mut registry = PluginRegistry::new();
535        let plugin = Box::new(EchoPlugin::new());
536
537        // Before registering, plugin shouldn't have permissions
538        assert!(!registry.has_permission("echo", &Permission::All));
539
540        // After registering, plugin gets all permissions by default
541        registry.register(plugin).unwrap();
542        assert!(registry.has_permission("echo", &Permission::All));
543        assert!(registry.has_permission("echo", &Permission::ReadMessages));
544
545        // Non-existent plugins don't have permissions
546        assert!(!registry.has_permission("nonexistent", &Permission::ReadMessages));
547    }
548
549    #[tokio::test]
550    async fn test_echo_plugin() {
551        let plugin = EchoPlugin::new();
552        let message = Message::text("Hello, world!");
553
554        let request = PluginRequest {
555            id: "test-123".to_string(),
556            request_type: RequestType::ProcessMessage,
557            data: serde_json::to_value(message).unwrap(),
558            metadata: HashMap::new(),
559        };
560
561        let response = plugin.process(request).await.unwrap();
562        assert!(response.success);
563
564        let echo_message: Message = serde_json::from_value(response.data).unwrap();
565        assert_eq!(echo_message.content, "Echo: Hello, world!");
566    }
567
568    #[test]
569    fn test_plugin_response() {
570        let response = PluginResponse::success("test", serde_json::json!({"key": "value"}));
571        assert!(response.success);
572        assert!(response.error.is_none());
573
574        let error_response = PluginResponse::error("test", "Something went wrong");
575        assert!(!error_response.success);
576        assert_eq!(
577            error_response.error.as_deref(),
578            Some("Something went wrong")
579        );
580    }
581}