turbomcp_protocol/context/
request.rs1use 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#[derive(Debug, Clone)]
24pub struct ResponseContext {
25 pub request_id: String,
27
28 pub timestamp: Timestamp,
30
31 pub duration: std::time::Duration,
33
34 pub status: ResponseStatus,
36
37 pub metadata: Arc<HashMap<String, serde_json::Value>>,
39}
40
41#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
43pub enum ResponseStatus {
44 Success,
46 Error {
48 code: i32,
50 message: String,
52 },
53 Partial,
55 Cancelled,
57}
58
59#[derive(Debug, Clone, Serialize, Deserialize)]
61pub struct RequestInfo {
62 pub timestamp: DateTime<Utc>,
64 pub client_id: String,
66 pub method_name: String,
68 pub parameters: serde_json::Value,
70 pub response_time_ms: Option<u64>,
72 pub success: bool,
74 pub error_message: Option<String>,
76 pub status_code: Option<u16>,
78 pub metadata: HashMap<String, serde_json::Value>,
80}
81
82impl ResponseContext {
83 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 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 #[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 #[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 #[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 #[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 #[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
165pub trait RequestContextExt {
172 #[must_use]
175 fn with_enhanced_client_id(self, client_id: super::client::ClientId) -> Self;
176
177 #[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 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}