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    #[must_use]
149    pub fn new(
150        client_name: String,
151        client_version: String,
152        capabilities: HashMap<String, Value>,
153        config: HashMap<String, Value>,
154        available_plugins: Vec<String>,
155    ) -> Self {
156        Self {
157            client_name,
158            client_version,
159            capabilities,
160            config,
161            available_plugins,
162        }
163    }
164
165    /// Check if a capability is available
166    #[must_use]
167    pub fn has_capability(&self, capability: &str) -> bool {
168        self.capabilities.contains_key(capability)
169    }
170
171    /// Get a configuration value
172    #[must_use]
173    pub fn get_config(&self, key: &str) -> Option<&Value> {
174        self.config.get(key)
175    }
176
177    /// Check if a plugin dependency is available
178    #[must_use]
179    pub fn has_plugin(&self, plugin_name: &str) -> bool {
180        self.available_plugins.contains(&plugin_name.to_string())
181    }
182}
183
184/// Context for request processing
185#[derive(Debug, Clone)]
186pub struct RequestContext {
187    /// The JSON-RPC request being processed
188    pub request: JsonRpcRequest,
189
190    /// Additional metadata (can be modified by plugins)
191    pub metadata: HashMap<String, Value>,
192
193    /// Request timestamp
194    pub timestamp: DateTime<Utc>,
195}
196
197impl RequestContext {
198    /// Create a new request context
199    #[must_use]
200    pub fn new(request: JsonRpcRequest, metadata: HashMap<String, Value>) -> Self {
201        Self {
202            request,
203            metadata,
204            timestamp: Utc::now(),
205        }
206    }
207
208    /// Get the request method
209    #[must_use]
210    pub fn method(&self) -> &str {
211        &self.request.method
212    }
213
214    /// Get request parameters
215    #[must_use]
216    pub fn params(&self) -> Option<&Value> {
217        self.request.params.as_ref()
218    }
219
220    /// Add metadata
221    pub fn add_metadata(&mut self, key: String, value: Value) {
222        self.metadata.insert(key, value);
223    }
224
225    /// Get metadata value
226    #[must_use]
227    pub fn get_metadata(&self, key: &str) -> Option<&Value> {
228        self.metadata.get(key)
229    }
230}
231
232/// Context for response processing
233#[derive(Debug, Clone)]
234pub struct ResponseContext {
235    /// The original request context
236    pub request_context: RequestContext,
237
238    /// The response data (if successful)
239    pub response: Option<Value>,
240
241    /// Error information (if failed)
242    pub error: Option<turbomcp_protocol::Error>,
243
244    /// Request duration
245    pub duration: Duration,
246
247    /// Additional metadata (can be modified by plugins)
248    pub metadata: HashMap<String, Value>,
249}
250
251impl ResponseContext {
252    /// Create a new response context
253    #[must_use]
254    pub fn new(
255        request_context: RequestContext,
256        response: Option<Value>,
257        error: Option<turbomcp_protocol::Error>,
258        duration: Duration,
259    ) -> Self {
260        Self {
261            request_context,
262            response,
263            error,
264            duration,
265            metadata: HashMap::new(),
266        }
267    }
268
269    /// Check if the response was successful
270    pub fn is_success(&self) -> bool {
271        self.error.is_none()
272    }
273
274    /// Check if the response was an error
275    pub fn is_error(&self) -> bool {
276        self.error.is_some()
277    }
278
279    /// Get the request method
280    pub fn method(&self) -> &str {
281        self.request_context.method()
282    }
283
284    /// Add metadata
285    pub fn add_metadata(&mut self, key: String, value: Value) {
286        self.metadata.insert(key, value);
287    }
288
289    /// Get metadata value
290    pub fn get_metadata(&self, key: &str) -> Option<&Value> {
291        self.metadata.get(key)
292    }
293}
294
295// ============================================================================
296// PLUGIN CONFIGURATION
297// ============================================================================
298
299/// Plugin configuration variants
300#[derive(Debug, Clone, Serialize, Deserialize)]
301#[serde(tag = "type")]
302pub enum PluginConfig {
303    /// Metrics plugin configuration
304    #[serde(rename = "metrics")]
305    Metrics,
306
307    /// Retry plugin configuration
308    #[serde(rename = "retry")]
309    Retry(super::examples::RetryConfig),
310
311    /// Cache plugin configuration
312    #[serde(rename = "cache")]
313    Cache(super::examples::CacheConfig),
314
315    /// Custom plugin configuration
316    #[serde(rename = "custom")]
317    Custom {
318        name: String,
319        config: HashMap<String, Value>,
320    },
321}
322
323impl fmt::Display for PluginConfig {
324    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
325        match self {
326            PluginConfig::Metrics => write!(f, "Metrics"),
327            PluginConfig::Retry(_) => write!(f, "Retry"),
328            PluginConfig::Cache(_) => write!(f, "Cache"),
329            PluginConfig::Custom { name, .. } => write!(f, "Custom({})", name),
330        }
331    }
332}
333
334// ============================================================================
335// CLIENT PLUGIN TRAIT
336// ============================================================================
337
338/// Core trait for client plugins
339///
340/// Plugins can hook into the client lifecycle at various points:
341/// - **initialization**: Called when the plugin is registered
342/// - **before_request**: Called before sending requests to the server
343/// - **after_response**: Called after receiving responses from the server
344/// - **handle_custom**: Called for custom method handling
345///
346/// All methods are async and return PluginResult to allow for error handling
347/// and async operations like network calls, database access, etc.
348///
349/// # Examples
350///
351/// ```rust,no_run
352/// use turbomcp_client::plugins::{ClientPlugin, PluginContext, RequestContext, ResponseContext, PluginResult};
353/// use async_trait::async_trait;
354/// use serde_json::Value;
355///
356/// #[derive(Debug)]
357/// struct LoggingPlugin;
358///
359/// #[async_trait]
360/// impl ClientPlugin for LoggingPlugin {
361///     fn name(&self) -> &str {
362///         "logging"
363///     }
364///
365///     fn version(&self) -> &str {
366///         "1.0.0"
367///     }
368///
369///     async fn initialize(&self, context: &PluginContext) -> PluginResult<()> {
370///         println!("Logging plugin initialized for client: {}", context.client_name);
371///         Ok(())
372///     }
373///
374///     async fn before_request(&self, context: &mut RequestContext) -> PluginResult<()> {
375///         println!("Request: {} {}", context.method(),
376///             context.params().unwrap_or(&Value::Null));
377///         Ok(())
378///     }
379///
380///     async fn after_response(&self, context: &mut ResponseContext) -> PluginResult<()> {
381///         println!("Response: {} took {:?}", context.method(), context.duration);
382///         Ok(())
383///     }
384///
385///     async fn handle_custom(&self, method: &str, params: Option<Value>) -> PluginResult<Option<Value>> {
386///         if method == "logging.get_stats" {
387///             Ok(Some(serde_json::json!({"logged_requests": 42})))
388///         } else {
389///             Ok(None) // Not handled by this plugin
390///         }
391///     }
392/// }
393/// ```
394#[async_trait]
395pub trait ClientPlugin: Send + Sync + fmt::Debug {
396    /// Plugin name - must be unique across all registered plugins
397    fn name(&self) -> &str;
398
399    /// Plugin version
400    fn version(&self) -> &str;
401
402    /// Optional plugin description
403    fn description(&self) -> Option<&str> {
404        None
405    }
406
407    /// Plugin dependencies (other plugins that must be registered first)
408    fn dependencies(&self) -> Vec<&str> {
409        Vec::new()
410    }
411
412    /// Initialize the plugin
413    ///
414    /// Called once when the plugin is registered with the client.
415    /// Use this to set up resources, validate configuration, check dependencies, etc.
416    ///
417    /// # Arguments
418    ///
419    /// * `context` - Plugin context with client info, capabilities, and configuration
420    ///
421    /// # Returns
422    ///
423    /// Returns `Ok(())` if initialization succeeds, or `PluginError` if it fails.
424    async fn initialize(&self, context: &PluginContext) -> PluginResult<()>;
425
426    /// Hook called before sending a request to the server
427    ///
428    /// This allows plugins to:
429    /// - Modify request parameters
430    /// - Add metadata for tracking
431    /// - Implement features like authentication, request logging, etc.
432    /// - Abort requests by returning an error
433    ///
434    /// # Arguments
435    ///
436    /// * `context` - Mutable request context that can be modified
437    ///
438    /// # Returns
439    ///
440    /// Returns `Ok(())` to continue processing, or `PluginError` to abort.
441    async fn before_request(&self, context: &mut RequestContext) -> PluginResult<()>;
442
443    /// Hook called after receiving a response from the server
444    ///
445    /// This allows plugins to:
446    /// - Process response data
447    /// - Log metrics and performance data
448    /// - Implement features like caching, retry logic, etc.
449    /// - Modify response metadata
450    ///
451    /// # Arguments
452    ///
453    /// * `context` - Mutable response context that can be modified
454    ///
455    /// # Returns
456    ///
457    /// Returns `Ok(())` if processing succeeds, or `PluginError` if it fails.
458    async fn after_response(&self, context: &mut ResponseContext) -> PluginResult<()>;
459
460    /// Handle custom methods not part of the standard MCP protocol
461    ///
462    /// This allows plugins to implement custom functionality that can be invoked
463    /// by clients. Each plugin can handle its own set of custom methods.
464    ///
465    /// # Arguments
466    ///
467    /// * `method` - The custom method name (e.g., "metrics.get_stats")
468    /// * `params` - Optional parameters for the method
469    ///
470    /// # Returns
471    ///
472    /// Returns `Some(Value)` if the method was handled, `None` if not handled by this plugin,
473    /// or `PluginError` if handling failed.
474    async fn handle_custom(
475        &self,
476        method: &str,
477        params: Option<Value>,
478    ) -> PluginResult<Option<Value>>;
479
480    /// Optional cleanup when plugin is unregistered
481    ///
482    /// Default implementation does nothing. Override to perform cleanup
483    /// like closing connections, flushing buffers, etc.
484    async fn cleanup(&self) -> PluginResult<()> {
485        Ok(())
486    }
487}
488
489#[cfg(test)]
490mod tests {
491    use super::*;
492    use serde_json::json;
493    use turbomcp_protocol::MessageId;
494    use turbomcp_protocol::jsonrpc::{JsonRpcRequest, JsonRpcVersion};
495
496    #[test]
497    fn test_plugin_error_creation() {
498        let error = PluginError::initialization("Test error");
499        assert!(error.to_string().contains("Plugin initialization failed"));
500
501        let config_error = PluginError::configuration("Invalid config");
502        assert!(
503            config_error
504                .to_string()
505                .contains("Invalid plugin configuration")
506        );
507
508        let request_error = PluginError::request_processing("Request failed");
509        assert!(
510            request_error
511                .to_string()
512                .contains("Request processing error")
513        );
514    }
515
516    #[test]
517    fn test_plugin_context_creation() {
518        let capabilities = HashMap::from([
519            ("tools".to_string(), json!(true)),
520            ("sampling".to_string(), json!(false)),
521        ]);
522
523        let config = HashMap::from([
524            ("debug".to_string(), json!(true)),
525            ("timeout".to_string(), json!(5000)),
526        ]);
527
528        let plugins = vec!["metrics".to_string(), "retry".to_string()];
529
530        let context = PluginContext::new(
531            "test-client".to_string(),
532            "1.0.0".to_string(),
533            capabilities,
534            config,
535            plugins,
536        );
537
538        assert_eq!(context.client_name, "test-client");
539        assert_eq!(context.client_version, "1.0.0");
540        assert!(context.has_capability("tools"));
541        assert!(!context.has_capability("nonexistent"));
542        assert_eq!(context.get_config("debug"), Some(&json!(true)));
543        assert!(context.has_plugin("metrics"));
544        assert!(!context.has_plugin("nonexistent"));
545    }
546
547    #[test]
548    fn test_request_context_creation() {
549        let request = JsonRpcRequest {
550            jsonrpc: JsonRpcVersion,
551            id: MessageId::from("test-123"),
552            method: "test/method".to_string(),
553            params: Some(json!({"key": "value"})),
554        };
555
556        let metadata = HashMap::from([("user_id".to_string(), json!("user123"))]);
557
558        let mut context = RequestContext::new(request, metadata);
559
560        assert_eq!(context.method(), "test/method");
561        assert_eq!(context.params(), Some(&json!({"key": "value"})));
562        assert_eq!(context.get_metadata("user_id"), Some(&json!("user123")));
563
564        context.add_metadata("request_id".to_string(), json!("req456"));
565        assert_eq!(context.get_metadata("request_id"), Some(&json!("req456")));
566    }
567
568    #[test]
569    fn test_response_context_creation() {
570        let request = JsonRpcRequest {
571            jsonrpc: JsonRpcVersion,
572            id: MessageId::from("test-123"),
573            method: "test/method".to_string(),
574            params: Some(json!({"key": "value"})),
575        };
576
577        let request_context = RequestContext::new(request, HashMap::new());
578        let response = Some(json!({"result": "success"}));
579        let duration = Duration::from_millis(150);
580
581        let mut context = ResponseContext::new(request_context, response, None, duration);
582
583        assert!(context.is_success());
584        assert!(!context.is_error());
585        assert_eq!(context.method(), "test/method");
586        assert_eq!(context.duration, Duration::from_millis(150));
587
588        context.add_metadata("cache_hit".to_string(), json!(true));
589        assert_eq!(context.get_metadata("cache_hit"), Some(&json!(true)));
590    }
591
592    #[test]
593    fn test_plugin_config_serialization() {
594        let config = PluginConfig::Metrics;
595        let json_str = serde_json::to_string(&config).unwrap();
596        assert!(json_str.contains("metrics"));
597
598        let deserialized: PluginConfig = serde_json::from_str(&json_str).unwrap();
599        match deserialized {
600            PluginConfig::Metrics => {}
601            _ => panic!("Expected Metrics config"),
602        }
603    }
604
605    #[test]
606    fn test_plugin_config_display() {
607        let metrics_config = PluginConfig::Metrics;
608        assert_eq!(format!("{}", metrics_config), "Metrics");
609
610        let custom_config = PluginConfig::Custom {
611            name: "test".to_string(),
612            config: HashMap::new(),
613        };
614        assert_eq!(format!("{}", custom_config), "Custom(test)");
615    }
616}