Skip to main content

turbomcp_protocol/context/
request.rs

1//! Request / response context types for the protocol layer.
2//!
3//! `RequestContext` is now a re-export of `turbomcp_core::RequestContext` —
4//! v3.2 unified the previously-triplicated context types into a single
5//! canonical one. The protocol-specific client-id and analytics helpers live
6//! on extension traits in this module so callers keep their existing import
7//! paths.
8//!
9//! Response-side analytics types (`ResponseContext`, `ResponseStatus`,
10//! `RequestInfo`) remain protocol-owned.
11
12use std::collections::HashMap;
13use std::sync::Arc;
14
15use chrono::{DateTime, Utc};
16use serde::{Deserialize, Serialize};
17
18use crate::types::Timestamp;
19
20pub use turbomcp_core::context::{RequestContext, TransportType};
21
22/// Context information generated after processing a request, containing response details.
23#[derive(Debug, Clone)]
24pub struct ResponseContext {
25    /// The ID of the original request this response is for.
26    pub request_id: String,
27
28    /// The timestamp when the response was generated.
29    pub timestamp: Timestamp,
30
31    /// The total time taken to process the request.
32    pub duration: std::time::Duration,
33
34    /// The status of the response (e.g., Success, Error).
35    pub status: ResponseStatus,
36
37    /// A collection of custom metadata for the response.
38    pub metadata: Arc<HashMap<String, serde_json::Value>>,
39}
40
41/// Represents the status of an MCP response.
42#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
43pub enum ResponseStatus {
44    /// The request was processed successfully.
45    Success,
46    /// An error occurred during request processing.
47    Error {
48        /// A numeric code indicating the error type.
49        code: i32,
50        /// A human-readable message describing the error.
51        message: String,
52    },
53    /// The response is partial, indicating more data will follow (for streaming).
54    Partial,
55    /// The request was cancelled before completion.
56    Cancelled,
57}
58
59/// Contains analytics information for a single request, used for monitoring and debugging.
60#[derive(Debug, Clone, Serialize, Deserialize)]
61pub struct RequestInfo {
62    /// The timestamp when the request was received.
63    pub timestamp: DateTime<Utc>,
64    /// The identifier of the client that made the request.
65    pub client_id: String,
66    /// The name of the tool or method that was called.
67    pub method_name: String,
68    /// The parameters provided in the request, potentially sanitized for privacy.
69    pub parameters: serde_json::Value,
70    /// The total time taken to generate a response, in milliseconds.
71    pub response_time_ms: Option<u64>,
72    /// A boolean indicating whether the request was successful.
73    pub success: bool,
74    /// The error message, if the request failed.
75    pub error_message: Option<String>,
76    /// The HTTP status code, if the request was handled over HTTP.
77    pub status_code: Option<u16>,
78    /// Additional custom metadata for analytics.
79    pub metadata: HashMap<String, serde_json::Value>,
80}
81
82impl ResponseContext {
83    /// Creates a new `ResponseContext` for a successful operation.
84    pub fn success(request_id: impl Into<String>, duration: std::time::Duration) -> Self {
85        Self {
86            request_id: request_id.into(),
87            timestamp: Timestamp::now(),
88            duration,
89            status: ResponseStatus::Success,
90            metadata: Arc::new(HashMap::new()),
91        }
92    }
93
94    /// Creates a new `ResponseContext` for a failed operation.
95    pub fn error(
96        request_id: impl Into<String>,
97        duration: std::time::Duration,
98        code: i32,
99        message: impl Into<String>,
100    ) -> Self {
101        Self {
102            request_id: request_id.into(),
103            timestamp: Timestamp::now(),
104            duration,
105            status: ResponseStatus::Error {
106                code,
107                message: message.into(),
108            },
109            metadata: Arc::new(HashMap::new()),
110        }
111    }
112}
113
114impl RequestInfo {
115    /// Creates a new `RequestInfo` for analytics.
116    #[must_use]
117    pub fn new(client_id: String, method_name: String, parameters: serde_json::Value) -> Self {
118        Self {
119            timestamp: Utc::now(),
120            client_id,
121            method_name,
122            parameters,
123            response_time_ms: None,
124            success: false,
125            error_message: None,
126            status_code: None,
127            metadata: HashMap::new(),
128        }
129    }
130
131    /// Marks the request as completed successfully and records the response time.
132    #[must_use]
133    pub const fn complete_success(mut self, response_time_ms: u64) -> Self {
134        self.response_time_ms = Some(response_time_ms);
135        self.success = true;
136        self.status_code = Some(200);
137        self
138    }
139
140    /// Marks the request as failed and records the response time and error message.
141    #[must_use]
142    pub fn complete_error(mut self, response_time_ms: u64, error: String) -> Self {
143        self.response_time_ms = Some(response_time_ms);
144        self.success = false;
145        self.error_message = Some(error);
146        self.status_code = Some(500);
147        self
148    }
149
150    /// Sets the HTTP status code for this request.
151    #[must_use]
152    pub const fn with_status_code(mut self, code: u16) -> Self {
153        self.status_code = Some(code);
154        self
155    }
156
157    /// Adds a key-value pair to the analytics metadata.
158    #[must_use]
159    pub fn with_metadata(mut self, key: String, value: serde_json::Value) -> Self {
160        self.metadata.insert(key, value);
161        self
162    }
163}
164
165/// Extension trait providing structured client-id handling.
166///
167/// The MCP transports capture a raw `client_id` string, but protocol-aware
168/// code often wants the richer [`super::client::ClientId`] enum (which tracks
169/// how the identity was proven — bearer token, session cookie, anonymous, etc.).
170/// This trait bridges the two.
171pub trait RequestContextExt {
172    /// Set `client_id` from a structured [`super::client::ClientId`] and record
173    /// the authentication method + authenticated flag in metadata.
174    #[must_use]
175    fn with_enhanced_client_id(self, client_id: super::client::ClientId) -> Self;
176
177    /// Extract a client ID from headers/query params and apply it via
178    /// [`Self::with_enhanced_client_id`].
179    #[must_use]
180    fn extract_client_id(
181        self,
182        extractor: &super::client::ClientIdExtractor,
183        headers: Option<&HashMap<String, String>>,
184        query_params: Option<&HashMap<String, String>>,
185    ) -> Self;
186
187    /// Rehydrate the structured [`super::client::ClientId`] from the context.
188    fn get_enhanced_client_id(&self) -> Option<super::client::ClientId>;
189}
190
191impl RequestContextExt for RequestContext {
192    fn with_enhanced_client_id(self, client_id: super::client::ClientId) -> Self {
193        self.with_client_id(client_id.as_str())
194            .with_metadata(
195                "client_id_method",
196                serde_json::Value::String(client_id.auth_method().to_string()),
197            )
198            .with_metadata(
199                "client_authenticated",
200                serde_json::Value::Bool(client_id.is_authenticated()),
201            )
202    }
203
204    fn extract_client_id(
205        self,
206        extractor: &super::client::ClientIdExtractor,
207        headers: Option<&HashMap<String, String>>,
208        query_params: Option<&HashMap<String, String>>,
209    ) -> Self {
210        let client_id = extractor.extract_client_id(headers, query_params);
211        self.with_enhanced_client_id(client_id)
212    }
213
214    fn get_enhanced_client_id(&self) -> Option<super::client::ClientId> {
215        self.client_id.as_ref().map(|id| {
216            let method = self
217                .get_metadata("client_id_method")
218                .and_then(|v| v.as_str())
219                .unwrap_or("header");
220
221            match method {
222                "bearer_token" => super::client::ClientId::Token(id.clone()),
223                "session_cookie" => super::client::ClientId::Session(id.clone()),
224                "query_param" => super::client::ClientId::QueryParam(id.clone()),
225                _ => super::client::ClientId::Header(id.clone()),
226            }
227        })
228    }
229}
230
231#[cfg(test)]
232mod tests {
233    use super::*;
234
235    #[test]
236    fn response_context_builders() {
237        let success = ResponseContext::success("req-1", std::time::Duration::from_millis(10));
238        assert_eq!(success.request_id, "req-1");
239        assert_eq!(success.status, ResponseStatus::Success);
240
241        let err =
242            ResponseContext::error("req-2", std::time::Duration::from_millis(5), -32000, "boom");
243        assert!(matches!(err.status, ResponseStatus::Error { .. }));
244    }
245
246    #[test]
247    fn request_info_lifecycle() {
248        let info = RequestInfo::new(
249            "client-1".into(),
250            "tools/list".into(),
251            serde_json::json!({}),
252        )
253        .complete_success(42)
254        .with_status_code(200)
255        .with_metadata("foo".into(), serde_json::json!("bar"));
256        assert!(info.success);
257        assert_eq!(info.response_time_ms, Some(42));
258        assert_eq!(info.metadata.get("foo"), Some(&serde_json::json!("bar")));
259    }
260}