1use 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#[derive(Error, Debug)]
22#[non_exhaustive]
23pub enum PluginError {
24 #[error("Plugin initialization failed: {message}")]
26 Initialization { message: String },
27
28 #[error("Invalid plugin configuration: {message}")]
30 Configuration { message: String },
31
32 #[error("Request processing error: {message}")]
34 RequestProcessing { message: String },
35
36 #[error("Response processing error: {message}")]
38 ResponseProcessing { message: String },
39
40 #[error("Custom handler error: {message}")]
42 CustomHandler { message: String },
43
44 #[error("Plugin dependency '{dependency}' not available")]
46 DependencyNotAvailable { dependency: String },
47
48 #[error("Plugin version incompatibility: {message}")]
50 VersionIncompatible { message: String },
51
52 #[error("Resource access error: {resource} - {message}")]
54 ResourceAccess { resource: String, message: String },
55
56 #[error("External system error: {source}")]
58 External {
59 #[from]
60 source: Box<dyn std::error::Error + Send + Sync>,
61 },
62}
63
64impl PluginError {
65 pub fn initialization(message: impl Into<String>) -> Self {
67 Self::Initialization {
68 message: message.into(),
69 }
70 }
71
72 pub fn configuration(message: impl Into<String>) -> Self {
74 Self::Configuration {
75 message: message.into(),
76 }
77 }
78
79 pub fn request_processing(message: impl Into<String>) -> Self {
81 Self::RequestProcessing {
82 message: message.into(),
83 }
84 }
85
86 pub fn response_processing(message: impl Into<String>) -> Self {
88 Self::ResponseProcessing {
89 message: message.into(),
90 }
91 }
92
93 pub fn custom_handler(message: impl Into<String>) -> Self {
95 Self::CustomHandler {
96 message: message.into(),
97 }
98 }
99
100 pub fn dependency_not_available(dependency: impl Into<String>) -> Self {
102 Self::DependencyNotAvailable {
103 dependency: dependency.into(),
104 }
105 }
106
107 pub fn version_incompatible(message: impl Into<String>) -> Self {
109 Self::VersionIncompatible {
110 message: message.into(),
111 }
112 }
113
114 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#[derive(Debug, Clone)]
131pub struct PluginContext {
132 pub client_name: String,
134 pub client_version: String,
135
136 pub capabilities: HashMap<String, Value>,
138
139 pub config: HashMap<String, Value>,
141
142 pub available_plugins: Vec<String>,
144}
145
146impl PluginContext {
147 #[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 #[must_use]
167 pub fn has_capability(&self, capability: &str) -> bool {
168 self.capabilities.contains_key(capability)
169 }
170
171 #[must_use]
173 pub fn get_config(&self, key: &str) -> Option<&Value> {
174 self.config.get(key)
175 }
176
177 #[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#[derive(Debug, Clone)]
186pub struct RequestContext {
187 pub request: JsonRpcRequest,
189
190 pub metadata: HashMap<String, Value>,
192
193 pub timestamp: DateTime<Utc>,
195}
196
197impl RequestContext {
198 #[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 #[must_use]
210 pub fn method(&self) -> &str {
211 &self.request.method
212 }
213
214 #[must_use]
216 pub fn params(&self) -> Option<&Value> {
217 self.request.params.as_ref()
218 }
219
220 pub fn add_metadata(&mut self, key: String, value: Value) {
222 self.metadata.insert(key, value);
223 }
224
225 #[must_use]
227 pub fn get_metadata(&self, key: &str) -> Option<&Value> {
228 self.metadata.get(key)
229 }
230}
231
232#[derive(Debug, Clone)]
234pub struct ResponseContext {
235 pub request_context: RequestContext,
237
238 pub response: Option<Value>,
240
241 pub error: Option<turbomcp_protocol::Error>,
243
244 pub duration: Duration,
246
247 pub metadata: HashMap<String, Value>,
249}
250
251impl ResponseContext {
252 #[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 pub fn is_success(&self) -> bool {
271 self.error.is_none()
272 }
273
274 pub fn is_error(&self) -> bool {
276 self.error.is_some()
277 }
278
279 pub fn method(&self) -> &str {
281 self.request_context.method()
282 }
283
284 pub fn add_metadata(&mut self, key: String, value: Value) {
286 self.metadata.insert(key, value);
287 }
288
289 pub fn get_metadata(&self, key: &str) -> Option<&Value> {
291 self.metadata.get(key)
292 }
293}
294
295#[derive(Debug, Clone, Serialize, Deserialize)]
301#[serde(tag = "type")]
302pub enum PluginConfig {
303 #[serde(rename = "metrics")]
305 Metrics,
306
307 #[serde(rename = "retry")]
309 Retry(super::examples::RetryConfig),
310
311 #[serde(rename = "cache")]
313 Cache(super::examples::CacheConfig),
314
315 #[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#[async_trait]
395pub trait ClientPlugin: Send + Sync + fmt::Debug {
396 fn name(&self) -> &str;
398
399 fn version(&self) -> &str;
401
402 fn description(&self) -> Option<&str> {
404 None
405 }
406
407 fn dependencies(&self) -> Vec<&str> {
409 Vec::new()
410 }
411
412 async fn initialize(&self, context: &PluginContext) -> PluginResult<()>;
425
426 async fn before_request(&self, context: &mut RequestContext) -> PluginResult<()>;
442
443 async fn after_response(&self, context: &mut ResponseContext) -> PluginResult<()>;
459
460 async fn handle_custom(
475 &self,
476 method: &str,
477 params: Option<Value>,
478 ) -> PluginResult<Option<Value>>;
479
480 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}