Skip to main content

trustformers_wasm/
plugin_framework.rs

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