turbomcp_core/
context.rs

1//! Request and response context for rich metadata handling.
2//!
3//! Enhanced with client identification, session management, and request analytics
4//! for comprehensive MCP application monitoring and management.
5//!
6//! # Examples
7//!
8//! ## Creating a basic request context
9//!
10//! ```
11//! use turbomcp_core::RequestContext;
12//!
13//! let ctx = RequestContext::new();
14//! println!("Request ID: {}", ctx.request_id);
15//! assert!(!ctx.request_id.is_empty());
16//! ```
17//!
18//! ## Building a context with metadata
19//!
20//! ```
21//! use turbomcp_core::RequestContext;
22//!
23//! let ctx = RequestContext::new()
24//!     .with_user_id("user123")
25//!     .with_session_id("session456")
26//!     .with_metadata("api_version", "2.0")
27//!     .with_metadata("client", "web_app");
28//!
29//! assert_eq!(ctx.user_id, Some("user123".to_string()));
30//! assert_eq!(ctx.session_id, Some("session456".to_string()));
31//! assert_eq!(ctx.get_metadata("api_version"), Some(&serde_json::json!("2.0")));
32//! ```
33//!
34//! ## Working with response contexts
35//!
36//! ```
37//! use turbomcp_core::{RequestContext, ResponseContext};
38//! use std::time::Duration;
39//!
40//! let request_ctx = RequestContext::with_id("req-123");
41//! let duration = Duration::from_millis(250);
42//!
43//! // Successful response
44//! let success_ctx = ResponseContext::success(&request_ctx.request_id, duration);
45//!
46//! // Error response
47//! let error_ctx = ResponseContext::error(&request_ctx.request_id, duration, -32600, "Invalid Request");
48//!
49//! assert_eq!(success_ctx.request_id, "req-123");
50//! assert_eq!(error_ctx.request_id, "req-123");
51//! ```
52
53use std::collections::HashMap;
54use std::fmt;
55use std::sync::Arc;
56use std::time::Instant;
57
58use chrono::{DateTime, Utc};
59use serde::{Deserialize, Serialize};
60use tokio_util::sync::CancellationToken;
61use uuid::Uuid;
62
63use crate::types::Timestamp;
64
65/// Context information for request processing
66#[derive(Debug, Clone)]
67pub struct RequestContext {
68    /// Unique request identifier
69    pub request_id: String,
70
71    /// User identifier (if authenticated)
72    pub user_id: Option<String>,
73
74    /// Session identifier
75    pub session_id: Option<String>,
76
77    /// Client identifier
78    pub client_id: Option<String>,
79
80    /// Request timestamp
81    pub timestamp: Timestamp,
82
83    /// Request start time for performance tracking
84    pub start_time: Instant,
85
86    /// Custom metadata
87    pub metadata: Arc<HashMap<String, serde_json::Value>>,
88
89    /// Tracing span context
90    #[cfg(feature = "tracing")]
91    pub span: Option<tracing::Span>,
92
93    /// Cancellation token
94    pub cancellation_token: Option<Arc<CancellationToken>>,
95}
96
97/// Context information for response processing
98#[derive(Debug, Clone)]
99pub struct ResponseContext {
100    /// Original request ID
101    pub request_id: String,
102
103    /// Response timestamp
104    pub timestamp: Timestamp,
105
106    /// Processing duration
107    pub duration: std::time::Duration,
108
109    /// Response status
110    pub status: ResponseStatus,
111
112    /// Custom metadata
113    pub metadata: Arc<HashMap<String, serde_json::Value>>,
114}
115
116/// Response status information
117#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
118pub enum ResponseStatus {
119    /// Successful response
120    Success,
121    /// Error response
122    Error {
123        /// Error code
124        code: i32,
125        /// Error message
126        message: String,
127    },
128    /// Partial response (streaming)
129    Partial,
130    /// Cancelled response
131    Cancelled,
132}
133
134impl RequestContext {
135    /// Create a new request context
136    ///
137    /// # Examples
138    ///
139    /// ```
140    /// use turbomcp_core::RequestContext;
141    ///
142    /// let ctx = RequestContext::new();
143    /// assert!(!ctx.request_id.is_empty());
144    /// assert!(ctx.user_id.is_none());
145    /// assert!(ctx.session_id.is_none());
146    /// assert!(ctx.metadata.is_empty());
147    /// ```
148    #[must_use]
149    pub fn new() -> Self {
150        Self {
151            request_id: Uuid::new_v4().to_string(),
152            user_id: None,
153            session_id: None,
154            client_id: None,
155            timestamp: Timestamp::now(),
156            start_time: Instant::now(),
157            metadata: Arc::new(HashMap::new()),
158            #[cfg(feature = "tracing")]
159            span: None,
160            cancellation_token: None,
161        }
162    }
163    /// Return true if the request is authenticated according to context metadata
164    ///
165    /// # Examples
166    ///
167    /// ```
168    /// use turbomcp_core::RequestContext;
169    ///
170    /// let ctx = RequestContext::new()
171    ///     .with_metadata("authenticated", true);
172    /// assert!(ctx.is_authenticated());
173    ///
174    /// let unauth_ctx = RequestContext::new();
175    /// assert!(!unauth_ctx.is_authenticated());
176    /// ```
177    #[must_use]
178    pub fn is_authenticated(&self) -> bool {
179        self.metadata
180            .get("authenticated")
181            .and_then(serde_json::Value::as_bool)
182            .unwrap_or(false)
183    }
184
185    /// Return user id if present
186    #[must_use]
187    pub fn user(&self) -> Option<&str> {
188        self.user_id.as_deref()
189    }
190
191    /// Return roles from `auth.roles` metadata, if present
192    #[must_use]
193    pub fn roles(&self) -> Vec<String> {
194        self.metadata
195            .get("auth")
196            .and_then(|v| v.get("roles"))
197            .and_then(|v| v.as_array())
198            .map(|arr| {
199                arr.iter()
200                    .filter_map(|v| v.as_str().map(std::string::ToString::to_string))
201                    .collect()
202            })
203            .unwrap_or_default()
204    }
205
206    /// Return true if the user has any of the required roles
207    pub fn has_any_role<S: AsRef<str>>(&self, required: &[S]) -> bool {
208        if required.is_empty() {
209            return true;
210        }
211        let user_roles = self.roles();
212        if user_roles.is_empty() {
213            return false;
214        }
215        let set: std::collections::HashSet<_> = user_roles.into_iter().collect();
216        required.iter().any(|r| set.contains(r.as_ref()))
217    }
218
219    /// Create a request context with specific ID
220    pub fn with_id(id: impl Into<String>) -> Self {
221        Self {
222            request_id: id.into(),
223            ..Self::new()
224        }
225    }
226
227    /// Set the user ID
228    #[must_use]
229    pub fn with_user_id(mut self, user_id: impl Into<String>) -> Self {
230        self.user_id = Some(user_id.into());
231        self
232    }
233
234    /// Set the session ID
235    #[must_use]
236    pub fn with_session_id(mut self, session_id: impl Into<String>) -> Self {
237        self.session_id = Some(session_id.into());
238        self
239    }
240
241    /// Set the client ID
242    #[must_use]
243    pub fn with_client_id(mut self, client_id: impl Into<String>) -> Self {
244        self.client_id = Some(client_id.into());
245        self
246    }
247
248    /// Add metadata
249    #[must_use]
250    pub fn with_metadata(
251        mut self,
252        key: impl Into<String>,
253        value: impl Into<serde_json::Value>,
254    ) -> Self {
255        Arc::make_mut(&mut self.metadata).insert(key.into(), value.into());
256        self
257    }
258
259    /// Set cancellation token
260    #[must_use]
261    pub fn with_cancellation_token(mut self, token: Arc<CancellationToken>) -> Self {
262        self.cancellation_token = Some(token);
263        self
264    }
265
266    /// Get elapsed time since request started
267    #[must_use]
268    pub fn elapsed(&self) -> std::time::Duration {
269        self.start_time.elapsed()
270    }
271
272    /// Check if request is cancelled
273    #[must_use]
274    pub fn is_cancelled(&self) -> bool {
275        self.cancellation_token
276            .as_ref()
277            .is_some_and(|token| token.is_cancelled())
278    }
279
280    /// Get metadata value
281    #[must_use]
282    pub fn get_metadata(&self, key: &str) -> Option<&serde_json::Value> {
283        self.metadata.get(key)
284    }
285
286    /// Clone with new request ID (for sub-requests)
287    #[must_use]
288    pub fn derive(&self) -> Self {
289        Self {
290            request_id: Uuid::new_v4().to_string(),
291            user_id: self.user_id.clone(),
292            session_id: self.session_id.clone(),
293            client_id: self.client_id.clone(),
294            timestamp: Timestamp::now(),
295            start_time: Instant::now(),
296            metadata: self.metadata.clone(),
297            #[cfg(feature = "tracing")]
298            span: None,
299            cancellation_token: self.cancellation_token.clone(),
300        }
301    }
302}
303
304impl ResponseContext {
305    /// Create a successful response context
306    pub fn success(request_id: impl Into<String>, duration: std::time::Duration) -> Self {
307        Self {
308            request_id: request_id.into(),
309            timestamp: Timestamp::now(),
310            duration,
311            status: ResponseStatus::Success,
312            metadata: Arc::new(HashMap::new()),
313        }
314    }
315
316    /// Create an error response context
317    pub fn error(
318        request_id: impl Into<String>,
319        duration: std::time::Duration,
320        code: i32,
321        message: impl Into<String>,
322    ) -> Self {
323        Self {
324            request_id: request_id.into(),
325            timestamp: Timestamp::now(),
326            duration,
327            status: ResponseStatus::Error {
328                code,
329                message: message.into(),
330            },
331            metadata: Arc::new(HashMap::new()),
332        }
333    }
334
335    /// Create a cancelled response context
336    pub fn cancelled(request_id: impl Into<String>, duration: std::time::Duration) -> Self {
337        Self {
338            request_id: request_id.into(),
339            timestamp: Timestamp::now(),
340            duration,
341            status: ResponseStatus::Cancelled,
342            metadata: Arc::new(HashMap::new()),
343        }
344    }
345
346    /// Add metadata
347    #[must_use]
348    pub fn with_metadata(
349        mut self,
350        key: impl Into<String>,
351        value: impl Into<serde_json::Value>,
352    ) -> Self {
353        Arc::make_mut(&mut self.metadata).insert(key.into(), value.into());
354        self
355    }
356
357    /// Check if response is successful
358    #[must_use]
359    pub const fn is_success(&self) -> bool {
360        matches!(self.status, ResponseStatus::Success)
361    }
362
363    /// Check if response is an error
364    #[must_use]
365    pub const fn is_error(&self) -> bool {
366        matches!(self.status, ResponseStatus::Error { .. })
367    }
368
369    /// Get error information if response is an error
370    #[must_use]
371    pub fn error_info(&self) -> Option<(i32, &str)> {
372        match &self.status {
373            ResponseStatus::Error { code, message } => Some((*code, message)),
374            _ => None,
375        }
376    }
377}
378
379impl Default for RequestContext {
380    fn default() -> Self {
381        Self::new()
382    }
383}
384
385impl fmt::Display for ResponseStatus {
386    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
387        match self {
388            Self::Success => write!(f, "Success"),
389            Self::Error { code, message } => write!(f, "Error({code}: {message})"),
390            Self::Partial => write!(f, "Partial"),
391            Self::Cancelled => write!(f, "Cancelled"),
392        }
393    }
394}
395
396// ============================================================================
397// Enhanced Client Management and Session Tracking
398// ============================================================================
399
400/// Client identification methods for enhanced request routing and analytics
401#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
402pub enum ClientId {
403    /// Explicit client ID from header
404    Header(String),
405    /// Bearer token from Authorization header
406    Token(String),
407    /// Session cookie
408    Session(String),
409    /// Query parameter
410    QueryParam(String),
411    /// Hash of User-Agent (fallback)
412    UserAgent(String),
413    /// Anonymous client
414    Anonymous,
415}
416
417impl ClientId {
418    /// Get the string representation of the client ID
419    #[must_use]
420    pub fn as_str(&self) -> &str {
421        match self {
422            Self::Header(id)
423            | Self::Token(id)
424            | Self::Session(id)
425            | Self::QueryParam(id)
426            | Self::UserAgent(id) => id,
427            Self::Anonymous => "anonymous",
428        }
429    }
430
431    /// Check if the client is authenticated
432    #[must_use]
433    pub const fn is_authenticated(&self) -> bool {
434        matches!(self, Self::Token(_) | Self::Session(_))
435    }
436
437    /// Get the authentication method
438    #[must_use]
439    pub const fn auth_method(&self) -> &'static str {
440        match self {
441            Self::Header(_) => "header",
442            Self::Token(_) => "bearer_token",
443            Self::Session(_) => "session_cookie",
444            Self::QueryParam(_) => "query_param",
445            Self::UserAgent(_) => "user_agent",
446            Self::Anonymous => "anonymous",
447        }
448    }
449}
450
451/// Client session information for tracking and analytics
452#[derive(Debug, Clone, Serialize, Deserialize)]
453pub struct ClientSession {
454    /// Unique client identifier
455    pub client_id: String,
456    /// Client name (optional, human-readable)
457    pub client_name: Option<String>,
458    /// When the client connected
459    pub connected_at: DateTime<Utc>,
460    /// Last activity timestamp
461    pub last_activity: DateTime<Utc>,
462    /// Number of requests made
463    pub request_count: usize,
464    /// Transport type (stdio, http, websocket, etc.)
465    pub transport_type: String,
466    /// Authentication status
467    pub authenticated: bool,
468    /// Client capabilities (optional)
469    pub capabilities: Option<serde_json::Value>,
470    /// Additional metadata
471    pub metadata: HashMap<String, serde_json::Value>,
472}
473
474impl ClientSession {
475    /// Create a new client session
476    #[must_use]
477    pub fn new(client_id: String, transport_type: String) -> Self {
478        let now = Utc::now();
479        Self {
480            client_id,
481            client_name: None,
482            connected_at: now,
483            last_activity: now,
484            request_count: 0,
485            transport_type,
486            authenticated: false,
487            capabilities: None,
488            metadata: HashMap::new(),
489        }
490    }
491
492    /// Update activity timestamp and increment request count
493    pub fn update_activity(&mut self) {
494        self.last_activity = Utc::now();
495        self.request_count += 1;
496    }
497
498    /// Set authentication status and client info
499    pub fn authenticate(&mut self, client_name: Option<String>) {
500        self.authenticated = true;
501        self.client_name = client_name;
502    }
503
504    /// Set client capabilities
505    pub fn set_capabilities(&mut self, capabilities: serde_json::Value) {
506        self.capabilities = Some(capabilities);
507    }
508
509    /// Get session duration
510    #[must_use]
511    pub fn session_duration(&self) -> chrono::Duration {
512        self.last_activity - self.connected_at
513    }
514
515    /// Check if session is idle (no activity for specified duration)
516    #[must_use]
517    pub fn is_idle(&self, idle_threshold: chrono::Duration) -> bool {
518        Utc::now() - self.last_activity > idle_threshold
519    }
520}
521
522/// Request analytics information for monitoring and debugging
523#[derive(Debug, Clone, Serialize, Deserialize)]
524pub struct RequestInfo {
525    /// Request timestamp
526    pub timestamp: DateTime<Utc>,
527    /// Client identifier
528    pub client_id: String,
529    /// Tool or method name
530    pub method_name: String,
531    /// Request parameters (sanitized for privacy)
532    pub parameters: serde_json::Value,
533    /// Response time in milliseconds
534    pub response_time_ms: Option<u64>,
535    /// Success status
536    pub success: bool,
537    /// Error message if failed
538    pub error_message: Option<String>,
539    /// HTTP status code (if applicable)
540    pub status_code: Option<u16>,
541    /// Additional metadata
542    pub metadata: HashMap<String, serde_json::Value>,
543}
544
545impl RequestInfo {
546    /// Create a new request info
547    #[must_use]
548    pub fn new(client_id: String, method_name: String, parameters: serde_json::Value) -> Self {
549        Self {
550            timestamp: Utc::now(),
551            client_id,
552            method_name,
553            parameters,
554            response_time_ms: None,
555            success: false,
556            error_message: None,
557            status_code: None,
558            metadata: HashMap::new(),
559        }
560    }
561
562    /// Mark the request as completed successfully
563    #[must_use]
564    pub const fn complete_success(mut self, response_time_ms: u64) -> Self {
565        self.response_time_ms = Some(response_time_ms);
566        self.success = true;
567        self.status_code = Some(200);
568        self
569    }
570
571    /// Mark the request as failed
572    #[must_use]
573    pub fn complete_error(mut self, response_time_ms: u64, error: String) -> Self {
574        self.response_time_ms = Some(response_time_ms);
575        self.success = false;
576        self.error_message = Some(error);
577        self.status_code = Some(500);
578        self
579    }
580
581    /// Set HTTP status code
582    #[must_use]
583    pub const fn with_status_code(mut self, code: u16) -> Self {
584        self.status_code = Some(code);
585        self
586    }
587
588    /// Add metadata
589    #[must_use]
590    pub fn with_metadata(mut self, key: String, value: serde_json::Value) -> Self {
591        self.metadata.insert(key, value);
592        self
593    }
594}
595
596/// Client identification extractor for various transport mechanisms
597#[derive(Debug)]
598pub struct ClientIdExtractor {
599    /// Authentication tokens mapping token -> `client_id`
600    auth_tokens: Arc<dashmap::DashMap<String, String>>,
601}
602
603impl ClientIdExtractor {
604    /// Create a new client ID extractor
605    #[must_use]
606    pub fn new() -> Self {
607        Self {
608            auth_tokens: Arc::new(dashmap::DashMap::new()),
609        }
610    }
611
612    /// Register an authentication token for a client
613    pub fn register_token(&self, token: String, client_id: String) {
614        self.auth_tokens.insert(token, client_id);
615    }
616
617    /// Remove an authentication token
618    pub fn revoke_token(&self, token: &str) {
619        self.auth_tokens.remove(token);
620    }
621
622    /// List all registered tokens (for admin purposes)
623    #[must_use]
624    pub fn list_tokens(&self) -> Vec<(String, String)> {
625        self.auth_tokens
626            .iter()
627            .map(|entry| (entry.key().clone(), entry.value().clone()))
628            .collect()
629    }
630
631    /// Extract client ID from HTTP headers
632    #[must_use]
633    #[allow(clippy::significant_drop_tightening)]
634    pub fn extract_from_http_headers(&self, headers: &HashMap<String, String>) -> ClientId {
635        // 1. Check for explicit client ID header
636        if let Some(client_id) = headers.get("x-client-id") {
637            return ClientId::Header(client_id.clone());
638        }
639
640        // 2. Check for Authorization header with Bearer token
641        if let Some(auth) = headers.get("authorization")
642            && let Some(token) = auth.strip_prefix("Bearer ")
643        {
644            // Look up client ID from token
645            let token_lookup = self.auth_tokens.iter().find(|e| e.key() == token);
646            if let Some(entry) = token_lookup {
647                let client_id = entry.value().clone();
648                drop(entry); // Explicitly drop the lock guard early
649                return ClientId::Token(client_id);
650            }
651            // Token not found - return the token itself as identifier
652            return ClientId::Token(token.to_string());
653        }
654
655        // 3. Check for session cookie
656        if let Some(cookie) = headers.get("cookie") {
657            for cookie_part in cookie.split(';') {
658                let parts: Vec<&str> = cookie_part.trim().splitn(2, '=').collect();
659                if parts.len() == 2 && (parts[0] == "session_id" || parts[0] == "sessionid") {
660                    return ClientId::Session(parts[1].to_string());
661                }
662            }
663        }
664
665        // 4. Use User-Agent hash as fallback
666        if let Some(user_agent) = headers.get("user-agent") {
667            use std::collections::hash_map::DefaultHasher;
668            use std::hash::{Hash, Hasher};
669            let mut hasher = DefaultHasher::new();
670            user_agent.hash(&mut hasher);
671            return ClientId::UserAgent(format!("ua_{:x}", hasher.finish()));
672        }
673
674        ClientId::Anonymous
675    }
676
677    /// Extract client ID from query parameters
678    #[must_use]
679    pub fn extract_from_query(&self, query_params: &HashMap<String, String>) -> Option<ClientId> {
680        query_params
681            .get("client_id")
682            .map(|client_id| ClientId::QueryParam(client_id.clone()))
683    }
684
685    /// Extract client ID from multiple sources (with priority)
686    #[must_use]
687    pub fn extract_client_id(
688        &self,
689        headers: Option<&HashMap<String, String>>,
690        query_params: Option<&HashMap<String, String>>,
691    ) -> ClientId {
692        // Try query parameters first (highest priority)
693        if let Some(params) = query_params
694            && let Some(client_id) = self.extract_from_query(params)
695        {
696            return client_id;
697        }
698
699        // Then try headers
700        if let Some(headers) = headers {
701            return self.extract_from_http_headers(headers);
702        }
703
704        ClientId::Anonymous
705    }
706}
707
708impl Default for ClientIdExtractor {
709    fn default() -> Self {
710        Self::new()
711    }
712}
713
714/// Extension trait to add enhanced client management to `RequestContext`
715pub trait RequestContextExt {
716    /// Set client ID using `ClientId` enum
717    #[must_use]
718    fn with_enhanced_client_id(self, client_id: ClientId) -> Self;
719
720    /// Extract and set client ID from headers and query params
721    #[must_use]
722    fn extract_client_id(
723        self,
724        extractor: &ClientIdExtractor,
725        headers: Option<&HashMap<String, String>>,
726        query_params: Option<&HashMap<String, String>>,
727    ) -> Self;
728
729    /// Get the enhanced client ID
730    fn get_enhanced_client_id(&self) -> Option<ClientId>;
731}
732
733impl RequestContextExt for RequestContext {
734    fn with_enhanced_client_id(self, client_id: ClientId) -> Self {
735        self.with_client_id(client_id.as_str())
736            .with_metadata("client_id_method", client_id.auth_method())
737            .with_metadata("client_authenticated", client_id.is_authenticated())
738    }
739
740    fn extract_client_id(
741        self,
742        extractor: &ClientIdExtractor,
743        headers: Option<&HashMap<String, String>>,
744        query_params: Option<&HashMap<String, String>>,
745    ) -> Self {
746        let client_id = extractor.extract_client_id(headers, query_params);
747        self.with_enhanced_client_id(client_id)
748    }
749
750    fn get_enhanced_client_id(&self) -> Option<ClientId> {
751        self.client_id.as_ref().map(|id| {
752            let method = self
753                .get_metadata("client_id_method")
754                .and_then(|v| v.as_str())
755                .unwrap_or("header");
756
757            match method {
758                "bearer_token" => ClientId::Token(id.clone()),
759                "session_cookie" => ClientId::Session(id.clone()),
760                "query_param" => ClientId::QueryParam(id.clone()),
761                "user_agent" => ClientId::UserAgent(id.clone()),
762                "anonymous" => ClientId::Anonymous,
763                _ => ClientId::Header(id.clone()), // Default to header for "header" and unknown methods
764            }
765        })
766    }
767}
768
769#[cfg(test)]
770mod tests {
771    use super::*;
772
773    #[test]
774    fn test_request_context_creation() {
775        let ctx = RequestContext::new();
776        assert!(!ctx.request_id.is_empty());
777        assert!(ctx.user_id.is_none());
778        assert!(ctx.elapsed() < std::time::Duration::from_millis(100));
779    }
780
781    #[test]
782    fn test_request_context_builder() {
783        let ctx = RequestContext::new()
784            .with_user_id("user123")
785            .with_session_id("session456")
786            .with_metadata("key", "value");
787
788        assert_eq!(ctx.user_id, Some("user123".to_string()));
789        assert_eq!(ctx.session_id, Some("session456".to_string()));
790        assert_eq!(
791            ctx.get_metadata("key"),
792            Some(&serde_json::Value::String("value".to_string()))
793        );
794    }
795
796    #[test]
797    fn test_response_context_creation() {
798        let duration = std::time::Duration::from_millis(100);
799
800        let success_ctx = ResponseContext::success("req1", duration);
801        assert!(success_ctx.is_success());
802        assert!(!success_ctx.is_error());
803
804        let error_ctx = ResponseContext::error("req2", duration, 500, "Internal error");
805        assert!(!error_ctx.is_success());
806        assert!(error_ctx.is_error());
807        assert_eq!(error_ctx.error_info(), Some((500, "Internal error")));
808    }
809
810    #[test]
811    fn test_context_derivation() {
812        let parent_ctx = RequestContext::new()
813            .with_user_id("user123")
814            .with_metadata("key", "value");
815
816        let child_ctx = parent_ctx.derive();
817
818        // Should have new request ID
819        assert_ne!(parent_ctx.request_id, child_ctx.request_id);
820
821        // Should inherit user info and metadata
822        assert_eq!(parent_ctx.user_id, child_ctx.user_id);
823        assert_eq!(
824            parent_ctx.get_metadata("key"),
825            child_ctx.get_metadata("key")
826        );
827    }
828
829    // Tests for enhanced client management
830
831    #[test]
832    fn test_client_id_extraction() {
833        let extractor = ClientIdExtractor::new();
834
835        // Test header extraction
836        let mut headers = HashMap::new();
837        headers.insert("x-client-id".to_string(), "test-client".to_string());
838
839        let client_id = extractor.extract_from_http_headers(&headers);
840        assert_eq!(client_id, ClientId::Header("test-client".to_string()));
841        assert_eq!(client_id.as_str(), "test-client");
842        assert_eq!(client_id.auth_method(), "header");
843        assert!(!client_id.is_authenticated());
844    }
845
846    #[test]
847    fn test_bearer_token_extraction() {
848        let extractor = ClientIdExtractor::new();
849        extractor.register_token("token123".to_string(), "client-1".to_string());
850
851        let mut headers = HashMap::new();
852        headers.insert("authorization".to_string(), "Bearer token123".to_string());
853
854        let client_id = extractor.extract_from_http_headers(&headers);
855        assert_eq!(client_id, ClientId::Token("client-1".to_string()));
856        assert!(client_id.is_authenticated());
857        assert_eq!(client_id.auth_method(), "bearer_token");
858    }
859
860    #[test]
861    fn test_session_cookie_extraction() {
862        let extractor = ClientIdExtractor::new();
863
864        let mut headers = HashMap::new();
865        headers.insert(
866            "cookie".to_string(),
867            "session_id=sess123; other=value".to_string(),
868        );
869
870        let client_id = extractor.extract_from_http_headers(&headers);
871        assert_eq!(client_id, ClientId::Session("sess123".to_string()));
872        assert!(client_id.is_authenticated());
873    }
874
875    #[test]
876    fn test_user_agent_fallback() {
877        let extractor = ClientIdExtractor::new();
878
879        let mut headers = HashMap::new();
880        headers.insert("user-agent".to_string(), "TestAgent/1.0".to_string());
881
882        let client_id = extractor.extract_from_http_headers(&headers);
883        if let ClientId::UserAgent(id) = client_id {
884            assert!(id.starts_with("ua_"));
885        } else {
886            // Ensure test failure without panicking in production codepaths
887            assert!(
888                matches!(client_id, ClientId::UserAgent(_)),
889                "Expected UserAgent ClientId"
890            );
891        }
892    }
893
894    #[test]
895    fn test_client_session() {
896        let mut session = ClientSession::new("test-client".to_string(), "http".to_string());
897        assert!(!session.authenticated);
898        assert_eq!(session.request_count, 0);
899
900        session.update_activity();
901        assert_eq!(session.request_count, 1);
902
903        session.authenticate(Some("Test Client".to_string()));
904        assert!(session.authenticated);
905        assert_eq!(session.client_name, Some("Test Client".to_string()));
906
907        // Test idle detection
908        assert!(!session.is_idle(chrono::Duration::seconds(1)));
909    }
910
911    #[test]
912    fn test_request_info() {
913        let params = serde_json::json!({"param": "value"});
914        let request = RequestInfo::new("client-1".to_string(), "test_method".to_string(), params);
915
916        assert!(!request.success);
917        assert!(request.response_time_ms.is_none());
918
919        let completed = request.complete_success(150);
920        assert!(completed.success);
921        assert_eq!(completed.response_time_ms, Some(150));
922        assert_eq!(completed.status_code, Some(200));
923    }
924
925    #[test]
926    fn test_request_context_ext() {
927        let extractor = ClientIdExtractor::new();
928
929        let mut headers = HashMap::new();
930        headers.insert("x-client-id".to_string(), "test-client".to_string());
931
932        let ctx = RequestContext::new().extract_client_id(&extractor, Some(&headers), None);
933
934        assert_eq!(ctx.client_id, Some("test-client".to_string()));
935        assert_eq!(
936            ctx.get_metadata("client_id_method"),
937            Some(&serde_json::Value::String("header".to_string()))
938        );
939        assert_eq!(
940            ctx.get_metadata("client_authenticated"),
941            Some(&serde_json::Value::Bool(false))
942        );
943
944        let enhanced_id = ctx.get_enhanced_client_id();
945        assert_eq!(
946            enhanced_id,
947            Some(ClientId::Header("test-client".to_string()))
948        );
949    }
950
951    #[test]
952    fn test_request_analytics() {
953        let start = std::time::Instant::now();
954        let request = RequestInfo::new(
955            "client-123".to_string(),
956            "get_data".to_string(),
957            serde_json::json!({"filter": "active"}),
958        );
959
960        let response_time = start.elapsed().as_millis() as u64;
961        let completed = request
962            .complete_success(response_time)
963            .with_metadata("cache_hit".to_string(), serde_json::json!(true));
964
965        assert!(completed.success);
966        assert!(completed.response_time_ms.is_some());
967        assert_eq!(
968            completed.metadata.get("cache_hit"),
969            Some(&serde_json::json!(true))
970        );
971    }
972}
turbomcp_core/context.rs

turbomcp_core/
context.rs
