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 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 pub fn has_capability(&self, capability: &str) -> bool {
166 self.capabilities.contains_key(capability)
167 }
168
169 pub fn get_config(&self, key: &str) -> Option<&Value> {
171 self.config.get(key)
172 }
173
174 pub fn has_plugin(&self, plugin_name: &str) -> bool {
176 self.available_plugins.contains(&plugin_name.to_string())
177 }
178}
179
180#[derive(Debug, Clone)]
182pub struct RequestContext {
183 pub request: JsonRpcRequest,
185
186 pub metadata: HashMap<String, Value>,
188
189 pub timestamp: DateTime<Utc>,
191}
192
193impl RequestContext {
194 pub fn new(request: JsonRpcRequest, metadata: HashMap<String, Value>) -> Self {
196 Self {
197 request,
198 metadata,
199 timestamp: Utc::now(),
200 }
201 }
202
203 pub fn method(&self) -> &str {
205 &self.request.method
206 }
207
208 pub fn params(&self) -> Option<&Value> {
210 self.request.params.as_ref()
211 }
212
213 pub fn add_metadata(&mut self, key: String, value: Value) {
215 self.metadata.insert(key, value);
216 }
217
218 pub fn get_metadata(&self, key: &str) -> Option<&Value> {
220 self.metadata.get(key)
221 }
222}
223
224#[derive(Debug, Clone)]
226pub struct ResponseContext {
227 pub request_context: RequestContext,
229
230 pub response: Option<Value>,
232
233 pub error: Option<turbomcp_protocol::Error>,
235
236 pub duration: Duration,
238
239 pub metadata: HashMap<String, Value>,
241}
242
243impl ResponseContext {
244 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 pub fn is_success(&self) -> bool {
262 self.error.is_none()
263 }
264
265 pub fn is_error(&self) -> bool {
267 self.error.is_some()
268 }
269
270 pub fn method(&self) -> &str {
272 self.request_context.method()
273 }
274
275 pub fn add_metadata(&mut self, key: String, value: Value) {
277 self.metadata.insert(key, value);
278 }
279
280 pub fn get_metadata(&self, key: &str) -> Option<&Value> {
282 self.metadata.get(key)
283 }
284}
285
286#[derive(Debug, Clone, Serialize, Deserialize)]
292#[serde(tag = "type")]
293pub enum PluginConfig {
294 #[serde(rename = "metrics")]
296 Metrics,
297
298 #[serde(rename = "retry")]
300 Retry(super::examples::RetryConfig),
301
302 #[serde(rename = "cache")]
304 Cache(super::examples::CacheConfig),
305
306 #[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#[async_trait]
386pub trait ClientPlugin: Send + Sync + fmt::Debug {
387 fn name(&self) -> &str;
389
390 fn version(&self) -> &str;
392
393 fn description(&self) -> Option<&str> {
395 None
396 }
397
398 fn dependencies(&self) -> Vec<&str> {
400 Vec::new()
401 }
402
403 async fn initialize(&self, context: &PluginContext) -> PluginResult<()>;
416
417 async fn before_request(&self, context: &mut RequestContext) -> PluginResult<()>;
433
434 async fn after_response(&self, context: &mut ResponseContext) -> PluginResult<()>;
450
451 async fn handle_custom(
466 &self,
467 method: &str,
468 params: Option<Value>,
469 ) -> PluginResult<Option<Value>>;
470
471 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}