Skip to main content

trustformers_wasm/
plugin_framework.rs

1// Plugin Framework for TrustformeRS WASM
2// Enables community extensions and custom functionality
3
4#![allow(clippy::missing_enforced_import_renames)]
5
6use serde::{Deserialize, Serialize};
7use std::collections::HashMap;
8use std::sync::{Arc, Mutex};
9use wasm_bindgen::prelude::*;
10
11/// Plugin trait that all plugins must implement
12pub trait Plugin: Send + Sync {
13    /// Plugin metadata
14    fn metadata(&self) -> PluginMetadata;
15
16    /// Initialize the plugin
17    fn initialize(&mut self, config: PluginConfig) -> Result<(), PluginError>;
18
19    /// Execute plugin functionality
20    fn execute(&self, context: &PluginContext) -> Result<PluginResult, PluginError>;
21
22    /// Cleanup resources when plugin is unloaded
23    fn cleanup(&mut self);
24}
25
26/// Plugin metadata information
27#[derive(Debug, Clone, Serialize, Deserialize)]
28pub struct PluginMetadata {
29    pub name: String,
30    pub version: String,
31    pub author: String,
32    pub description: String,
33    pub plugin_type: PluginType,
34    pub dependencies: Vec<String>,
35    pub permissions: Vec<PluginPermission>,
36}
37
38/// Types of plugins supported
39#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
40pub enum PluginType {
41    /// Preprocessing plugins (tokenization, data processing)
42    Preprocessor,
43    /// Model inference plugins (custom model formats)
44    InferenceEngine,
45    /// Postprocessing plugins (output formatting, analysis)
46    Postprocessor,
47    /// Optimization plugins (quantization, pruning)
48    Optimizer,
49    /// Visualization plugins (charts, graphs)
50    Visualizer,
51    /// Storage plugins (custom backends)
52    Storage,
53    /// Network plugins (custom protocols)
54    Network,
55    /// Utility plugins (general purpose)
56    Utility,
57}
58
59/// Plugin permissions system
60#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
61pub enum PluginPermission {
62    /// Access to read model data
63    ReadModelData,
64    /// Access to modify model data
65    WriteModelData,
66    /// Access to network operations
67    NetworkAccess,
68    /// Access to local storage
69    StorageAccess,
70    /// Access to GPU resources
71    GpuAccess,
72    /// Access to performance profiling
73    ProfilingAccess,
74    /// Access to debug information
75    DebugAccess,
76    /// Access to user interface
77    UiAccess,
78}
79
80/// Plugin configuration
81#[derive(Debug, Clone, Serialize, Deserialize)]
82pub struct PluginConfig {
83    pub settings: HashMap<String, String>,
84    pub enabled_features: Vec<String>,
85    pub resource_limits: ResourceLimits,
86}
87
88/// Resource limits for plugins
89#[derive(Debug, Clone, Serialize, Deserialize)]
90pub struct ResourceLimits {
91    pub max_memory_mb: Option<usize>,
92    pub max_execution_time_ms: Option<u64>,
93    pub max_network_requests: Option<usize>,
94    pub max_gpu_memory_mb: Option<usize>,
95}
96
97/// Plugin execution context
98#[derive(Debug)]
99pub struct PluginContext {
100    pub plugin_id: String,
101    pub session_id: String,
102    pub request_data: HashMap<String, String>,
103    pub model_metadata: Option<ModelMetadata>,
104    pub performance_budget: PerformanceBudget,
105}
106
107/// Serializable version of PluginContext for JavaScript interop
108#[derive(Debug, Clone, Serialize, Deserialize)]
109pub struct SerializablePluginContext {
110    pub plugin_id: String,
111    pub session_id: String,
112    pub request_data: HashMap<String, String>,
113    pub model_metadata: Option<SerializableModelMetadata>,
114    pub performance_budget: SerializablePerformanceBudget,
115}
116
117/// Model metadata for plugin context
118#[derive(Debug, Clone)]
119pub struct ModelMetadata {
120    pub model_type: String,
121    pub size_mb: f64,
122    pub architecture: String,
123    pub precision: String,
124}
125
126/// Serializable version of ModelMetadata
127#[derive(Debug, Clone, Serialize, Deserialize)]
128pub struct SerializableModelMetadata {
129    pub model_type: String,
130    pub size_mb: f64,
131    pub architecture: String,
132    pub precision: String,
133}
134
135/// Performance budget for plugin execution
136#[derive(Debug, Clone)]
137pub struct PerformanceBudget {
138    pub max_latency_ms: u64,
139    pub max_memory_mb: usize,
140    pub priority: ExecutionPriority,
141}
142
143/// Serializable version of PerformanceBudget
144#[derive(Debug, Clone, Serialize, Deserialize)]
145pub struct SerializablePerformanceBudget {
146    pub max_latency_ms: u64,
147    pub max_memory_mb: usize,
148    pub priority: SerializableExecutionPriority,
149}
150
151/// Execution priority levels
152#[derive(Debug, Clone, PartialEq)]
153pub enum ExecutionPriority {
154    Critical,
155    High,
156    Normal,
157    Low,
158    Background,
159}
160
161/// Serializable version of ExecutionPriority
162#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
163pub enum SerializableExecutionPriority {
164    Critical,
165    High,
166    Normal,
167    Low,
168    Background,
169}
170
171// Conversion implementations
172impl From<ExecutionPriority> for SerializableExecutionPriority {
173    fn from(priority: ExecutionPriority) -> Self {
174        match priority {
175            ExecutionPriority::Critical => SerializableExecutionPriority::Critical,
176            ExecutionPriority::High => SerializableExecutionPriority::High,
177            ExecutionPriority::Normal => SerializableExecutionPriority::Normal,
178            ExecutionPriority::Low => SerializableExecutionPriority::Low,
179            ExecutionPriority::Background => SerializableExecutionPriority::Background,
180        }
181    }
182}
183
184impl From<ModelMetadata> for SerializableModelMetadata {
185    fn from(metadata: ModelMetadata) -> Self {
186        SerializableModelMetadata {
187            model_type: metadata.model_type,
188            size_mb: metadata.size_mb,
189            architecture: metadata.architecture,
190            precision: metadata.precision,
191        }
192    }
193}
194
195impl From<PerformanceBudget> for SerializablePerformanceBudget {
196    fn from(budget: PerformanceBudget) -> Self {
197        SerializablePerformanceBudget {
198            max_latency_ms: budget.max_latency_ms,
199            max_memory_mb: budget.max_memory_mb,
200            priority: budget.priority.into(),
201        }
202    }
203}
204
205impl From<PluginContext> for SerializablePluginContext {
206    fn from(context: PluginContext) -> Self {
207        SerializablePluginContext {
208            plugin_id: context.plugin_id,
209            session_id: context.session_id,
210            request_data: context.request_data,
211            model_metadata: context.model_metadata.map(|m| m.into()),
212            performance_budget: context.performance_budget.into(),
213        }
214    }
215}
216
217/// Plugin execution result
218#[derive(Debug, Clone, Serialize, Deserialize)]
219pub struct PluginResult {
220    pub success: bool,
221    pub data: HashMap<String, String>,
222    pub metrics: ExecutionMetrics,
223    pub messages: Vec<String>,
224}
225
226/// Execution metrics
227#[derive(Debug, Clone, Serialize, Deserialize)]
228pub struct ExecutionMetrics {
229    pub execution_time_ms: f64,
230    pub memory_used_mb: f64,
231    pub cpu_usage_percent: f64,
232    pub gpu_memory_used_mb: Option<f64>,
233}
234
235/// Plugin errors
236#[derive(Debug, Clone, Serialize, Deserialize)]
237pub struct PluginError {
238    pub code: PluginErrorCode,
239    pub message: String,
240    pub details: Option<String>,
241}
242
243/// Plugin error codes
244#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
245pub enum PluginErrorCode {
246    InitializationFailed,
247    ExecutionFailed,
248    PermissionDenied,
249    ResourceExhausted,
250    InvalidConfiguration,
251    DependencyMissing,
252    UnsupportedOperation,
253    Internal,
254}
255
256/// Plugin registry for managing loaded plugins
257pub struct PluginRegistry {
258    plugins: Arc<Mutex<HashMap<String, Box<dyn Plugin>>>>,
259    enabled_plugins: Arc<Mutex<Vec<String>>>,
260    plugin_configs: Arc<Mutex<HashMap<String, PluginConfig>>>,
261}
262
263impl Default for PluginRegistry {
264    fn default() -> Self {
265        Self::new()
266    }
267}
268
269impl PluginRegistry {
270    /// Create a new plugin registry
271    pub fn new() -> Self {
272        Self {
273            plugins: Arc::new(Mutex::new(HashMap::new())),
274            enabled_plugins: Arc::new(Mutex::new(Vec::new())),
275            plugin_configs: Arc::new(Mutex::new(HashMap::new())),
276        }
277    }
278
279    /// Register a new plugin
280    pub fn register_plugin(
281        &self,
282        plugin_id: String,
283        plugin: Box<dyn Plugin>,
284        config: PluginConfig,
285    ) -> Result<(), PluginError> {
286        // Validate plugin metadata
287        let metadata = plugin.metadata();
288        self.validate_plugin_metadata(&metadata)?;
289
290        // Check dependencies
291        self.check_plugin_dependencies(&metadata.dependencies)?;
292
293        // Store plugin and configuration
294        {
295            let mut plugins = self.plugins.lock().expect("lock should not be poisoned");
296            plugins.insert(plugin_id.clone(), plugin);
297        }
298
299        {
300            let mut configs = self.plugin_configs.lock().expect("lock should not be poisoned");
301            configs.insert(plugin_id, config);
302        }
303
304        Ok(())
305    }
306
307    /// Enable a plugin
308    pub fn enable_plugin(&self, plugin_id: &str) -> Result<(), PluginError> {
309        // Check if plugin exists
310        {
311            let plugins = self.plugins.lock().expect("lock should not be poisoned");
312            if !plugins.contains_key(plugin_id) {
313                return Err(PluginError {
314                    code: PluginErrorCode::DependencyMissing,
315                    message: format!("Plugin '{plugin_id}' not found"),
316                    details: None,
317                });
318            }
319        }
320
321        // Initialize plugin
322        {
323            let mut plugins = self.plugins.lock().expect("lock should not be poisoned");
324            let configs = self.plugin_configs.lock().expect("lock should not be poisoned");
325
326            if let (Some(plugin), Some(config)) =
327                (plugins.get_mut(plugin_id), configs.get(plugin_id))
328            {
329                plugin.initialize(config.clone())?;
330            }
331        }
332
333        // Add to enabled list
334        {
335            let mut enabled = self.enabled_plugins.lock().expect("lock should not be poisoned");
336            if !enabled.contains(&plugin_id.to_string()) {
337                enabled.push(plugin_id.to_string());
338            }
339        }
340
341        Ok(())
342    }
343
344    /// Disable a plugin
345    pub fn disable_plugin(&self, plugin_id: &str) -> Result<(), PluginError> {
346        // Remove from enabled list
347        {
348            let mut enabled = self.enabled_plugins.lock().expect("lock should not be poisoned");
349            enabled.retain(|id| id != plugin_id);
350        }
351
352        // Cleanup plugin
353        {
354            let mut plugins = self.plugins.lock().expect("lock should not be poisoned");
355            if let Some(plugin) = plugins.get_mut(plugin_id) {
356                plugin.cleanup();
357            }
358        }
359
360        Ok(())
361    }
362
363    /// Execute plugins of a specific type
364    pub fn execute_plugins(
365        &self,
366        plugin_type: PluginType,
367        context: &PluginContext,
368    ) -> Vec<Result<PluginResult, PluginError>> {
369        let enabled = self.enabled_plugins.lock().expect("lock should not be poisoned").clone();
370        let plugins = self.plugins.lock().expect("lock should not be poisoned");
371
372        let mut results = Vec::new();
373
374        for plugin_id in enabled {
375            if let Some(plugin) = plugins.get(&plugin_id) {
376                let metadata = plugin.metadata();
377                if metadata.plugin_type == plugin_type {
378                    let result = plugin.execute(context);
379                    results.push(result);
380                }
381            }
382        }
383
384        results
385    }
386
387    /// Get enabled plugins
388    pub fn get_enabled_plugins(&self) -> Vec<String> {
389        self.enabled_plugins.lock().expect("lock should not be poisoned").clone()
390    }
391
392    /// Get plugin metadata
393    pub fn get_plugin_metadata(&self, plugin_id: &str) -> Option<PluginMetadata> {
394        let plugins = self.plugins.lock().expect("lock should not be poisoned");
395        plugins.get(plugin_id).map(|plugin| plugin.metadata())
396    }
397
398    /// List all registered plugins
399    pub fn list_plugins(&self) -> Vec<PluginMetadata> {
400        let plugins = self.plugins.lock().expect("lock should not be poisoned");
401        plugins.values().map(|plugin| plugin.metadata()).collect()
402    }
403
404    fn validate_plugin_metadata(&self, metadata: &PluginMetadata) -> Result<(), PluginError> {
405        if metadata.name.is_empty() {
406            return Err(PluginError {
407                code: PluginErrorCode::InvalidConfiguration,
408                message: "Plugin name cannot be empty".to_string(),
409                details: None,
410            });
411        }
412
413        if metadata.version.is_empty() {
414            return Err(PluginError {
415                code: PluginErrorCode::InvalidConfiguration,
416                message: "Plugin version cannot be empty".to_string(),
417                details: None,
418            });
419        }
420
421        Ok(())
422    }
423
424    fn check_plugin_dependencies(&self, dependencies: &[String]) -> Result<(), PluginError> {
425        let plugins = self.plugins.lock().expect("lock should not be poisoned");
426
427        for dep in dependencies {
428            if !plugins.contains_key(dep) {
429                return Err(PluginError {
430                    code: PluginErrorCode::DependencyMissing,
431                    message: format!("Missing dependency: {dep}"),
432                    details: None,
433                });
434            }
435        }
436
437        Ok(())
438    }
439}
440
441/// WASM bindings for plugin framework
442#[wasm_bindgen]
443pub struct PluginManager {
444    registry: PluginRegistry,
445}
446
447impl Default for PluginManager {
448    fn default() -> Self {
449        Self::new()
450    }
451}
452
453#[wasm_bindgen]
454impl PluginManager {
455    #[wasm_bindgen(constructor)]
456    pub fn new() -> Self {
457        Self {
458            registry: PluginRegistry::new(),
459        }
460    }
461
462    /// Get list of enabled plugins as JSON
463    #[wasm_bindgen(getter)]
464    pub fn enabled_plugins(&self) -> String {
465        serde_json::to_string(&self.registry.get_enabled_plugins()).unwrap_or_default()
466    }
467
468    /// Get list of all plugins as JSON
469    pub fn list_plugins(&self) -> String {
470        serde_json::to_string(&self.registry.list_plugins()).unwrap_or_default()
471    }
472
473    /// Enable a plugin by ID
474    pub fn enable_plugin(&self, plugin_id: &str) -> Result<(), JsValue> {
475        self.registry
476            .enable_plugin(plugin_id)
477            .map_err(|e| JsValue::from_str(&e.message))
478    }
479
480    /// Disable a plugin by ID
481    pub fn disable_plugin(&self, plugin_id: &str) -> Result<(), JsValue> {
482        self.registry
483            .disable_plugin(plugin_id)
484            .map_err(|e| JsValue::from_str(&e.message))
485    }
486
487    /// Get plugin metadata as JSON
488    pub fn get_plugin_metadata(&self, plugin_id: &str) -> Option<String> {
489        self.registry
490            .get_plugin_metadata(plugin_id)
491            .and_then(|metadata| serde_json::to_string(&metadata).ok())
492    }
493}
494
495/// Default resource limits for plugins
496impl Default for ResourceLimits {
497    fn default() -> Self {
498        Self {
499            max_memory_mb: Some(100),
500            max_execution_time_ms: Some(5000),
501            max_network_requests: Some(10),
502            max_gpu_memory_mb: Some(50),
503        }
504    }
505}
506
507/// Default plugin configuration
508impl Default for PluginConfig {
509    fn default() -> Self {
510        Self {
511            settings: HashMap::new(),
512            enabled_features: Vec::new(),
513            resource_limits: ResourceLimits::default(),
514        }
515    }
516}
517
518/// Utility functions for creating plugin configs
519#[wasm_bindgen]
520pub fn create_default_plugin_config() -> String {
521    serde_json::to_string(&PluginConfig::default()).unwrap_or_default()
522}
523
524#[wasm_bindgen]
525pub fn create_plugin_context(plugin_id: &str, session_id: &str, request_data: &str) -> String {
526    let request_map: HashMap<String, String> =
527        serde_json::from_str(request_data).unwrap_or_default();
528
529    let context = PluginContext {
530        plugin_id: plugin_id.to_string(),
531        session_id: session_id.to_string(),
532        request_data: request_map,
533        model_metadata: Some(ModelMetadata {
534            model_type: "transformer".to_string(),
535            size_mb: 125.5,
536            architecture: "gpt-2".to_string(),
537            precision: "fp16".to_string(),
538        }),
539        performance_budget: PerformanceBudget {
540            max_latency_ms: 1000,
541            max_memory_mb: 50,
542            priority: ExecutionPriority::Normal,
543        },
544    };
545
546    // Convert to serializable version and serialize to JSON
547    let serializable_context: SerializablePluginContext = context.into();
548    serde_json::to_string(&serializable_context).unwrap_or_else(|e| {
549        web_sys::console::log_1(&format!("Failed to serialize plugin context: {}", e).into());
550        format!(
551            r#"{{"plugin_id":"{}","session_id":"{}","error":"serialization_failed"}}"#,
552            plugin_id, session_id
553        )
554    })
555}