turbomcp_client/plugins/
core.rs

1//! Core plugin system traits and types
2//!
3//! Defines the fundamental abstractions for the plugin system including the ClientPlugin trait,
4//! context objects, error types, and configuration structures.
5
6use async_trait::async_trait;
7use chrono::{DateTime, Utc};
8use serde::{Deserialize, Serialize};
9use serde_json::Value;
10use std::collections::HashMap;
11use std::fmt;
12use std::time::Duration;
13use thiserror::Error;
14use turbomcp_protocol::jsonrpc::JsonRpcRequest;
15
16// ============================================================================
17// ERROR TYPES
18// ============================================================================
19
20/// Errors that can occur during plugin operations
21#[derive(Error, Debug)]
22#[non_exhaustive]
23pub enum PluginError {
24    /// Plugin initialization failed
25    #[error("Plugin initialization failed: {message}")]
26    Initialization { message: String },
27
28    /// Plugin configuration is invalid
29    #[error("Invalid plugin configuration: {message}")]
30    Configuration { message: String },
31
32    /// Error during request processing
33    #[error("Request processing error: {message}")]
34    RequestProcessing { message: String },
35
36    /// Error during response processing
37    #[error("Response processing error: {message}")]
38    ResponseProcessing { message: String },
39
40    /// Error in custom method handler
41    #[error("Custom handler error: {message}")]
42    CustomHandler { message: String },
43
44    /// Plugin dependency not available
45    #[error("Plugin dependency '{dependency}' not available")]
46    DependencyNotAvailable { dependency: String },
47
48    /// Plugin version compatibility issue
49    #[error("Plugin version incompatibility: {message}")]
50    VersionIncompatible { message: String },
51
52    /// Resource access error
53    #[error("Resource access error: {resource} - {message}")]
54    ResourceAccess { resource: String, message: String },
55
56    /// External system error
57    #[error("External system error: {source}")]
58    External {
59        #[from]
60        source: Box<dyn std::error::Error + Send + Sync>,
61    },
62}
63
64impl PluginError {
65    /// Create an initialization error
66    pub fn initialization(message: impl Into<String>) -> Self {
67        Self::Initialization {
68            message: message.into(),
69        }
70    }
71
72    /// Create a configuration error
73    pub fn configuration(message: impl Into<String>) -> Self {
74        Self::Configuration {
75            message: message.into(),
76        }
77    }
78
79    /// Create a request processing error
80    pub fn request_processing(message: impl Into<String>) -> Self {
81        Self::RequestProcessing {
82            message: message.into(),
83        }
84    }
85
86    /// Create a response processing error
87    pub fn response_processing(message: impl Into<String>) -> Self {
88        Self::ResponseProcessing {
89            message: message.into(),
90        }
91    }
92
93    /// Create a custom handler error
94    pub fn custom_handler(message: impl Into<String>) -> Self {
95        Self::CustomHandler {
96            message: message.into(),
97        }
98    }
99
100    /// Create a dependency error
101    pub fn dependency_not_available(dependency: impl Into<String>) -> Self {
102        Self::DependencyNotAvailable {
103            dependency: dependency.into(),
104        }
105    }
106
107    /// Create a version incompatibility error
108    pub fn version_incompatible(message: impl Into<String>) -> Self {
109        Self::VersionIncompatible {
110            message: message.into(),
111        }
112    }
113
114    /// Create a resource access error
115    pub fn resource_access(resource: impl Into<String>, message: impl Into<String>) -> Self {
116        Self::ResourceAccess {
117            resource: resource.into(),
118            message: message.into(),
119        }
120    }
121}
122
123pub type PluginResult<T> = Result<T, PluginError>;
124
125// ============================================================================
126// CONTEXT TYPES
127// ============================================================================
128
129/// Context information available to plugins during initialization
130#[derive(Debug, Clone)]
131pub struct PluginContext {
132    /// Client information
133    pub client_name: String,
134    pub client_version: String,
135
136    /// Available capabilities
137    pub capabilities: HashMap<String, Value>,
138
139    /// Configuration values
140    pub config: HashMap<String, Value>,
141
142    /// Registered plugin names (for dependency checking)
143    pub available_plugins: Vec<String>,
144}
145
146impl PluginContext {
147    /// Create a new plugin context
148    pub fn new(
149        client_name: String,
150        client_version: String,
151        capabilities: HashMap<String, Value>,
152        config: HashMap<String, Value>,
153        available_plugins: Vec<String>,
154    ) -> Self {
155        Self {
156            client_name,
157            client_version,
158            capabilities,
159            config,
160            available_plugins,
161        }
162    }
163
164    /// Check if a capability is available
165    pub fn has_capability(&self, capability: &str) -> bool {
166        self.capabilities.contains_key(capability)
167    }
168
169    /// Get a configuration value
170    pub fn get_config(&self, key: &str) -> Option<&Value> {
171        self.config.get(key)
172    }
173
174    /// Check if a plugin dependency is available
175    pub fn has_plugin(&self, plugin_name: &str) -> bool {
176        self.available_plugins.contains(&plugin_name.to_string())
177    }
178}
179
180/// Context for request processing
181#[derive(Debug, Clone)]
182pub struct RequestContext {
183    /// The JSON-RPC request being processed
184    pub request: JsonRpcRequest,
185
186    /// Additional metadata (can be modified by plugins)
187    pub metadata: HashMap<String, Value>,
188
189    /// Request timestamp
190    pub timestamp: DateTime<Utc>,
191}
192
193impl RequestContext {
194    /// Create a new request context
195    pub fn new(request: JsonRpcRequest, metadata: HashMap<String, Value>) -> Self {
196        Self {
197            request,
198            metadata,
199            timestamp: Utc::now(),
200        }
201    }
202
203    /// Get the request method
204    pub fn method(&self) -> &str {
205        &self.request.method
206    }
207
208    /// Get request parameters
209    pub fn params(&self) -> Option<&Value> {
210        self.request.params.as_ref()
211    }
212
213    /// Add metadata
214    pub fn add_metadata(&mut self, key: String, value: Value) {
215        self.metadata.insert(key, value);
216    }
217
218    /// Get metadata value
219    pub fn get_metadata(&self, key: &str) -> Option<&Value> {
220        self.metadata.get(key)
221    }
222}
223
224/// Context for response processing
225#[derive(Debug, Clone)]
226pub struct ResponseContext {
227    /// The original request context
228    pub request_context: RequestContext,
229
230    /// The response data (if successful)
231    pub response: Option<Value>,
232
233    /// Error information (if failed)
234    pub error: Option<turbomcp_protocol::Error>,
235
236    /// Request duration
237    pub duration: Duration,
238
239    /// Additional metadata (can be modified by plugins)
240    pub metadata: HashMap<String, Value>,
241}
242
243impl ResponseContext {
244    /// Create a new response context
245    pub fn new(
246        request_context: RequestContext,
247        response: Option<Value>,
248        error: Option<turbomcp_protocol::Error>,
249        duration: Duration,
250    ) -> Self {
251        Self {
252            request_context,
253            response,
254            error,
255            duration,
256            metadata: HashMap::new(),
257        }
258    }
259
260    /// Check if the response was successful
261    pub fn is_success(&self) -> bool {
262        self.error.is_none()
263    }
264
265    /// Check if the response was an error
266    pub fn is_error(&self) -> bool {
267        self.error.is_some()
268    }
269
270    /// Get the request method
271    pub fn method(&self) -> &str {
272        self.request_context.method()
273    }
274
275    /// Add metadata
276    pub fn add_metadata(&mut self, key: String, value: Value) {
277        self.metadata.insert(key, value);
278    }
279
280    /// Get metadata value
281    pub fn get_metadata(&self, key: &str) -> Option<&Value> {
282        self.metadata.get(key)
283    }
284}
285
286// ============================================================================
287// PLUGIN CONFIGURATION
288// ============================================================================
289
290/// Plugin configuration variants
291#[derive(Debug, Clone, Serialize, Deserialize)]
292#[serde(tag = "type")]
293pub enum PluginConfig {
294    /// Metrics plugin configuration
295    #[serde(rename = "metrics")]
296    Metrics,
297
298    /// Retry plugin configuration
299    #[serde(rename = "retry")]
300    Retry(super::examples::RetryConfig),
301
302    /// Cache plugin configuration
303    #[serde(rename = "cache")]
304    Cache(super::examples::CacheConfig),
305
306    /// Custom plugin configuration
307    #[serde(rename = "custom")]
308    Custom {
309        name: String,
310        config: HashMap<String, Value>,
311    },
312}
313
314impl fmt::Display for PluginConfig {
315    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
316        match self {
317            PluginConfig::Metrics => write!(f, "Metrics"),
318            PluginConfig::Retry(_) => write!(f, "Retry"),
319            PluginConfig::Cache(_) => write!(f, "Cache"),
320            PluginConfig::Custom { name, .. } => write!(f, "Custom({})", name),
321        }
322    }
323}
324
325// ============================================================================
326// CLIENT PLUGIN TRAIT
327// ============================================================================
328
329/// Core trait for client plugins
330///
331/// Plugins can hook into the client lifecycle at various points:
332/// - **initialization**: Called when the plugin is registered
333/// - **before_request**: Called before sending requests to the server
334/// - **after_response**: Called after receiving responses from the server
335/// - **handle_custom**: Called for custom method handling
336///
337/// All methods are async and return PluginResult to allow for error handling
338/// and async operations like network calls, database access, etc.
339///
340/// # Examples
341///
342/// ```rust,no_run
343/// use turbomcp_client::plugins::{ClientPlugin, PluginContext, RequestContext, ResponseContext, PluginResult};
344/// use async_trait::async_trait;
345/// use serde_json::Value;
346///
347/// #[derive(Debug)]
348/// struct LoggingPlugin;
349///
350/// #[async_trait]
351/// impl ClientPlugin for LoggingPlugin {
352///     fn name(&self) -> &str {
353///         "logging"
354///     }
355///
356///     fn version(&self) -> &str {
357///         "1.0.0"
358///     }
359///
360///     async fn initialize(&self, context: &PluginContext) -> PluginResult<()> {
361///         println!("Logging plugin initialized for client: {}", context.client_name);
362///         Ok(())
363///     }
364///
365///     async fn before_request(&self, context: &mut RequestContext) -> PluginResult<()> {
366///         println!("Request: {} {}", context.method(),
367///             context.params().unwrap_or(&Value::Null));
368///         Ok(())
369///     }
370///
371///     async fn after_response(&self, context: &mut ResponseContext) -> PluginResult<()> {
372///         println!("Response: {} took {:?}", context.method(), context.duration);
373///         Ok(())
374///     }
375///
376///     async fn handle_custom(&self, method: &str, params: Option<Value>) -> PluginResult<Option<Value>> {
377///         if method == "logging.get_stats" {
378///             Ok(Some(serde_json::json!({"logged_requests": 42})))
379///         } else {
380///             Ok(None) // Not handled by this plugin
381///         }
382///     }
383/// }
384/// ```
385#[async_trait]
386pub trait ClientPlugin: Send + Sync + fmt::Debug {
387    /// Plugin name - must be unique across all registered plugins
388    fn name(&self) -> &str;
389
390    /// Plugin version
391    fn version(&self) -> &str;
392
393    /// Optional plugin description
394    fn description(&self) -> Option<&str> {
395        None
396    }
397
398    /// Plugin dependencies (other plugins that must be registered first)
399    fn dependencies(&self) -> Vec<&str> {
400        Vec::new()
401    }
402
403    /// Initialize the plugin
404    ///
405    /// Called once when the plugin is registered with the client.
406    /// Use this to set up resources, validate configuration, check dependencies, etc.
407    ///
408    /// # Arguments
409    ///
410    /// * `context` - Plugin context with client info, capabilities, and configuration
411    ///
412    /// # Returns
413    ///
414    /// Returns `Ok(())` if initialization succeeds, or `PluginError` if it fails.
415    async fn initialize(&self, context: &PluginContext) -> PluginResult<()>;
416
417    /// Hook called before sending a request to the server
418    ///
419    /// This allows plugins to:
420    /// - Modify request parameters
421    /// - Add metadata for tracking
422    /// - Implement features like authentication, request logging, etc.
423    /// - Abort requests by returning an error
424    ///
425    /// # Arguments
426    ///
427    /// * `context` - Mutable request context that can be modified
428    ///
429    /// # Returns
430    ///
431    /// Returns `Ok(())` to continue processing, or `PluginError` to abort.
432    async fn before_request(&self, context: &mut RequestContext) -> PluginResult<()>;
433
434    /// Hook called after receiving a response from the server
435    ///
436    /// This allows plugins to:
437    /// - Process response data
438    /// - Log metrics and performance data
439    /// - Implement features like caching, retry logic, etc.
440    /// - Modify response metadata
441    ///
442    /// # Arguments
443    ///
444    /// * `context` - Mutable response context that can be modified
445    ///
446    /// # Returns
447    ///
448    /// Returns `Ok(())` if processing succeeds, or `PluginError` if it fails.
449    async fn after_response(&self, context: &mut ResponseContext) -> PluginResult<()>;
450
451    /// Handle custom methods not part of the standard MCP protocol
452    ///
453    /// This allows plugins to implement custom functionality that can be invoked
454    /// by clients. Each plugin can handle its own set of custom methods.
455    ///
456    /// # Arguments
457    ///
458    /// * `method` - The custom method name (e.g., "metrics.get_stats")
459    /// * `params` - Optional parameters for the method
460    ///
461    /// # Returns
462    ///
463    /// Returns `Some(Value)` if the method was handled, `None` if not handled by this plugin,
464    /// or `PluginError` if handling failed.
465    async fn handle_custom(
466        &self,
467        method: &str,
468        params: Option<Value>,
469    ) -> PluginResult<Option<Value>>;
470
471    /// Optional cleanup when plugin is unregistered
472    ///
473    /// Default implementation does nothing. Override to perform cleanup
474    /// like closing connections, flushing buffers, etc.
475    async fn cleanup(&self) -> PluginResult<()> {
476        Ok(())
477    }
478}
479
480#[cfg(test)]
481mod tests {
482    use super::*;
483    use serde_json::json;
484    use turbomcp_protocol::MessageId;
485    use turbomcp_protocol::jsonrpc::{JsonRpcRequest, JsonRpcVersion};
486
487    #[test]
488    fn test_plugin_error_creation() {
489        let error = PluginError::initialization("Test error");
490        assert!(error.to_string().contains("Plugin initialization failed"));
491
492        let config_error = PluginError::configuration("Invalid config");
493        assert!(
494            config_error
495                .to_string()
496                .contains("Invalid plugin configuration")
497        );
498
499        let request_error = PluginError::request_processing("Request failed");
500        assert!(
501            request_error
502                .to_string()
503                .contains("Request processing error")
504        );
505    }
506
507    #[test]
508    fn test_plugin_context_creation() {
509        let capabilities = HashMap::from([
510            ("tools".to_string(), json!(true)),
511            ("sampling".to_string(), json!(false)),
512        ]);
513
514        let config = HashMap::from([
515            ("debug".to_string(), json!(true)),
516            ("timeout".to_string(), json!(5000)),
517        ]);
518
519        let plugins = vec!["metrics".to_string(), "retry".to_string()];
520
521        let context = PluginContext::new(
522            "test-client".to_string(),
523            "1.0.0".to_string(),
524            capabilities,
525            config,
526            plugins,
527        );
528
529        assert_eq!(context.client_name, "test-client");
530        assert_eq!(context.client_version, "1.0.0");
531        assert!(context.has_capability("tools"));
532        assert!(!context.has_capability("nonexistent"));
533        assert_eq!(context.get_config("debug"), Some(&json!(true)));
534        assert!(context.has_plugin("metrics"));
535        assert!(!context.has_plugin("nonexistent"));
536    }
537
538    #[test]
539    fn test_request_context_creation() {
540        let request = JsonRpcRequest {
541            jsonrpc: JsonRpcVersion,
542            id: MessageId::from("test-123"),
543            method: "test/method".to_string(),
544            params: Some(json!({"key": "value"})),
545        };
546
547        let metadata = HashMap::from([("user_id".to_string(), json!("user123"))]);
548
549        let mut context = RequestContext::new(request, metadata);
550
551        assert_eq!(context.method(), "test/method");
552        assert_eq!(context.params(), Some(&json!({"key": "value"})));
553        assert_eq!(context.get_metadata("user_id"), Some(&json!("user123")));
554
555        context.add_metadata("request_id".to_string(), json!("req456"));
556        assert_eq!(context.get_metadata("request_id"), Some(&json!("req456")));
557    }
558
559    #[test]
560    fn test_response_context_creation() {
561        let request = JsonRpcRequest {
562            jsonrpc: JsonRpcVersion,
563            id: MessageId::from("test-123"),
564            method: "test/method".to_string(),
565            params: Some(json!({"key": "value"})),
566        };
567
568        let request_context = RequestContext::new(request, HashMap::new());
569        let response = Some(json!({"result": "success"}));
570        let duration = Duration::from_millis(150);
571
572        let mut context = ResponseContext::new(request_context, response, None, duration);
573
574        assert!(context.is_success());
575        assert!(!context.is_error());
576        assert_eq!(context.method(), "test/method");
577        assert_eq!(context.duration, Duration::from_millis(150));
578
579        context.add_metadata("cache_hit".to_string(), json!(true));
580        assert_eq!(context.get_metadata("cache_hit"), Some(&json!(true)));
581    }
582
583    #[test]
584    fn test_plugin_config_serialization() {
585        let config = PluginConfig::Metrics;
586        let json_str = serde_json::to_string(&config).unwrap();
587        assert!(json_str.contains("metrics"));
588
589        let deserialized: PluginConfig = serde_json::from_str(&json_str).unwrap();
590        match deserialized {
591            PluginConfig::Metrics => {}
592            _ => panic!("Expected Metrics config"),
593        }
594    }
595
596    #[test]
597    fn test_plugin_config_display() {
598        let metrics_config = PluginConfig::Metrics;
599        assert_eq!(format!("{}", metrics_config), "Metrics");
600
601        let custom_config = PluginConfig::Custom {
602            name: "test".to_string(),
603            config: HashMap::new(),
604        };
605        assert_eq!(format!("{}", custom_config), "Custom(test)");
606    }
607}