1use std::future::Future;
7use std::pin::Pin;
8use std::sync::atomic::{AtomicU8, AtomicU32, Ordering};
9use std::sync::{Arc, Mutex};
10
11use asupersync::sync::Notify;
12use asupersync::types::{CancelReason, MAX_MASK_DEPTH};
13use asupersync::{Budget, Cx, Outcome, RegionId, TaskId, Time};
14
15#[cfg(test)]
16use asupersync::time::wall_now;
17
18use crate::{AuthContext, SessionState};
19
20const REQUEST_LEASE_UNMANAGED: u8 = 0;
21const REQUEST_LEASE_ACTIVE: u8 = 1;
22const REQUEST_LEASE_CLOSED: u8 = 2;
23const REQUEST_CANCELLATION_ACTIVE: u8 = 0;
24const REQUEST_CANCELLATION_CANCELLED: u8 = 1;
25const REQUEST_CANCELLATION_FINALIZING: u8 = 2;
26const REQUEST_AUTH_UNCOMMITTED: u8 = 0;
27const REQUEST_AUTH_ANONYMOUS: u8 = 1;
28const REQUEST_AUTH_AUTHENTICATED: u8 = 2;
29
30#[derive(Debug, Default)]
37struct McpRequestCancellationInner {
38 state: AtomicU8,
39 notify: Notify,
40}
41
42#[derive(Clone, Debug, Default)]
43#[doc(hidden)]
44pub struct McpRequestCancellation {
45 inner: Arc<McpRequestCancellationInner>,
46}
47
48impl McpRequestCancellation {
49 #[must_use]
51 pub fn new() -> Self {
52 Self::default()
53 }
54
55 pub fn cancel(&self) -> bool {
57 let cancelled = self
58 .inner
59 .state
60 .compare_exchange(
61 REQUEST_CANCELLATION_ACTIVE,
62 REQUEST_CANCELLATION_CANCELLED,
63 Ordering::AcqRel,
64 Ordering::Acquire,
65 )
66 .is_ok();
67 if cancelled {
68 self.notify_terminal_waiters();
69 }
70 cancelled
71 }
72
73 fn notify_terminal_waiters(&self) {
74 let _ = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
78 self.inner.notify.notify_waiters();
79 }));
80 }
81
82 #[must_use]
84 pub fn is_cancel_requested(&self) -> bool {
85 self.inner.state.load(Ordering::Acquire) == REQUEST_CANCELLATION_CANCELLED
86 }
87
88 #[must_use]
94 pub fn is_terminal(&self) -> bool {
95 self.inner.state.load(Ordering::Acquire) != REQUEST_CANCELLATION_ACTIVE
96 }
97
98 pub async fn cancelled(&self) {
104 self.inner
105 .notify
106 .wait_until(|| self.is_cancel_requested())
107 .await;
108 }
109
110 pub async fn terminated(&self) {
115 self.inner.notify.wait_until(|| self.is_terminal()).await;
116 }
117
118 #[must_use]
124 pub fn begin_finalization(&self) -> bool {
125 loop {
126 match self.inner.state.load(Ordering::Acquire) {
127 REQUEST_CANCELLATION_ACTIVE => {
128 if self
129 .inner
130 .state
131 .compare_exchange(
132 REQUEST_CANCELLATION_ACTIVE,
133 REQUEST_CANCELLATION_FINALIZING,
134 Ordering::AcqRel,
135 Ordering::Acquire,
136 )
137 .is_ok()
138 {
139 self.notify_terminal_waiters();
140 return true;
141 }
142 }
143 REQUEST_CANCELLATION_CANCELLED => return false,
144 REQUEST_CANCELLATION_FINALIZING => return true,
145 _ => return false,
146 }
147 }
148 }
149
150 #[must_use]
152 pub fn is_finalizing(&self) -> bool {
153 self.inner.state.load(Ordering::Acquire) == REQUEST_CANCELLATION_FINALIZING
154 }
155}
156
157pub trait NotificationSender: Send + Sync {
166 fn send_progress(&self, progress: f64, total: Option<f64>, message: Option<&str>);
174
175 fn send_progress_exact(
183 &self,
184 _progress: serde_json::Number,
185 _total: Option<serde_json::Number>,
186 _message: Option<&str>,
187 ) {
188 }
189
190 fn send_log(&self, _level: McpLogLevel, _logger: Option<&str>, _data: serde_json::Value) {}
196
197 fn send_catalog_changed(&self, _kind: McpCatalogKind) {}
203
204 fn send_resource_updated(&self, _uri: &str) {}
209}
210
211#[derive(Debug, Clone, Copy, PartialEq, Eq)]
213pub enum McpCatalogKind {
214 Tools,
215 Resources,
216 Prompts,
217}
218
219pub trait CatalogChangePublisher: Send + Sync {
222 fn publish_catalog_changed(&self, kind: McpCatalogKind) -> bool;
224 fn publish_resource_updated(&self, uri: &str) -> bool;
226}
227
228#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
234pub enum McpLogLevel {
235 Debug,
236 Info,
237 Notice,
238 Warning,
239 Error,
240 Critical,
241 Alert,
242 Emergency,
243}
244
245impl McpLogLevel {
246 #[must_use]
248 pub const fn as_str(self) -> &'static str {
249 match self {
250 Self::Debug => "debug",
251 Self::Info => "info",
252 Self::Notice => "notice",
253 Self::Warning => "warning",
254 Self::Error => "error",
255 Self::Critical => "critical",
256 Self::Alert => "alert",
257 Self::Emergency => "emergency",
258 }
259 }
260
261 #[must_use]
263 pub const fn rank(self) -> u8 {
264 match self {
265 Self::Debug => 1,
266 Self::Info => 2,
267 Self::Notice => 3,
268 Self::Warning => 4,
269 Self::Error => 5,
270 Self::Critical => 6,
271 Self::Alert => 7,
272 Self::Emergency => 8,
273 }
274 }
275}
276
277#[derive(Debug, Clone, PartialEq, Eq)]
287pub struct ClientRoot {
288 pub uri: String,
290 pub name: Option<String>,
292}
293
294impl ClientRoot {
295 #[must_use]
297 pub fn new(uri: impl Into<String>) -> Self {
298 Self {
299 uri: uri.into(),
300 name: None,
301 }
302 }
303
304 #[must_use]
306 pub fn with_name(uri: impl Into<String>, name: impl Into<String>) -> Self {
307 Self {
308 uri: uri.into(),
309 name: Some(name.into()),
310 }
311 }
312}
313
314pub trait RootsProvider: Send + Sync {
316 fn list_roots(
318 &self,
319 ) -> std::pin::Pin<
320 Box<dyn std::future::Future<Output = crate::McpResult<Vec<ClientRoot>>> + Send + '_>,
321 >;
322}
323
324pub trait SamplingSender: Send + Sync {
334 fn create_message(
345 &self,
346 request: SamplingRequest,
347 ) -> std::pin::Pin<
348 Box<dyn std::future::Future<Output = crate::McpResult<SamplingResponse>> + Send + '_>,
349 >;
350}
351
352#[derive(Debug, Clone)]
354pub struct SamplingRequest {
355 pub messages: Vec<SamplingRequestMessage>,
357 pub max_tokens: u32,
359 pub system_prompt: Option<String>,
361 pub temperature: Option<f64>,
363 pub stop_sequences: Vec<String>,
365 pub model_hints: Vec<String>,
367}
368
369impl SamplingRequest {
370 #[must_use]
372 pub fn new(messages: Vec<SamplingRequestMessage>, max_tokens: u32) -> Self {
373 Self {
374 messages,
375 max_tokens,
376 system_prompt: None,
377 temperature: None,
378 stop_sequences: Vec::new(),
379 model_hints: Vec::new(),
380 }
381 }
382
383 #[must_use]
385 pub fn prompt(text: impl Into<String>, max_tokens: u32) -> Self {
386 Self::new(vec![SamplingRequestMessage::user(text)], max_tokens)
387 }
388
389 #[must_use]
391 pub fn with_system_prompt(mut self, prompt: impl Into<String>) -> Self {
392 self.system_prompt = Some(prompt.into());
393 self
394 }
395
396 #[must_use]
398 pub fn with_temperature(mut self, temp: f64) -> Self {
399 self.temperature = Some(temp);
400 self
401 }
402
403 #[must_use]
405 pub fn with_stop_sequences(mut self, sequences: Vec<String>) -> Self {
406 self.stop_sequences = sequences;
407 self
408 }
409
410 #[must_use]
412 pub fn with_model_hints(mut self, hints: Vec<String>) -> Self {
413 self.model_hints = hints;
414 self
415 }
416}
417
418#[derive(Debug, Clone)]
420pub struct SamplingRequestMessage {
421 pub role: SamplingRole,
423 pub text: String,
425}
426
427impl SamplingRequestMessage {
428 #[must_use]
430 pub fn user(text: impl Into<String>) -> Self {
431 Self {
432 role: SamplingRole::User,
433 text: text.into(),
434 }
435 }
436
437 #[must_use]
439 pub fn assistant(text: impl Into<String>) -> Self {
440 Self {
441 role: SamplingRole::Assistant,
442 text: text.into(),
443 }
444 }
445}
446
447#[derive(Debug, Clone, Copy, PartialEq, Eq)]
449pub enum SamplingRole {
450 User,
452 Assistant,
454}
455
456#[derive(Debug, Clone)]
458pub struct SamplingResponse {
459 pub text: String,
461 pub model: String,
463 pub stop_reason: SamplingStopReason,
465}
466
467impl SamplingResponse {
468 #[must_use]
470 pub fn new(text: impl Into<String>, model: impl Into<String>) -> Self {
471 Self {
472 text: text.into(),
473 model: model.into(),
474 stop_reason: SamplingStopReason::EndTurn,
475 }
476 }
477}
478
479#[derive(Debug, Clone, PartialEq, Eq, Default)]
481pub enum SamplingStopReason {
482 #[default]
484 EndTurn,
485 StopSequence,
487 MaxTokens,
489 Unspecified,
491 Other(String),
493}
494
495impl SamplingStopReason {
496 #[must_use]
499 pub fn from_wire_value(value: Option<String>) -> Self {
500 match value {
501 Some(value) => match value.as_str() {
502 "endTurn" => Self::EndTurn,
503 "stopSequence" => Self::StopSequence,
504 "maxTokens" => Self::MaxTokens,
505 _ => Self::Other(value),
506 },
507 None => Self::Unspecified,
508 }
509 }
510
511 #[must_use]
514 pub fn as_wire_value(&self) -> Option<&str> {
515 match self {
516 Self::EndTurn => Some("endTurn"),
517 Self::StopSequence => Some("stopSequence"),
518 Self::MaxTokens => Some("maxTokens"),
519 Self::Unspecified => None,
520 Self::Other(value) => Some(value),
521 }
522 }
523}
524
525#[derive(Debug, Clone, Copy, Default)]
529pub struct NoOpSamplingSender;
530
531impl SamplingSender for NoOpSamplingSender {
532 fn create_message(
533 &self,
534 _request: SamplingRequest,
535 ) -> std::pin::Pin<
536 Box<dyn std::future::Future<Output = crate::McpResult<SamplingResponse>> + Send + '_>,
537 > {
538 Box::pin(async {
539 Err(crate::McpError::new(
540 crate::McpErrorCode::InvalidRequest,
541 "Sampling not supported: client does not have sampling capability",
542 ))
543 })
544 }
545}
546
547pub trait ElicitationSender: Send + Sync {
557 fn elicit(
568 &self,
569 request: ElicitationRequest,
570 ) -> std::pin::Pin<
571 Box<dyn std::future::Future<Output = crate::McpResult<ElicitationResponse>> + Send + '_>,
572 >;
573}
574
575#[derive(Debug, Clone)]
577pub struct ElicitationRequest {
578 pub mode: ElicitationMode,
580 pub message: String,
582 pub schema: Option<serde_json::Value>,
584 pub url: Option<String>,
586 pub elicitation_id: Option<String>,
588}
589
590impl ElicitationRequest {
591 #[must_use]
593 pub fn form(message: impl Into<String>, schema: serde_json::Value) -> Self {
594 Self {
595 mode: ElicitationMode::Form,
596 message: message.into(),
597 schema: Some(schema),
598 url: None,
599 elicitation_id: None,
600 }
601 }
602
603 #[must_use]
605 pub fn url(
606 message: impl Into<String>,
607 url: impl Into<String>,
608 elicitation_id: impl Into<String>,
609 ) -> Self {
610 Self {
611 mode: ElicitationMode::Url,
612 message: message.into(),
613 schema: None,
614 url: Some(url.into()),
615 elicitation_id: Some(elicitation_id.into()),
616 }
617 }
618}
619
620#[derive(Debug, Clone, Copy, PartialEq, Eq)]
622pub enum ElicitationMode {
623 Form,
625 Url,
627}
628
629#[derive(Debug, Clone)]
631pub struct ElicitationResponse {
632 pub action: ElicitationAction,
634 pub content: Option<std::collections::HashMap<String, serde_json::Value>>,
636}
637
638impl ElicitationResponse {
639 #[must_use]
641 pub fn accept(content: std::collections::HashMap<String, serde_json::Value>) -> Self {
642 Self {
643 action: ElicitationAction::Accept,
644 content: Some(content),
645 }
646 }
647
648 #[must_use]
650 pub fn accept_url() -> Self {
651 Self {
652 action: ElicitationAction::Accept,
653 content: None,
654 }
655 }
656
657 #[must_use]
659 pub fn decline() -> Self {
660 Self {
661 action: ElicitationAction::Decline,
662 content: None,
663 }
664 }
665
666 #[must_use]
668 pub fn cancel() -> Self {
669 Self {
670 action: ElicitationAction::Cancel,
671 content: None,
672 }
673 }
674
675 #[must_use]
677 pub fn is_accepted(&self) -> bool {
678 matches!(self.action, ElicitationAction::Accept)
679 }
680
681 #[must_use]
683 pub fn is_declined(&self) -> bool {
684 matches!(self.action, ElicitationAction::Decline)
685 }
686
687 #[must_use]
689 pub fn is_cancelled(&self) -> bool {
690 matches!(self.action, ElicitationAction::Cancel)
691 }
692
693 #[must_use]
695 pub fn get_string(&self, key: &str) -> Option<&str> {
696 self.content.as_ref()?.get(key)?.as_str()
697 }
698
699 #[must_use]
701 pub fn get_bool(&self, key: &str) -> Option<bool> {
702 self.content.as_ref()?.get(key)?.as_bool()
703 }
704
705 #[must_use]
707 pub fn get_int(&self, key: &str) -> Option<i64> {
708 self.content.as_ref()?.get(key)?.as_i64()
709 }
710}
711
712#[derive(Debug, Clone, Copy, PartialEq, Eq)]
714pub enum ElicitationAction {
715 Accept,
717 Decline,
719 Cancel,
721}
722
723#[derive(Debug, Clone, Copy, Default)]
727pub struct NoOpElicitationSender;
728
729impl ElicitationSender for NoOpElicitationSender {
730 fn elicit(
731 &self,
732 _request: ElicitationRequest,
733 ) -> std::pin::Pin<
734 Box<dyn std::future::Future<Output = crate::McpResult<ElicitationResponse>> + Send + '_>,
735 > {
736 Box::pin(async {
737 Err(crate::McpError::new(
738 crate::McpErrorCode::InvalidRequest,
739 "Elicitation not supported: client does not have elicitation capability",
740 ))
741 })
742 }
743}
744
745pub const MAX_RESOURCE_READ_DEPTH: u32 = 10;
751
752#[derive(Debug, Clone)]
757pub struct ResourceContentItem {
758 pub uri: String,
760 pub mime_type: Option<String>,
762 pub text: Option<String>,
764 pub blob: Option<String>,
766}
767
768impl ResourceContentItem {
769 #[must_use]
771 pub fn text(uri: impl Into<String>, text: impl Into<String>) -> Self {
772 Self {
773 uri: uri.into(),
774 mime_type: Some("text/plain".to_string()),
775 text: Some(text.into()),
776 blob: None,
777 }
778 }
779
780 #[must_use]
782 pub fn json(uri: impl Into<String>, text: impl Into<String>) -> Self {
783 Self {
784 uri: uri.into(),
785 mime_type: Some("application/json".to_string()),
786 text: Some(text.into()),
787 blob: None,
788 }
789 }
790
791 #[must_use]
793 pub fn blob(
794 uri: impl Into<String>,
795 mime_type: impl Into<String>,
796 blob: impl Into<String>,
797 ) -> Self {
798 Self {
799 uri: uri.into(),
800 mime_type: Some(mime_type.into()),
801 text: None,
802 blob: Some(blob.into()),
803 }
804 }
805
806 #[must_use]
808 pub fn as_text(&self) -> Option<&str> {
809 self.text.as_deref()
810 }
811
812 #[must_use]
814 pub fn as_blob(&self) -> Option<&str> {
815 self.blob.as_deref()
816 }
817
818 #[must_use]
820 pub fn is_text(&self) -> bool {
821 self.text.is_some()
822 }
823
824 #[must_use]
826 pub fn is_blob(&self) -> bool {
827 self.blob.is_some()
828 }
829}
830
831#[derive(Debug, Clone)]
833pub struct ResourceReadResult {
834 pub contents: Vec<ResourceContentItem>,
836}
837
838impl ResourceReadResult {
839 #[must_use]
841 pub fn new(contents: Vec<ResourceContentItem>) -> Self {
842 Self { contents }
843 }
844
845 #[must_use]
847 pub fn text(uri: impl Into<String>, text: impl Into<String>) -> Self {
848 Self {
849 contents: vec![ResourceContentItem::text(uri, text)],
850 }
851 }
852
853 #[must_use]
855 pub fn first_text(&self) -> Option<&str> {
856 self.contents.first().and_then(|c| c.as_text())
857 }
858
859 #[must_use]
861 pub fn first_blob(&self) -> Option<&str> {
862 self.contents.first().and_then(|c| c.as_blob())
863 }
864}
865
866pub trait ResourceReader: Send + Sync {
875 fn read_resource<'a>(
888 &'a self,
889 context: &'a McpContext,
890 uri: &'a str,
891 depth: u32,
892 ) -> Pin<Box<dyn Future<Output = crate::McpResult<ResourceReadResult>> + Send + 'a>>;
893}
894
895pub const MAX_TOOL_CALL_DEPTH: u32 = 10;
901
902#[derive(Debug, Clone)]
907pub enum ToolContentItem {
908 Text {
910 text: String,
912 },
913 Image {
915 data: String,
917 mime_type: String,
919 },
920 Audio {
922 data: String,
924 mime_type: String,
926 },
927 Resource {
929 uri: String,
931 mime_type: Option<String>,
933 text: Option<String>,
935 blob: Option<String>,
937 },
938}
939
940impl ToolContentItem {
941 #[must_use]
943 pub fn text(text: impl Into<String>) -> Self {
944 Self::Text { text: text.into() }
945 }
946
947 #[must_use]
949 pub fn as_text(&self) -> Option<&str> {
950 match self {
951 Self::Text { text } => Some(text),
952 _ => None,
953 }
954 }
955
956 #[must_use]
958 pub fn is_text(&self) -> bool {
959 matches!(self, Self::Text { .. })
960 }
961}
962
963#[derive(Debug, Clone)]
965pub struct ToolCallResult {
966 pub content: Vec<ToolContentItem>,
968 pub is_error: bool,
970}
971
972impl ToolCallResult {
973 #[must_use]
975 pub fn success(content: Vec<ToolContentItem>) -> Self {
976 Self {
977 content,
978 is_error: false,
979 }
980 }
981
982 #[must_use]
984 pub fn text(text: impl Into<String>) -> Self {
985 Self {
986 content: vec![ToolContentItem::text(text)],
987 is_error: false,
988 }
989 }
990
991 #[must_use]
993 pub fn error(message: impl Into<String>) -> Self {
994 Self {
995 content: vec![ToolContentItem::text(message)],
996 is_error: true,
997 }
998 }
999
1000 #[must_use]
1002 pub fn first_text(&self) -> Option<&str> {
1003 self.content.first().and_then(|c| c.as_text())
1004 }
1005}
1006
1007pub trait ToolCaller: Send + Sync {
1016 fn call_tool<'a>(
1030 &'a self,
1031 context: &'a McpContext,
1032 name: &'a str,
1033 args: serde_json::Value,
1034 depth: u32,
1035 ) -> Pin<Box<dyn Future<Output = crate::McpResult<ToolCallResult>> + Send + 'a>>;
1036}
1037
1038pub const MAX_PROMPT_GET_DEPTH: u32 = 10;
1044
1045#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1047pub enum PromptMessageRole {
1048 User,
1050 Assistant,
1052}
1053
1054#[derive(Debug, Clone, PartialEq, Eq)]
1056pub struct PromptMessageItem {
1057 pub role: PromptMessageRole,
1059 pub text: Option<String>,
1061}
1062
1063impl PromptMessageItem {
1064 #[must_use]
1066 pub fn user_text(text: impl Into<String>) -> Self {
1067 Self {
1068 role: PromptMessageRole::User,
1069 text: Some(text.into()),
1070 }
1071 }
1072
1073 #[must_use]
1075 pub fn as_text(&self) -> Option<&str> {
1076 self.text.as_deref()
1077 }
1078}
1079
1080#[derive(Debug, Clone, PartialEq, Eq)]
1082pub struct PromptGetResult {
1083 pub description: Option<String>,
1085 pub messages: Vec<PromptMessageItem>,
1087}
1088
1089impl PromptGetResult {
1090 #[must_use]
1092 pub fn new(messages: Vec<PromptMessageItem>) -> Self {
1093 Self {
1094 description: None,
1095 messages,
1096 }
1097 }
1098
1099 #[must_use]
1101 pub fn first_text(&self) -> Option<&str> {
1102 self.messages.iter().find_map(PromptMessageItem::as_text)
1103 }
1104}
1105
1106pub trait PromptCaller: Send + Sync {
1112 fn get_prompt<'a>(
1114 &'a self,
1115 context: &'a McpContext,
1116 name: &'a str,
1117 arguments: std::collections::HashMap<String, String>,
1118 depth: u32,
1119 ) -> Pin<Box<dyn Future<Output = crate::McpResult<PromptGetResult>> + Send + 'a>>;
1120}
1121
1122#[derive(Debug, Clone, Default)]
1131pub struct ClientCapabilityInfo {
1132 pub sampling: bool,
1134 pub elicitation: bool,
1136 pub elicitation_form: bool,
1138 pub elicitation_url: bool,
1140 pub roots: bool,
1142 pub roots_list_changed: bool,
1144}
1145
1146impl ClientCapabilityInfo {
1147 #[must_use]
1149 pub fn new() -> Self {
1150 Self::default()
1151 }
1152
1153 #[must_use]
1155 pub fn with_sampling(mut self) -> Self {
1156 self.sampling = true;
1157 self
1158 }
1159
1160 #[must_use]
1162 pub fn with_elicitation(mut self, form: bool, url: bool) -> Self {
1163 self.elicitation = form || url;
1164 self.elicitation_form = form;
1165 self.elicitation_url = url;
1166 self
1167 }
1168
1169 #[must_use]
1171 pub fn with_roots(mut self, list_changed: bool) -> Self {
1172 self.roots = true;
1173 self.roots_list_changed = list_changed;
1174 self
1175 }
1176}
1177
1178#[derive(Debug, Clone, PartialEq, Eq)]
1183pub struct ClientImplementationInfo {
1184 pub name: String,
1186 pub version: String,
1188 pub title: Option<String>,
1190 pub description: Option<String>,
1192 pub website_url: Option<String>,
1194 pub icon_sources: Vec<String>,
1196}
1197
1198impl ClientImplementationInfo {
1199 #[must_use]
1201 pub fn new(name: impl Into<String>, version: impl Into<String>) -> Self {
1202 Self {
1203 name: name.into(),
1204 version: version.into(),
1205 title: None,
1206 description: None,
1207 website_url: None,
1208 icon_sources: Vec::new(),
1209 }
1210 }
1211
1212 #[must_use]
1214 pub fn has_extras(&self) -> bool {
1215 self.title.is_some()
1216 || self.description.is_some()
1217 || self.website_url.is_some()
1218 || !self.icon_sources.is_empty()
1219 }
1220}
1221
1222#[derive(Debug, Clone, Default)]
1226pub struct ServerCapabilityInfo {
1227 pub tools: bool,
1229 pub resources: bool,
1231 pub resources_subscribe: bool,
1233 pub prompts: bool,
1235 pub logging: bool,
1237}
1238
1239impl ServerCapabilityInfo {
1240 #[must_use]
1242 pub fn new() -> Self {
1243 Self::default()
1244 }
1245
1246 #[must_use]
1248 pub fn with_tools(mut self) -> Self {
1249 self.tools = true;
1250 self
1251 }
1252
1253 #[must_use]
1255 pub fn with_resources(mut self, subscribe: bool) -> Self {
1256 self.resources = true;
1257 self.resources_subscribe = subscribe;
1258 self
1259 }
1260
1261 #[must_use]
1263 pub fn with_prompts(mut self) -> Self {
1264 self.prompts = true;
1265 self
1266 }
1267
1268 #[must_use]
1270 pub fn with_logging(mut self) -> Self {
1271 self.logging = true;
1272 self
1273 }
1274}
1275
1276#[derive(Debug, Clone, Copy, Default)]
1278pub struct NoOpNotificationSender;
1279
1280impl NotificationSender for NoOpNotificationSender {
1281 fn send_progress(&self, _progress: f64, _total: Option<f64>, _message: Option<&str>) {
1282 }
1284}
1285
1286#[derive(Clone)]
1291pub struct ProgressReporter {
1292 sender: Arc<dyn NotificationSender>,
1293 marker: Option<serde_json::Value>,
1298}
1299
1300impl ProgressReporter {
1301 pub fn new(sender: Arc<dyn NotificationSender>) -> Self {
1303 Self {
1304 sender,
1305 marker: None,
1306 }
1307 }
1308
1309 #[must_use]
1314 pub fn with_marker(marker: serde_json::Value, sender: Arc<dyn NotificationSender>) -> Self {
1315 Self {
1316 sender,
1317 marker: Some(marker),
1318 }
1319 }
1320
1321 #[must_use]
1323 pub fn marker(&self) -> Option<&serde_json::Value> {
1324 self.marker.as_ref()
1325 }
1326
1327 pub fn report(&self, progress: f64, message: Option<&str>) {
1334 self.sender.send_progress(progress, None, message);
1335 }
1336
1337 pub fn report_with_total(&self, progress: f64, total: f64, message: Option<&str>) {
1345 self.sender.send_progress(progress, Some(total), message);
1346 }
1347
1348 pub fn report_exact(
1355 &self,
1356 progress: serde_json::Number,
1357 total: Option<serde_json::Number>,
1358 message: Option<&str>,
1359 ) {
1360 self.sender.send_progress_exact(progress, total, message);
1361 }
1362}
1363
1364impl std::fmt::Debug for ProgressReporter {
1365 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1366 f.debug_struct("ProgressReporter").finish_non_exhaustive()
1367 }
1368}
1369
1370#[derive(Clone)]
1413pub struct McpContext {
1414 cx: Cx,
1416 budget_state: Arc<Mutex<FrameworkBudgetState>>,
1421 framework_mask_depth: Arc<AtomicU32>,
1424 mask_transition: Arc<Mutex<()>>,
1430 operation_deadline: Option<Time>,
1437 request_lease: Arc<AtomicU8>,
1445 request_cancellation: McpRequestCancellation,
1447 request_id: u64,
1449 final_request_surface: bool,
1454 progress_reporter: Option<ProgressReporter>,
1456 state: Option<SessionState>,
1458 cache_admission_partition: Arc<Mutex<Option<([u8; 32], u64)>>>,
1460 response_cache_hits: Arc<Mutex<Vec<u64>>>,
1462 auth: Arc<Mutex<Option<AuthContext>>>,
1464 auth_state: Arc<AtomicU8>,
1467 sampling_sender: Option<Arc<dyn SamplingSender>>,
1469 elicitation_sender: Option<Arc<dyn ElicitationSender>>,
1471 roots_provider: Option<Arc<dyn RootsProvider>>,
1473 resource_reader: Option<Arc<dyn ResourceReader>>,
1475 resource_read_depth: u32,
1477 tool_caller: Option<Arc<dyn ToolCaller>>,
1479 tool_call_depth: u32,
1481 prompt_caller: Option<Arc<dyn PromptCaller>>,
1483 prompt_get_depth: u32,
1485 client_capabilities: Option<ClientCapabilityInfo>,
1487 client_implementation: Option<ClientImplementationInfo>,
1489 server_capabilities: Option<ServerCapabilityInfo>,
1491 log_sender: Option<Arc<dyn NotificationSender>>,
1493 min_log_level: Option<McpLogLevel>,
1498 resource_subscriptions: Option<Arc<std::collections::HashSet<String>>>,
1500 catalog_publisher: Option<Arc<dyn CatalogChangePublisher>>,
1502}
1503
1504impl std::fmt::Debug for McpContext {
1505 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1506 let budget_state = *self
1507 .budget_state
1508 .lock()
1509 .unwrap_or_else(std::sync::PoisonError::into_inner);
1510 f.debug_struct("McpContext")
1511 .field("cx", &self.cx)
1512 .field("budget_ceiling", &budget_state.ceiling)
1513 .field("ambient_poll_debits", &budget_state.ambient_poll_debits)
1514 .field("ambient_cost_debits", &budget_state.ambient_cost_debits)
1515 .field("deferred_overrun", &budget_state.deferred_overrun)
1516 .field(
1517 "framework_mask_depth",
1518 &self.framework_mask_depth.load(Ordering::Relaxed),
1519 )
1520 .field("operation_deadline", &self.operation_deadline)
1521 .field("request_lease_active", &self.request_scope_is_active())
1522 .field(
1523 "request_cancel_requested",
1524 &self.request_cancellation.is_cancel_requested(),
1525 )
1526 .field("request_id", &self.request_id)
1527 .field("final_request_surface", &self.final_request_surface)
1528 .field("progress_reporter", &self.progress_reporter)
1529 .field("state", &self.state.is_some())
1530 .field(
1531 "cache_admission_partition",
1532 &self
1533 .cache_admission_partition
1534 .lock()
1535 .unwrap_or_else(std::sync::PoisonError::into_inner)
1536 .is_some(),
1537 )
1538 .field(
1539 "response_cache_hit_count",
1540 &self
1541 .response_cache_hits
1542 .lock()
1543 .unwrap_or_else(std::sync::PoisonError::into_inner)
1544 .len(),
1545 )
1546 .field(
1547 "auth",
1548 &self
1549 .auth
1550 .lock()
1551 .unwrap_or_else(std::sync::PoisonError::into_inner)
1552 .is_some(),
1553 )
1554 .field(
1555 "auth_committed",
1556 &(self.auth_state.load(Ordering::Acquire) != REQUEST_AUTH_UNCOMMITTED),
1557 )
1558 .field("sampling_sender", &self.sampling_sender.is_some())
1559 .field("elicitation_sender", &self.elicitation_sender.is_some())
1560 .field("roots_provider", &self.roots_provider.is_some())
1561 .field("resource_reader", &self.resource_reader.is_some())
1562 .field("resource_read_depth", &self.resource_read_depth)
1563 .field("tool_caller", &self.tool_caller.is_some())
1564 .field("tool_call_depth", &self.tool_call_depth)
1565 .field("prompt_caller", &self.prompt_caller.is_some())
1566 .field("prompt_get_depth", &self.prompt_get_depth)
1567 .field("client_capabilities", &self.client_capabilities)
1568 .field("client_implementation", &self.client_implementation)
1569 .field("server_capabilities", &self.server_capabilities)
1570 .field("log_sender", &self.log_sender.is_some())
1571 .field("min_log_level", &self.min_log_level)
1572 .field(
1573 "resource_subscription_count",
1574 &self
1575 .resource_subscriptions
1576 .as_ref()
1577 .map_or(0, |uris| uris.len()),
1578 )
1579 .field("catalog_publisher", &self.catalog_publisher.is_some())
1580 .finish()
1581 }
1582}
1583
1584#[derive(Clone, Copy, Debug, Default)]
1585struct FrameworkBudgetState {
1586 ceiling: Option<Budget>,
1587 ambient_poll_debits: u32,
1590 ambient_cost_debits: u64,
1597 deferred_overrun: bool,
1601}
1602
1603impl FrameworkBudgetState {
1604 fn adjusted_ambient(self, mut ambient: Budget) -> Budget {
1605 if ambient.poll_quota != u32::MAX {
1606 ambient.poll_quota = ambient.poll_quota.saturating_sub(self.ambient_poll_debits);
1607 }
1608 if let Some(remaining) = ambient.cost_quota.as_mut() {
1609 *remaining = remaining.saturating_sub(self.ambient_cost_debits);
1610 }
1611 ambient
1612 }
1613
1614 fn effective(self, ambient: Budget) -> Budget {
1615 let ambient = self.adjusted_ambient(ambient);
1616 self.ceiling
1617 .map_or(ambient, |ceiling| ambient.meet(ceiling))
1618 }
1619}
1620
1621struct FrameworkMaskGuard<'a> {
1622 depth: &'a AtomicU32,
1623}
1624
1625#[doc(hidden)]
1630pub struct McpContextLeaseGuard {
1631 lease: Arc<AtomicU8>,
1632}
1633
1634impl std::fmt::Debug for McpContextLeaseGuard {
1635 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1636 f.debug_struct("McpContextLeaseGuard")
1637 .field(
1638 "active",
1639 &(self.lease.load(Ordering::Acquire) == REQUEST_LEASE_ACTIVE),
1640 )
1641 .finish()
1642 }
1643}
1644
1645impl Drop for McpContextLeaseGuard {
1646 fn drop(&mut self) {
1647 self.lease.store(REQUEST_LEASE_CLOSED, Ordering::Release);
1648 }
1649}
1650
1651impl Drop for FrameworkMaskGuard<'_> {
1652 fn drop(&mut self) {
1653 self.depth.fetch_sub(1, Ordering::SeqCst);
1654 }
1655}
1656
1657impl McpContext {
1658 #[must_use]
1666 pub fn new(cx: Cx, request_id: u64) -> Self {
1667 Self {
1668 cx,
1669 budget_state: Arc::new(Mutex::new(FrameworkBudgetState::default())),
1670 framework_mask_depth: Arc::new(AtomicU32::new(0)),
1671 mask_transition: Arc::new(Mutex::new(())),
1672 operation_deadline: None,
1673 request_lease: Arc::new(AtomicU8::new(REQUEST_LEASE_UNMANAGED)),
1674 request_cancellation: McpRequestCancellation::new(),
1675 request_id,
1676 final_request_surface: false,
1677 progress_reporter: None,
1678 state: None,
1679 cache_admission_partition: Arc::new(Mutex::new(None)),
1680 response_cache_hits: Arc::new(Mutex::new(Vec::new())),
1681 auth: Arc::new(Mutex::new(None)),
1682 auth_state: Arc::new(AtomicU8::new(REQUEST_AUTH_UNCOMMITTED)),
1683 sampling_sender: None,
1684 elicitation_sender: None,
1685 roots_provider: None,
1686 resource_reader: None,
1687 resource_read_depth: 0,
1688 tool_caller: None,
1689 tool_call_depth: 0,
1690 prompt_caller: None,
1691 prompt_get_depth: 0,
1692 client_capabilities: None,
1693 client_implementation: None,
1694 server_capabilities: None,
1695 log_sender: None,
1696 min_log_level: None,
1697 resource_subscriptions: None,
1698 catalog_publisher: None,
1699 }
1700 }
1701
1702 #[must_use]
1708 pub fn with_state(cx: Cx, request_id: u64, state: SessionState) -> Self {
1709 Self {
1710 cx,
1711 budget_state: Arc::new(Mutex::new(FrameworkBudgetState::default())),
1712 framework_mask_depth: Arc::new(AtomicU32::new(0)),
1713 mask_transition: Arc::new(Mutex::new(())),
1714 operation_deadline: None,
1715 request_lease: Arc::new(AtomicU8::new(REQUEST_LEASE_UNMANAGED)),
1716 request_cancellation: McpRequestCancellation::new(),
1717 request_id,
1718 final_request_surface: false,
1719 progress_reporter: None,
1720 state: Some(state),
1721 cache_admission_partition: Arc::new(Mutex::new(None)),
1722 response_cache_hits: Arc::new(Mutex::new(Vec::new())),
1723 auth: Arc::new(Mutex::new(None)),
1724 auth_state: Arc::new(AtomicU8::new(REQUEST_AUTH_UNCOMMITTED)),
1725 sampling_sender: None,
1726 elicitation_sender: None,
1727 roots_provider: None,
1728 resource_reader: None,
1729 resource_read_depth: 0,
1730 tool_caller: None,
1731 tool_call_depth: 0,
1732 prompt_caller: None,
1733 prompt_get_depth: 0,
1734 client_capabilities: None,
1735 client_implementation: None,
1736 server_capabilities: None,
1737 log_sender: None,
1738 min_log_level: None,
1739 resource_subscriptions: None,
1740 catalog_publisher: None,
1741 }
1742 }
1743
1744 #[must_use]
1751 pub fn with_progress(cx: Cx, request_id: u64, reporter: ProgressReporter) -> Self {
1752 Self {
1753 cx,
1754 budget_state: Arc::new(Mutex::new(FrameworkBudgetState::default())),
1755 framework_mask_depth: Arc::new(AtomicU32::new(0)),
1756 mask_transition: Arc::new(Mutex::new(())),
1757 operation_deadline: None,
1758 request_lease: Arc::new(AtomicU8::new(REQUEST_LEASE_UNMANAGED)),
1759 request_cancellation: McpRequestCancellation::new(),
1760 request_id,
1761 final_request_surface: false,
1762 progress_reporter: Some(reporter),
1763 state: None,
1764 cache_admission_partition: Arc::new(Mutex::new(None)),
1765 response_cache_hits: Arc::new(Mutex::new(Vec::new())),
1766 auth: Arc::new(Mutex::new(None)),
1767 auth_state: Arc::new(AtomicU8::new(REQUEST_AUTH_UNCOMMITTED)),
1768 sampling_sender: None,
1769 elicitation_sender: None,
1770 roots_provider: None,
1771 resource_reader: None,
1772 resource_read_depth: 0,
1773 tool_caller: None,
1774 tool_call_depth: 0,
1775 prompt_caller: None,
1776 prompt_get_depth: 0,
1777 client_capabilities: None,
1778 client_implementation: None,
1779 server_capabilities: None,
1780 log_sender: None,
1781 min_log_level: None,
1782 resource_subscriptions: None,
1783 catalog_publisher: None,
1784 }
1785 }
1786
1787 #[must_use]
1793 pub fn with_state_and_progress(
1794 cx: Cx,
1795 request_id: u64,
1796 state: SessionState,
1797 reporter: ProgressReporter,
1798 ) -> Self {
1799 Self {
1800 cx,
1801 budget_state: Arc::new(Mutex::new(FrameworkBudgetState::default())),
1802 framework_mask_depth: Arc::new(AtomicU32::new(0)),
1803 mask_transition: Arc::new(Mutex::new(())),
1804 operation_deadline: None,
1805 request_lease: Arc::new(AtomicU8::new(REQUEST_LEASE_UNMANAGED)),
1806 request_cancellation: McpRequestCancellation::new(),
1807 request_id,
1808 final_request_surface: false,
1809 progress_reporter: Some(reporter),
1810 state: Some(state),
1811 cache_admission_partition: Arc::new(Mutex::new(None)),
1812 response_cache_hits: Arc::new(Mutex::new(Vec::new())),
1813 auth: Arc::new(Mutex::new(None)),
1814 auth_state: Arc::new(AtomicU8::new(REQUEST_AUTH_UNCOMMITTED)),
1815 sampling_sender: None,
1816 elicitation_sender: None,
1817 roots_provider: None,
1818 resource_reader: None,
1819 resource_read_depth: 0,
1820 tool_caller: None,
1821 tool_call_depth: 0,
1822 prompt_caller: None,
1823 prompt_get_depth: 0,
1824 client_capabilities: None,
1825 client_implementation: None,
1826 server_capabilities: None,
1827 log_sender: None,
1828 min_log_level: None,
1829 resource_subscriptions: None,
1830 catalog_publisher: None,
1831 }
1832 }
1833
1834 #[must_use]
1841 pub fn with_progress_reporter(mut self, reporter: ProgressReporter) -> Self {
1842 self.progress_reporter = Some(reporter);
1843 self
1844 }
1845
1846 #[must_use]
1848 pub fn with_log_sender(mut self, sender: Arc<dyn NotificationSender>) -> Self {
1849 self.log_sender = Some(sender);
1850 self
1851 }
1852
1853 #[must_use]
1858 pub fn with_min_log_level(mut self, level: Option<McpLogLevel>) -> Self {
1859 self.min_log_level = level;
1860 self
1861 }
1862
1863 #[must_use]
1867 pub fn min_log_level(&self) -> Option<McpLogLevel> {
1868 self.min_log_level
1869 }
1870
1871 #[must_use]
1873 pub fn with_resource_subscriptions(
1874 mut self,
1875 uris: impl IntoIterator<Item = impl Into<String>>,
1876 ) -> Self {
1877 self.resource_subscriptions = Some(Arc::new(uris.into_iter().map(Into::into).collect()));
1878 self
1879 }
1880
1881 #[must_use]
1883 pub fn with_catalog_publisher(mut self, publisher: Arc<dyn CatalogChangePublisher>) -> Self {
1884 self.catalog_publisher = Some(publisher);
1885 self
1886 }
1887
1888 #[must_use]
1893 pub fn with_sampling(mut self, sender: Arc<dyn SamplingSender>) -> Self {
1894 self.sampling_sender = Some(sender);
1895 self
1896 }
1897
1898 #[must_use]
1903 pub fn with_elicitation(mut self, sender: Arc<dyn ElicitationSender>) -> Self {
1904 self.elicitation_sender = Some(sender);
1905 self
1906 }
1907
1908 #[must_use]
1912 pub fn with_roots_provider(mut self, provider: Arc<dyn RootsProvider>) -> Self {
1913 self.roots_provider = Some(provider);
1914 self
1915 }
1916
1917 #[must_use]
1926 pub fn with_budget_ceiling(self, ceiling: Budget) -> Self {
1927 {
1928 let mut current = self
1929 .budget_state
1930 .lock()
1931 .unwrap_or_else(std::sync::PoisonError::into_inner);
1932 current.ceiling = Some(
1933 current
1934 .ceiling
1935 .map_or(ceiling, |budget| budget.meet(ceiling)),
1936 );
1937 }
1938 self
1939 }
1940
1941 #[must_use]
1949 pub fn with_operation_deadline(mut self, deadline: Option<Time>) -> Self {
1950 if let Some(deadline) = deadline {
1951 self.operation_deadline = Some(
1952 self.operation_deadline
1953 .map_or(deadline, |current| current.min(deadline)),
1954 );
1955 }
1956 self
1957 }
1958
1959 #[doc(hidden)]
1964 #[must_use]
1965 pub fn with_request_cancellation(mut self, cancellation: McpRequestCancellation) -> Self {
1966 if self.request_lease.load(Ordering::Acquire) == REQUEST_LEASE_UNMANAGED {
1970 self.request_cancellation = cancellation;
1971 }
1972 self
1973 }
1974
1975 #[doc(hidden)]
1985 #[must_use]
1986 pub fn begin_request_scope(self) -> Option<(Self, McpContextLeaseGuard)> {
1987 if self
1988 .request_lease
1989 .compare_exchange(
1990 REQUEST_LEASE_UNMANAGED,
1991 REQUEST_LEASE_ACTIVE,
1992 Ordering::AcqRel,
1993 Ordering::Acquire,
1994 )
1995 .is_err()
1996 {
1997 return None;
1998 }
1999 let guard = McpContextLeaseGuard {
2000 lease: Arc::clone(&self.request_lease),
2001 };
2002 Some((self, guard))
2003 }
2004
2005 #[must_use]
2010 pub fn with_resource_reader(mut self, reader: Arc<dyn ResourceReader>) -> Self {
2011 self.resource_reader = Some(reader);
2012 self
2013 }
2014
2015 #[must_use]
2020 pub fn with_resource_read_depth(mut self, depth: u32) -> Self {
2021 self.resource_read_depth = self.resource_read_depth.max(depth);
2022 self
2023 }
2024
2025 #[must_use]
2030 pub fn with_tool_caller(mut self, caller: Arc<dyn ToolCaller>) -> Self {
2031 self.tool_caller = Some(caller);
2032 self
2033 }
2034
2035 #[must_use]
2040 pub fn with_tool_call_depth(mut self, depth: u32) -> Self {
2041 self.tool_call_depth = self.tool_call_depth.max(depth);
2042 self
2043 }
2044
2045 #[must_use]
2050 pub fn with_prompt_caller(mut self, caller: Arc<dyn PromptCaller>) -> Self {
2051 self.prompt_caller = Some(caller);
2052 self
2053 }
2054
2055 #[must_use]
2060 pub fn with_final_request_surface(mut self, final_surface: bool) -> Self {
2061 self.final_request_surface = final_surface;
2062 self
2063 }
2064
2065 #[must_use]
2067 pub fn is_final_request_surface(&self) -> bool {
2068 self.final_request_surface
2069 }
2070
2071 #[must_use]
2076 pub fn with_prompt_get_depth(mut self, depth: u32) -> Self {
2077 self.prompt_get_depth = self.prompt_get_depth.max(depth);
2078 self
2079 }
2080
2081 #[must_use]
2086 pub fn with_client_capabilities(mut self, capabilities: ClientCapabilityInfo) -> Self {
2087 self.client_capabilities = Some(capabilities);
2088 self
2089 }
2090
2091 #[must_use]
2093 pub fn with_client_implementation(mut self, identity: ClientImplementationInfo) -> Self {
2094 self.client_implementation = Some(identity);
2095 self
2096 }
2097
2098 #[must_use]
2103 pub fn with_server_capabilities(mut self, capabilities: ServerCapabilityInfo) -> Self {
2104 self.server_capabilities = Some(capabilities);
2105 self
2106 }
2107
2108 #[must_use]
2110 pub fn has_progress_reporter(&self) -> bool {
2111 self.ensure_live().is_ok() && self.progress_reporter.is_some()
2112 }
2113
2114 #[must_use]
2119 pub fn progress_marker(&self) -> Option<&serde_json::Value> {
2120 self.ensure_live()
2121 .ok()
2122 .and_then(|()| self.progress_reporter.as_ref()?.marker())
2123 }
2124
2125 pub fn report_progress(&self, progress: f64, message: Option<&str>) {
2148 if self.ensure_live().is_ok()
2149 && let Some(ref reporter) = self.progress_reporter
2150 {
2151 reporter.report(progress, message);
2152 }
2153 }
2154
2155 pub fn report_progress_with_total(&self, progress: f64, total: f64, message: Option<&str>) {
2178 if self.ensure_live().is_ok()
2179 && let Some(ref reporter) = self.progress_reporter
2180 {
2181 reporter.report_with_total(progress, total, message);
2182 }
2183 }
2184
2185 pub fn report_progress_exact(
2191 &self,
2192 progress: serde_json::Number,
2193 total: Option<serde_json::Number>,
2194 message: Option<&str>,
2195 ) {
2196 if self.ensure_live().is_ok()
2197 && let Some(ref reporter) = self.progress_reporter
2198 {
2199 reporter.report_exact(progress, total, message);
2200 }
2201 }
2202
2203 #[must_use]
2208 pub fn request_id(&self) -> u64 {
2209 self.request_id
2210 }
2211
2212 #[must_use]
2219 pub fn region_id(&self) -> RegionId {
2220 self.cx.region_id()
2221 }
2222
2223 #[must_use]
2225 pub fn task_id(&self) -> TaskId {
2226 self.cx.task_id()
2227 }
2228
2229 fn apply_operation_deadline(&self, budget: Budget) -> Budget {
2230 self.operation_deadline.map_or(budget, |deadline| {
2231 budget.meet(Budget::new().with_deadline(deadline))
2232 })
2233 }
2234
2235 fn request_scope_is_active(&self) -> bool {
2236 self.request_lease.load(Ordering::Acquire) != REQUEST_LEASE_CLOSED
2237 }
2238
2239 #[must_use]
2247 pub fn budget(&self) -> Budget {
2248 let ambient = self.cx.budget();
2249 let state = *self
2250 .budget_state
2251 .lock()
2252 .unwrap_or_else(std::sync::PoisonError::into_inner);
2253 self.apply_operation_deadline(state.effective(ambient))
2254 }
2255
2256 #[must_use]
2261 pub fn is_cancelled(&self) -> bool {
2262 self.ensure_live().is_err()
2263 }
2264
2265 #[must_use]
2272 pub fn request_cancellation(&self) -> McpRequestCancellation {
2273 self.request_cancellation.clone()
2274 }
2275
2276 pub fn ensure_live(&self) -> Result<(), CancelledError> {
2291 if !self.request_scope_is_active() {
2292 return Err(CancelledError);
2293 }
2294 let _mask_transition = self
2295 .mask_transition
2296 .lock()
2297 .unwrap_or_else(std::sync::PoisonError::into_inner);
2298 if self.framework_mask_depth.load(Ordering::SeqCst) > 0 {
2299 return Ok(());
2300 }
2301
2302 let ambient = self.cx.budget();
2303 let now = self.cx.now();
2304 let state = *self
2305 .budget_state
2306 .lock()
2307 .unwrap_or_else(std::sync::PoisonError::into_inner);
2308 let effective = self.apply_operation_deadline(state.effective(ambient));
2309 if self.request_cancellation.is_cancel_requested()
2310 || self.cx.is_cancel_requested()
2311 || effective.is_past_deadline(now)
2312 || state.deferred_overrun
2313 {
2314 return Err(CancelledError);
2315 }
2316 Ok(())
2317 }
2318
2319 pub fn checkpoint(&self) -> Result<(), CancelledError> {
2351 if !self.request_scope_is_active() {
2352 return Err(CancelledError);
2353 }
2354 let _mask_transition = self
2355 .mask_transition
2356 .lock()
2357 .unwrap_or_else(std::sync::PoisonError::into_inner);
2358 let masked = self.framework_mask_depth.load(Ordering::SeqCst) > 0;
2359 let ambient = self.cx.budget();
2360 let now = self.cx.now();
2361 let mut state = self
2362 .budget_state
2363 .lock()
2364 .unwrap_or_else(std::sync::PoisonError::into_inner);
2365 let adjusted_ambient = state.adjusted_ambient(ambient);
2366 let effective = self.apply_operation_deadline(
2367 state
2368 .ceiling
2369 .map_or(adjusted_ambient, |ceiling| adjusted_ambient.meet(ceiling)),
2370 );
2371 let poll_unavailable = effective.poll_quota == 0;
2372 let past_deadline = effective.is_past_deadline(now);
2373 let cancelled =
2374 self.request_cancellation.is_cancel_requested() || self.cx.is_cancel_requested();
2375 let deferred_overrun = state.deferred_overrun;
2376
2377 if !masked && (cancelled || poll_unavailable || past_deadline || deferred_overrun) {
2378 return Err(CancelledError);
2379 }
2380
2381 if poll_unavailable {
2382 debug_assert!(masked);
2383 state.deferred_overrun = true;
2384 }
2385
2386 if adjusted_ambient.poll_quota != u32::MAX {
2387 state.ambient_poll_debits = state.ambient_poll_debits.saturating_add(1);
2388 }
2389
2390 if let Some(budget) = state.ceiling.as_mut()
2391 && budget.poll_quota != u32::MAX
2392 {
2393 if budget.consume_poll().is_none() {
2394 debug_assert!(masked);
2395 state.deferred_overrun = true;
2396 }
2397 }
2398
2399 Ok(())
2400 }
2401
2402 pub fn consume_cost(&self, cost: u64) -> Result<(), CancelledError> {
2433 if !self.request_scope_is_active() {
2434 return Err(CancelledError);
2435 }
2436 let _mask_transition = self
2437 .mask_transition
2438 .lock()
2439 .unwrap_or_else(std::sync::PoisonError::into_inner);
2440 let masked = self.framework_mask_depth.load(Ordering::SeqCst) > 0;
2441 let ambient = self.cx.budget();
2442 let now = self.cx.now();
2443 let mut state = self
2444 .budget_state
2445 .lock()
2446 .unwrap_or_else(std::sync::PoisonError::into_inner);
2447 let effective = self.apply_operation_deadline(state.effective(ambient));
2448 let enough_cost = effective
2449 .cost_quota
2450 .is_none_or(|remaining| remaining >= cost);
2451 let past_deadline = effective.is_past_deadline(now);
2452 let cancelled =
2453 self.request_cancellation.is_cancel_requested() || self.cx.is_cancel_requested();
2454
2455 if !masked && (cancelled || past_deadline || state.deferred_overrun || !enough_cost) {
2456 return Err(CancelledError);
2457 }
2458
2459 if !enough_cost {
2460 debug_assert!(masked);
2461 state.deferred_overrun = true;
2462 }
2463 state.ambient_cost_debits = state.ambient_cost_debits.saturating_add(cost);
2464 if let Some(budget) = state.ceiling.as_mut()
2465 && !budget.consume_cost(cost)
2466 {
2467 debug_assert!(masked);
2468 budget.cost_quota = Some(0);
2469 }
2470
2471 Ok(())
2472 }
2473
2474 pub fn masked<F, R>(&self, f: F) -> Result<R, CancelledError>
2502 where
2503 F: FnOnce() -> R,
2504 {
2505 if !self.request_scope_is_active() {
2506 return Err(CancelledError);
2507 }
2508 let entry_transition = self
2509 .mask_transition
2510 .lock()
2511 .unwrap_or_else(std::sync::PoisonError::into_inner);
2512 if self.framework_mask_depth.load(Ordering::SeqCst) >= MAX_MASK_DEPTH {
2513 return Err(CancelledError);
2514 }
2515 if self
2516 .framework_mask_depth
2517 .try_update(Ordering::SeqCst, Ordering::SeqCst, |depth| {
2518 depth.checked_add(1)
2519 })
2520 .is_err()
2521 {
2522 return Err(CancelledError);
2523 }
2524 let framework_mask = FrameworkMaskGuard {
2525 depth: &self.framework_mask_depth,
2526 };
2527 let masked_outcome = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
2528 self.cx.masked(|| {
2529 drop(entry_transition);
2530 let outcome = std::panic::catch_unwind(std::panic::AssertUnwindSafe(f));
2531 let exit_transition = self
2532 .mask_transition
2533 .lock()
2534 .unwrap_or_else(std::sync::PoisonError::into_inner);
2535 (outcome, exit_transition)
2536 })
2537 }));
2538 let (outcome, exit_transition) = match masked_outcome {
2539 Ok(result) => result,
2540 Err(_runtime_mask_failure) => {
2541 let exit_transition = self
2542 .mask_transition
2543 .lock()
2544 .unwrap_or_else(std::sync::PoisonError::into_inner);
2545 drop(framework_mask);
2546 drop(exit_transition);
2547 return Err(CancelledError);
2548 }
2549 };
2550 drop(framework_mask);
2551 drop(exit_transition);
2552
2553 match outcome {
2554 Ok(result) => Ok(result),
2555 Err(payload) => std::panic::resume_unwind(payload),
2556 }
2557 }
2558
2559 pub fn trace(&self, message: &str) {
2564 if self.ensure_live().is_ok() {
2565 self.cx.trace(message);
2566 }
2567 }
2568
2569 pub fn debug(&self, message: impl AsRef<str>) {
2571 self.log(McpLogLevel::Debug, message);
2572 }
2573
2574 pub fn info(&self, message: impl AsRef<str>) {
2576 self.log(McpLogLevel::Info, message);
2577 }
2578
2579 pub fn notice(&self, message: impl AsRef<str>) {
2581 self.log(McpLogLevel::Notice, message);
2582 }
2583
2584 pub fn warning(&self, message: impl AsRef<str>) {
2586 self.log(McpLogLevel::Warning, message);
2587 }
2588
2589 pub fn error(&self, message: impl AsRef<str>) {
2591 self.log(McpLogLevel::Error, message);
2592 }
2593
2594 pub fn log(&self, level: McpLogLevel, message: impl AsRef<str>) {
2599 self.log_data(
2600 level,
2601 serde_json::Value::String(message.as_ref().to_owned()),
2602 );
2603 }
2604
2605 pub fn log_data(&self, level: McpLogLevel, data: serde_json::Value) {
2607 if self.ensure_live().is_err() {
2608 return;
2609 }
2610 let Some(min_level) = self.min_log_level else {
2611 return;
2612 };
2613 if level.rank() < min_level.rank() {
2614 return;
2615 }
2616 if let Some(sender) = self.log_sender.as_ref() {
2617 sender.send_log(level, Some("fastmcp"), data);
2618 }
2619 }
2620
2621 pub fn notify_resource_updated(&self, uri: impl AsRef<str>) -> bool {
2627 if self.ensure_live().is_err() {
2628 return false;
2629 }
2630 let uri = uri.as_ref();
2631 let mut delivered = false;
2632 if self
2633 .resource_subscriptions
2634 .as_ref()
2635 .is_some_and(|uris| uris.contains(uri))
2636 && let Some(sender) = self.log_sender.as_ref()
2637 {
2638 sender.send_resource_updated(uri);
2639 delivered = true;
2640 }
2641 if let Some(publisher) = self.catalog_publisher.as_ref()
2642 && publisher.publish_resource_updated(uri)
2643 {
2644 delivered = true;
2645 }
2646 delivered
2647 }
2648
2649 #[must_use]
2662 pub fn cx(&self) -> &Cx {
2663 &self.cx
2664 }
2665
2666 #[must_use]
2674 pub fn final_result_outcome<TypedResult, LegacyResult, TerminalReason>(
2675 &self,
2676 result: crate::combinator::FinalRequestResult<TypedResult, LegacyResult, TerminalReason>,
2677 ) -> crate::McpOutcome<
2678 crate::combinator::FinalRequestResult<TypedResult, LegacyResult, TerminalReason>,
2679 > {
2680 if self.ensure_live().is_err() {
2681 return Outcome::Cancelled(self.final_result_cancellation_reason());
2682 }
2683 Outcome::Ok(result)
2684 }
2685
2686 #[must_use]
2691 pub fn adapt_final_request_outcome<TypedResult, LegacyResult, TerminalReason>(
2692 &self,
2693 outcome: crate::McpOutcome<
2694 crate::combinator::FinalRequestResult<TypedResult, LegacyResult, TerminalReason>,
2695 >,
2696 ) -> crate::McpOutcome<
2697 crate::combinator::FinalRequestResult<TypedResult, LegacyResult, TerminalReason>,
2698 > {
2699 match outcome {
2700 Outcome::Ok(result) => self.final_result_outcome(result),
2701 Outcome::Err(error) => Outcome::Err(error),
2702 Outcome::Cancelled(reason) => Outcome::Cancelled(reason),
2703 Outcome::Panicked(payload) => Outcome::Panicked(payload),
2704 }
2705 }
2706
2707 fn final_result_cancellation_reason(&self) -> CancelReason {
2708 self.cx.cancel_reason().unwrap_or_else(|| {
2709 if self.request_cancellation.is_cancel_requested() {
2710 CancelReason::user("FastMCP request-local cancellation")
2711 } else if !self.request_scope_is_active() {
2712 CancelReason::user("FastMCP request lease closed")
2713 } else {
2714 CancelReason::user("FastMCP request liveness rejected final result")
2715 }
2716 })
2717 }
2718
2719 #[must_use]
2742 pub fn get_state<T: serde::de::DeserializeOwned>(&self, key: &str) -> Option<T> {
2743 if !self.request_scope_is_active() {
2744 return None;
2745 }
2746 self.state.as_ref()?.get(key)
2747 }
2748
2749 #[must_use]
2751 pub fn auth(&self) -> Option<AuthContext> {
2752 if !self.request_scope_is_active() {
2753 return None;
2754 }
2755 self.auth
2756 .lock()
2757 .unwrap_or_else(std::sync::PoisonError::into_inner)
2758 .clone()
2759 }
2760
2761 pub fn set_auth(&self, auth: AuthContext) -> bool {
2770 if self.ensure_live().is_err() {
2771 return false;
2772 }
2773 let mut slot = self
2774 .auth
2775 .lock()
2776 .unwrap_or_else(std::sync::PoisonError::into_inner);
2777 if self.auth_state.load(Ordering::Acquire) != REQUEST_AUTH_UNCOMMITTED {
2778 return false;
2779 }
2780 *slot = Some(auth);
2781 self.auth_state
2782 .store(REQUEST_AUTH_AUTHENTICATED, Ordering::Release);
2783 true
2784 }
2785
2786 #[doc(hidden)]
2794 pub fn commit_anonymous_auth(&self) -> bool {
2795 if self.ensure_live().is_err() {
2796 return false;
2797 }
2798 let slot = self
2799 .auth
2800 .lock()
2801 .unwrap_or_else(std::sync::PoisonError::into_inner);
2802 if self.auth_state.load(Ordering::Acquire) != REQUEST_AUTH_UNCOMMITTED || slot.is_some() {
2803 return false;
2804 }
2805 self.auth_state
2806 .store(REQUEST_AUTH_ANONYMOUS, Ordering::Release);
2807 true
2808 }
2809
2810 #[doc(hidden)]
2816 #[must_use]
2817 pub fn cache_auth_partition(&self) -> Option<Option<AuthContext>> {
2818 if !self.request_scope_is_active() {
2819 return None;
2820 }
2821 let slot = self
2822 .auth
2823 .lock()
2824 .unwrap_or_else(std::sync::PoisonError::into_inner);
2825 match self.auth_state.load(Ordering::Acquire) {
2826 REQUEST_AUTH_ANONYMOUS => Some(None),
2827 REQUEST_AUTH_AUTHENTICATED => slot.clone().map(Some),
2828 _ => None,
2829 }
2830 }
2831
2832 #[must_use]
2834 pub fn with_auth(self, auth: AuthContext) -> Self {
2835 let _ = self.set_auth(auth);
2836 self
2837 }
2838
2839 #[must_use]
2847 pub fn with_isolated_auth(mut self) -> Self {
2848 let already_committed = self.auth_state.load(Ordering::Acquire) != REQUEST_AUTH_UNCOMMITTED;
2849 if already_committed {
2850 return self;
2851 }
2852 self.auth = Arc::new(Mutex::new(None));
2853 self.auth_state = Arc::new(AtomicU8::new(REQUEST_AUTH_UNCOMMITTED));
2854 self.state = None;
2855 self.progress_reporter = None;
2856 self.sampling_sender = None;
2857 self.elicitation_sender = None;
2858 self.roots_provider = None;
2859 self.resource_reader = None;
2860 self.tool_caller = None;
2861 self.prompt_caller = None;
2862 self
2863 }
2864
2865 pub fn set_state<T: serde::Serialize>(&self, key: impl Into<String>, value: T) -> bool {
2882 if self.ensure_live().is_err() {
2883 return false;
2884 }
2885 match &self.state {
2886 Some(state) => state.set(key, value),
2887 None => false,
2888 }
2889 }
2890
2891 pub fn remove_state(&self, key: &str) -> Option<serde_json::Value> {
2897 if self.ensure_live().is_err() {
2898 return None;
2899 }
2900 self.state.as_ref()?.remove(key)
2901 }
2902
2903 #[must_use]
2907 pub fn has_state(&self, key: &str) -> bool {
2908 self.request_scope_is_active() && self.state.as_ref().is_some_and(|s| s.contains(key))
2909 }
2910
2911 #[must_use]
2913 pub fn has_session_state(&self) -> bool {
2914 self.request_scope_is_active() && self.state.is_some()
2915 }
2916
2917 #[doc(hidden)]
2919 #[must_use]
2920 pub fn session_is_ephemeral(&self) -> bool {
2921 self.request_scope_is_active()
2922 && self.state.as_ref().is_some_and(SessionState::is_ephemeral)
2923 }
2924
2925 #[must_use]
2931 pub fn session_state(&self) -> Option<&SessionState> {
2932 self.state.as_ref()
2933 }
2934
2935 #[doc(hidden)]
2943 #[must_use]
2944 pub fn session_cache_partition(&self) -> Option<([u8; 32], u64)> {
2945 if !self.request_scope_is_active() {
2946 return None;
2947 }
2948 self.state.as_ref()?.cache_partition()
2949 }
2950
2951 #[doc(hidden)]
2957 #[must_use]
2958 pub fn begin_session_cache_partition(&self) -> Option<([u8; 32], u64)> {
2959 let current = self.session_cache_partition()?;
2960 let mut admitted = self
2961 .cache_admission_partition
2962 .lock()
2963 .unwrap_or_else(std::sync::PoisonError::into_inner);
2964 match *admitted {
2965 None => {
2966 *admitted = Some(current);
2967 Some(current)
2968 }
2969 Some(existing) if existing == current => Some(existing),
2970 Some(_) => None,
2971 }
2972 }
2973
2974 #[doc(hidden)]
2977 #[must_use]
2978 pub fn complete_session_cache_partition(&self) -> Option<([u8; 32], u64)> {
2979 if !self.request_scope_is_active() {
2980 return None;
2981 }
2982 let admitted = *self
2983 .cache_admission_partition
2984 .lock()
2985 .unwrap_or_else(std::sync::PoisonError::into_inner);
2986 let admitted = admitted?;
2987 (self.state.as_ref()?.cache_partition() == Some(admitted)).then_some(admitted)
2988 }
2989
2990 #[doc(hidden)]
2993 pub fn mark_response_cache_hit(&self, cache_id: u64) -> bool {
2994 const MAX_CACHE_MIDDLEWARE_PER_REQUEST: usize = 64;
2995 if !self.request_scope_is_active() || cache_id == 0 {
2996 return false;
2997 }
2998 let mut hits = self
2999 .response_cache_hits
3000 .lock()
3001 .unwrap_or_else(std::sync::PoisonError::into_inner);
3002 if hits.contains(&cache_id) {
3003 return true;
3004 }
3005 if hits.len() >= MAX_CACHE_MIDDLEWARE_PER_REQUEST || hits.try_reserve(1).is_err() {
3006 return false;
3007 }
3008 hits.push(cache_id);
3009 true
3010 }
3011
3012 #[doc(hidden)]
3015 #[must_use]
3016 pub fn response_was_cache_hit(&self, cache_id: u64) -> bool {
3017 self.request_scope_is_active()
3018 && cache_id != 0
3019 && self
3020 .response_cache_hits
3021 .lock()
3022 .unwrap_or_else(std::sync::PoisonError::into_inner)
3023 .contains(&cache_id)
3024 }
3025
3026 #[doc(hidden)]
3028 #[must_use]
3029 pub fn response_was_served_from_cache(&self) -> bool {
3030 self.request_scope_is_active()
3031 && !self
3032 .response_cache_hits
3033 .lock()
3034 .unwrap_or_else(std::sync::PoisonError::into_inner)
3035 .is_empty()
3036 }
3037
3038 #[must_use]
3047 pub fn client_capabilities(&self) -> Option<&ClientCapabilityInfo> {
3048 self.client_capabilities.as_ref()
3049 }
3050
3051 #[must_use]
3056 pub fn client_implementation(&self) -> Option<&ClientImplementationInfo> {
3057 self.client_implementation.as_ref()
3058 }
3059
3060 #[must_use]
3064 pub fn server_capabilities(&self) -> Option<&ServerCapabilityInfo> {
3065 self.server_capabilities.as_ref()
3066 }
3067
3068 #[must_use]
3073 pub fn client_supports_sampling(&self) -> bool {
3074 self.client_capabilities
3075 .as_ref()
3076 .is_some_and(|c| c.sampling)
3077 }
3078
3079 #[must_use]
3084 pub fn client_supports_elicitation(&self) -> bool {
3085 self.client_capabilities
3086 .as_ref()
3087 .is_some_and(|c| c.elicitation)
3088 }
3089
3090 #[must_use]
3092 pub fn client_supports_elicitation_form(&self) -> bool {
3093 self.client_capabilities
3094 .as_ref()
3095 .is_some_and(|c| c.elicitation_form)
3096 }
3097
3098 #[must_use]
3100 pub fn client_supports_elicitation_url(&self) -> bool {
3101 self.client_capabilities
3102 .as_ref()
3103 .is_some_and(|c| c.elicitation_url)
3104 }
3105
3106 #[must_use]
3111 pub fn client_supports_roots(&self) -> bool {
3112 self.client_capabilities.as_ref().is_some_and(|c| c.roots)
3113 }
3114
3115 const DISABLED_TOOLS_KEY: &'static str = "fastmcp.disabled_tools";
3121 const DISABLED_RESOURCES_KEY: &'static str = "fastmcp.disabled_resources";
3123 const DISABLED_PROMPTS_KEY: &'static str = "fastmcp.disabled_prompts";
3125
3126 pub fn disable_tool(&self, name: impl Into<String>) -> bool {
3144 self.add_to_disabled_set(Self::DISABLED_TOOLS_KEY, name.into(), McpCatalogKind::Tools)
3145 }
3146
3147 pub fn enable_tool(&self, name: &str) -> bool {
3151 self.remove_from_disabled_set(Self::DISABLED_TOOLS_KEY, name, McpCatalogKind::Tools)
3152 }
3153
3154 #[must_use]
3158 pub fn is_tool_enabled(&self, name: &str) -> bool {
3159 self.request_scope_is_active() && !self.is_in_disabled_set(Self::DISABLED_TOOLS_KEY, name)
3160 }
3161
3162 pub fn disable_resource(&self, uri: impl Into<String>) -> bool {
3169 self.add_to_disabled_set(
3170 Self::DISABLED_RESOURCES_KEY,
3171 uri.into(),
3172 McpCatalogKind::Resources,
3173 )
3174 }
3175
3176 pub fn enable_resource(&self, uri: &str) -> bool {
3180 self.remove_from_disabled_set(Self::DISABLED_RESOURCES_KEY, uri, McpCatalogKind::Resources)
3181 }
3182
3183 #[must_use]
3187 pub fn is_resource_enabled(&self, uri: &str) -> bool {
3188 self.request_scope_is_active()
3189 && !self.is_in_disabled_set(Self::DISABLED_RESOURCES_KEY, uri)
3190 }
3191
3192 pub fn disable_prompt(&self, name: impl Into<String>) -> bool {
3199 self.add_to_disabled_set(
3200 Self::DISABLED_PROMPTS_KEY,
3201 name.into(),
3202 McpCatalogKind::Prompts,
3203 )
3204 }
3205
3206 pub fn enable_prompt(&self, name: &str) -> bool {
3210 self.remove_from_disabled_set(Self::DISABLED_PROMPTS_KEY, name, McpCatalogKind::Prompts)
3211 }
3212
3213 #[must_use]
3217 pub fn is_prompt_enabled(&self, name: &str) -> bool {
3218 self.request_scope_is_active() && !self.is_in_disabled_set(Self::DISABLED_PROMPTS_KEY, name)
3219 }
3220
3221 #[must_use]
3223 pub fn disabled_tools(&self) -> std::collections::HashSet<String> {
3224 self.get_disabled_set(Self::DISABLED_TOOLS_KEY)
3225 }
3226
3227 #[must_use]
3229 pub fn disabled_resources(&self) -> std::collections::HashSet<String> {
3230 self.get_disabled_set(Self::DISABLED_RESOURCES_KEY)
3231 }
3232
3233 #[must_use]
3235 pub fn disabled_prompts(&self) -> std::collections::HashSet<String> {
3236 self.get_disabled_set(Self::DISABLED_PROMPTS_KEY)
3237 }
3238
3239 fn add_to_disabled_set(&self, key: &str, name: String, kind: McpCatalogKind) -> bool {
3241 if self.ensure_live().is_err() {
3242 return false;
3243 }
3244 let Some(state) = self.state.as_ref() else {
3245 return false;
3246 };
3247 let mut set: std::collections::HashSet<String> = state.get(key).unwrap_or_default();
3248 let changed = set.insert(name);
3249 let stored = state.set(key, set);
3250 if stored && changed {
3251 self.emit_catalog_changed(kind);
3252 }
3253 stored
3254 }
3255
3256 fn remove_from_disabled_set(&self, key: &str, name: &str, kind: McpCatalogKind) -> bool {
3258 if self.ensure_live().is_err() {
3259 return false;
3260 }
3261 let Some(state) = self.state.as_ref() else {
3262 return false;
3263 };
3264 let mut set: std::collections::HashSet<String> = state.get(key).unwrap_or_default();
3265 let changed = set.remove(name);
3266 let stored = state.set(key, set);
3267 if stored && changed {
3268 self.emit_catalog_changed(kind);
3269 }
3270 stored
3271 }
3272
3273 fn emit_catalog_changed(&self, kind: McpCatalogKind) {
3274 if let Some(sender) = self.log_sender.as_ref() {
3275 sender.send_catalog_changed(kind);
3276 }
3277 if let Some(publisher) = self.catalog_publisher.as_ref() {
3278 let _ = publisher.publish_catalog_changed(kind);
3279 }
3280 }
3281
3282 fn is_in_disabled_set(&self, key: &str, name: &str) -> bool {
3284 if !self.request_scope_is_active() {
3285 return false;
3286 }
3287 let Some(state) = self.state.as_ref() else {
3288 return false;
3289 };
3290 let set: std::collections::HashSet<String> = state.get(key).unwrap_or_default();
3291 set.contains(name)
3292 }
3293
3294 fn get_disabled_set(&self, key: &str) -> std::collections::HashSet<String> {
3296 if !self.request_scope_is_active() {
3297 return std::collections::HashSet::new();
3298 }
3299 self.state
3300 .as_ref()
3301 .and_then(|s| s.get(key))
3302 .unwrap_or_default()
3303 }
3304
3305 #[must_use]
3311 pub fn can_list_roots(&self) -> bool {
3312 self.ensure_live().is_ok() && self.roots_provider.is_some()
3313 }
3314
3315 pub async fn list_roots(&self) -> crate::McpResult<Vec<ClientRoot>> {
3322 self.ensure_live()
3323 .map_err(|_| crate::McpError::request_cancelled())?;
3324 let provider = self.roots_provider.as_ref().ok_or_else(|| {
3325 crate::McpError::new(
3326 crate::McpErrorCode::InvalidRequest,
3327 "Roots not available: client does not support roots capability",
3328 )
3329 })?;
3330
3331 let roots = provider.list_roots().await?;
3332 self.ensure_live()
3333 .map_err(|_| crate::McpError::request_cancelled())?;
3334 Ok(roots)
3335 }
3336
3337 #[must_use]
3346 pub fn can_sample(&self) -> bool {
3347 self.ensure_live().is_ok() && self.sampling_sender.is_some()
3348 }
3349
3350 pub async fn sample(
3375 &self,
3376 prompt: impl Into<String>,
3377 max_tokens: u32,
3378 ) -> crate::McpResult<SamplingResponse> {
3379 let request = SamplingRequest::prompt(prompt, max_tokens);
3380 self.sample_with_request(request).await
3381 }
3382
3383 pub async fn sample_with_request(
3415 &self,
3416 request: SamplingRequest,
3417 ) -> crate::McpResult<SamplingResponse> {
3418 self.ensure_live()
3419 .map_err(|_| crate::McpError::request_cancelled())?;
3420 let sender = self.sampling_sender.as_ref().ok_or_else(|| {
3421 crate::McpError::new(
3422 crate::McpErrorCode::InvalidRequest,
3423 "Sampling not available: client does not support sampling capability",
3424 )
3425 })?;
3426
3427 let response = sender.create_message(request).await?;
3428 self.ensure_live()
3429 .map_err(|_| crate::McpError::request_cancelled())?;
3430 Ok(response)
3431 }
3432
3433 #[must_use]
3442 pub fn can_elicit(&self) -> bool {
3443 self.ensure_live().is_ok() && self.elicitation_sender.is_some()
3444 }
3445
3446 pub async fn elicit_form(
3484 &self,
3485 message: impl Into<String>,
3486 schema: serde_json::Value,
3487 ) -> crate::McpResult<ElicitationResponse> {
3488 let request = ElicitationRequest::form(message, schema);
3489 self.elicit_with_request(request).await
3490 }
3491
3492 pub async fn elicit_url(
3526 &self,
3527 message: impl Into<String>,
3528 url: impl Into<String>,
3529 elicitation_id: impl Into<String>,
3530 ) -> crate::McpResult<ElicitationResponse> {
3531 let request = ElicitationRequest::url(message, url, elicitation_id);
3532 self.elicit_with_request(request).await
3533 }
3534
3535 pub async fn elicit_with_request(
3547 &self,
3548 request: ElicitationRequest,
3549 ) -> crate::McpResult<ElicitationResponse> {
3550 self.ensure_live()
3551 .map_err(|_| crate::McpError::request_cancelled())?;
3552 let sender = self.elicitation_sender.as_ref().ok_or_else(|| {
3553 crate::McpError::new(
3554 crate::McpErrorCode::InvalidRequest,
3555 "Elicitation not available: client does not support elicitation capability",
3556 )
3557 })?;
3558
3559 let response = sender.elicit(request).await?;
3560 self.ensure_live()
3561 .map_err(|_| crate::McpError::request_cancelled())?;
3562 Ok(response)
3563 }
3564
3565 #[must_use]
3574 pub fn can_read_resources(&self) -> bool {
3575 self.ensure_live().is_ok() && self.resource_reader.is_some()
3576 }
3577
3578 #[must_use]
3582 pub fn resource_read_depth(&self) -> u32 {
3583 self.resource_read_depth
3584 }
3585
3586 pub async fn read_resource(&self, uri: &str) -> crate::McpResult<ResourceReadResult> {
3615 self.ensure_live()
3616 .map_err(|_| crate::McpError::request_cancelled())?;
3617 let reader = self.resource_reader.as_ref().ok_or_else(|| {
3619 crate::McpError::new(
3620 crate::McpErrorCode::InternalError,
3621 "Resource reading not available: no router attached to context",
3622 )
3623 })?;
3624
3625 let nested_dispatch_depth = self.nested_dispatch_depth();
3629 if nested_dispatch_depth >= MAX_RESOURCE_READ_DEPTH {
3630 return Err(crate::McpError::new(
3631 crate::McpErrorCode::InternalError,
3632 format!(
3633 "Maximum resource read depth ({}) exceeded; possible infinite recursion",
3634 MAX_RESOURCE_READ_DEPTH
3635 ),
3636 ));
3637 }
3638
3639 let result = reader
3641 .read_resource(self, uri, nested_dispatch_depth + 1)
3642 .await?;
3643 self.ensure_live()
3644 .map_err(|_| crate::McpError::request_cancelled())?;
3645 Ok(result)
3646 }
3647
3648 pub async fn read_resource_text(&self, uri: &str) -> crate::McpResult<String> {
3666 let result = self.read_resource(uri).await?;
3667 result.first_text().map(String::from).ok_or_else(|| {
3668 crate::McpError::new(
3669 crate::McpErrorCode::InternalError,
3670 format!("Resource '{}' has no text content", uri),
3671 )
3672 })
3673 }
3674
3675 pub async fn read_resource_json<T: serde::de::DeserializeOwned>(
3699 &self,
3700 uri: &str,
3701 ) -> crate::McpResult<T> {
3702 let text = self.read_resource_text(uri).await?;
3703 serde_json::from_str(&text).map_err(|e| {
3704 crate::McpError::new(
3705 crate::McpErrorCode::InternalError,
3706 format!("Failed to parse resource '{}' as JSON: {}", uri, e),
3707 )
3708 })
3709 }
3710
3711 #[must_use]
3720 pub fn can_call_tools(&self) -> bool {
3721 self.ensure_live().is_ok() && self.tool_caller.is_some()
3722 }
3723
3724 #[must_use]
3728 pub fn tool_call_depth(&self) -> u32 {
3729 self.tool_call_depth
3730 }
3731
3732 pub async fn call_tool(
3760 &self,
3761 name: &str,
3762 args: serde_json::Value,
3763 ) -> crate::McpResult<ToolCallResult> {
3764 self.ensure_live()
3765 .map_err(|_| crate::McpError::request_cancelled())?;
3766 let caller = self.tool_caller.as_ref().ok_or_else(|| {
3768 crate::McpError::new(
3769 crate::McpErrorCode::InternalError,
3770 "Tool calling not available: no router attached to context",
3771 )
3772 })?;
3773
3774 let nested_dispatch_depth = self.nested_dispatch_depth();
3777 if nested_dispatch_depth >= MAX_TOOL_CALL_DEPTH {
3778 return Err(crate::McpError::new(
3779 crate::McpErrorCode::InternalError,
3780 format!(
3781 "Maximum tool call depth ({}) exceeded calling '{}'; possible infinite recursion",
3782 MAX_TOOL_CALL_DEPTH, name
3783 ),
3784 ));
3785 }
3786
3787 let result = caller
3789 .call_tool(self, name, args, nested_dispatch_depth + 1)
3790 .await?;
3791 self.ensure_live()
3792 .map_err(|_| crate::McpError::request_cancelled())?;
3793 Ok(result)
3794 }
3795
3796 pub async fn call_tool_text(
3815 &self,
3816 name: &str,
3817 args: serde_json::Value,
3818 ) -> crate::McpResult<String> {
3819 let result = self.call_tool(name, args).await?;
3820
3821 if result.is_error {
3823 let error_msg = result.first_text().unwrap_or("Tool returned an error");
3824 return Err(crate::McpError::new(
3825 crate::McpErrorCode::InternalError,
3826 format!("Tool '{}' failed: {}", name, error_msg),
3827 ));
3828 }
3829
3830 result.first_text().map(String::from).ok_or_else(|| {
3831 crate::McpError::new(
3832 crate::McpErrorCode::InternalError,
3833 format!("Tool '{}' returned no text content", name),
3834 )
3835 })
3836 }
3837
3838 pub async fn call_tool_json<T: serde::de::DeserializeOwned>(
3863 &self,
3864 name: &str,
3865 args: serde_json::Value,
3866 ) -> crate::McpResult<T> {
3867 let text = self.call_tool_text(name, args).await?;
3868 serde_json::from_str(&text).map_err(|e| {
3869 crate::McpError::new(
3870 crate::McpErrorCode::InternalError,
3871 format!("Failed to parse tool '{}' result as JSON: {}", name, e),
3872 )
3873 })
3874 }
3875
3876 #[must_use]
3882 pub fn can_get_prompts(&self) -> bool {
3883 self.ensure_live().is_ok() && self.prompt_caller.is_some()
3884 }
3885
3886 #[must_use]
3888 pub fn prompt_get_depth(&self) -> u32 {
3889 self.prompt_get_depth
3890 }
3891
3892 fn nested_dispatch_depth(&self) -> u32 {
3893 self.resource_read_depth
3894 .max(self.tool_call_depth)
3895 .max(self.prompt_get_depth)
3896 }
3897
3898 pub async fn get_prompt(
3903 &self,
3904 name: &str,
3905 arguments: std::collections::HashMap<String, String>,
3906 ) -> crate::McpResult<PromptGetResult> {
3907 self.ensure_live()
3908 .map_err(|_| crate::McpError::request_cancelled())?;
3909 let caller = self.prompt_caller.as_ref().ok_or_else(|| {
3910 crate::McpError::new(
3911 crate::McpErrorCode::InternalError,
3912 "Prompt getting not available: no router attached to context",
3913 )
3914 })?;
3915
3916 let nested_dispatch_depth = self.nested_dispatch_depth();
3917 if nested_dispatch_depth >= MAX_PROMPT_GET_DEPTH {
3918 return Err(crate::McpError::new(
3919 crate::McpErrorCode::InternalError,
3920 format!(
3921 "Maximum prompt get depth ({}) exceeded getting '{}'; possible infinite recursion",
3922 MAX_PROMPT_GET_DEPTH, name
3923 ),
3924 ));
3925 }
3926
3927 let result = caller
3928 .get_prompt(self, name, arguments, nested_dispatch_depth + 1)
3929 .await?;
3930 self.ensure_live()
3931 .map_err(|_| crate::McpError::request_cancelled())?;
3932 Ok(result)
3933 }
3934
3935 pub async fn get_prompt_text(
3937 &self,
3938 name: &str,
3939 arguments: std::collections::HashMap<String, String>,
3940 ) -> crate::McpResult<String> {
3941 let result = self.get_prompt(name, arguments).await?;
3942 result.first_text().map(String::from).ok_or_else(|| {
3943 crate::McpError::new(
3944 crate::McpErrorCode::InternalError,
3945 format!("Prompt '{}' returned no text content", name),
3946 )
3947 })
3948 }
3949
3950 pub async fn join_all<T: Send + 'static>(
3970 &self,
3971 futures: Vec<crate::combinator::BoxFuture<'_, T>>,
3972 ) -> crate::McpResult<Vec<T>> {
3973 self.ensure_live()
3974 .map_err(|_| crate::McpError::request_cancelled())?;
3975 let results = crate::combinator::join_all(&self.cx, futures).await;
3976 self.ensure_live()
3977 .map_err(|_| crate::McpError::request_cancelled())?;
3978 Ok(results)
3979 }
3980
3981 pub async fn race<T: Send + 'static>(
3998 &self,
3999 futures: Vec<crate::combinator::BoxFuture<'_, T>>,
4000 ) -> crate::McpResult<T> {
4001 self.ensure_live()
4002 .map_err(|_| crate::McpError::request_cancelled())?;
4003 let result = crate::combinator::race(&self.cx, futures).await;
4004 self.ensure_live()
4005 .map_err(|_| crate::McpError::request_cancelled())?;
4006 result
4007 }
4008
4009 pub async fn quorum<T: Send + 'static>(
4026 &self,
4027 required: usize,
4028 futures: Vec<crate::combinator::BoxFuture<'_, crate::McpResult<T>>>,
4029 ) -> crate::McpResult<crate::combinator::QuorumResult<T>> {
4030 self.ensure_live()
4031 .map_err(|_| crate::McpError::request_cancelled())?;
4032 let result = crate::combinator::quorum(&self.cx, required, futures).await;
4033 self.ensure_live()
4034 .map_err(|_| crate::McpError::request_cancelled())?;
4035 result
4036 }
4037
4038 pub async fn first_ok<T: Send + 'static>(
4055 &self,
4056 futures: Vec<crate::combinator::BoxFuture<'_, crate::McpResult<T>>>,
4057 ) -> crate::McpResult<T> {
4058 self.ensure_live()
4059 .map_err(|_| crate::McpError::request_cancelled())?;
4060 let result = crate::combinator::first_ok(&self.cx, futures).await;
4061 self.ensure_live()
4062 .map_err(|_| crate::McpError::request_cancelled())?;
4063 result
4064 }
4065}
4066
4067#[derive(Debug, Clone, Copy)]
4073pub struct CancelledError;
4074
4075impl std::fmt::Display for CancelledError {
4076 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
4077 write!(f, "request cancelled")
4078 }
4079}
4080
4081impl std::error::Error for CancelledError {}
4082
4083pub trait IntoOutcome<T, E> {
4088 fn into_outcome(self) -> Outcome<T, E>;
4090}
4091
4092impl<T, E> IntoOutcome<T, E> for Result<T, E> {
4093 fn into_outcome(self) -> Outcome<T, E> {
4094 match self {
4095 Ok(v) => Outcome::Ok(v),
4096 Err(e) => Outcome::Err(e),
4097 }
4098 }
4099}
4100
4101impl<T, E> IntoOutcome<T, E> for Result<T, CancelledError>
4102where
4103 E: Default,
4104{
4105 fn into_outcome(self) -> Outcome<T, E> {
4106 match self {
4107 Ok(v) => Outcome::Ok(v),
4108 Err(CancelledError) => Outcome::Cancelled(CancelReason::user("request cancelled")),
4109 }
4110 }
4111}
4112
4113#[cfg(test)]
4114mod tests {
4115 use super::*;
4116
4117 #[test]
4118 fn test_mcp_context_creation() {
4119 let cx = Cx::for_testing();
4120 let ctx = McpContext::new(cx, 42);
4121
4122 assert_eq!(ctx.request_id(), 42);
4123 }
4124
4125 #[test]
4126 fn test_mcp_context_not_cancelled_initially() {
4127 let cx = Cx::for_testing();
4128 let ctx = McpContext::new(cx, 1);
4129
4130 assert!(!ctx.is_cancelled());
4131 }
4132
4133 #[test]
4134 fn test_mcp_context_checkpoint_success() {
4135 let cx = Cx::for_testing();
4136 let ctx = McpContext::new(cx, 1);
4137
4138 assert!(ctx.checkpoint().is_ok());
4140 }
4141
4142 #[test]
4143 fn test_mcp_context_checkpoint_cancelled() {
4144 let cx = Cx::for_testing();
4145 cx.set_cancel_requested(true);
4146 let ctx = McpContext::new(cx, 1);
4147
4148 assert!(ctx.checkpoint().is_err());
4150 }
4151
4152 #[test]
4153 fn request_local_cancellation_does_not_cancel_shared_ambient_context() {
4154 let cx = Cx::for_testing();
4155 let cancellation = McpRequestCancellation::new();
4156 let request =
4157 McpContext::new(cx.clone(), 1).with_request_cancellation(cancellation.clone());
4158 let sibling = McpContext::new(cx.clone(), 2);
4159
4160 cancellation.cancel();
4161
4162 assert!(request.ensure_live().is_err());
4163 assert!(request.checkpoint().is_err());
4164 assert!(sibling.ensure_live().is_ok());
4165 assert!(!cx.is_cancel_requested());
4166 }
4167
4168 #[test]
4169 fn context_exposes_its_request_local_cancellation_handle() {
4170 let cancellation = McpRequestCancellation::new();
4171 let context =
4172 McpContext::new(Cx::for_testing(), 1).with_request_cancellation(cancellation.clone());
4173
4174 let observed = context.request_cancellation();
4175 assert!(observed.cancel());
4176 assert!(cancellation.is_cancel_requested());
4177 assert!(context.is_cancelled());
4178 }
4179
4180 #[test]
4181 fn request_local_cancelled_future_registers_and_is_woken_without_polling() {
4182 use std::sync::atomic::AtomicBool;
4183
4184 struct WakeFlag(AtomicBool);
4185
4186 impl std::task::Wake for WakeFlag {
4187 fn wake(self: Arc<Self>) {
4188 self.0.store(true, Ordering::Release);
4189 }
4190 }
4191
4192 let cancellation = McpRequestCancellation::new();
4193 let wake_flag = Arc::new(WakeFlag(AtomicBool::new(false)));
4194 let waker = std::task::Waker::from(Arc::clone(&wake_flag));
4195 let mut task_cx = std::task::Context::from_waker(&waker);
4196 let mut future = Box::pin(cancellation.cancelled());
4197
4198 assert!(std::future::Future::poll(future.as_mut(), &mut task_cx).is_pending());
4199 assert!(cancellation.cancel());
4200 assert!(wake_flag.0.load(Ordering::Acquire));
4201 assert!(std::future::Future::poll(future.as_mut(), &mut task_cx).is_ready());
4202 }
4203
4204 #[test]
4205 fn request_local_cancelled_future_observes_preexisting_cancellation() {
4206 let cancellation = McpRequestCancellation::new();
4207 assert!(cancellation.cancel());
4208
4209 let mut future = Box::pin(cancellation.cancelled());
4210 let waker = std::task::Waker::noop();
4211 let mut task_cx = std::task::Context::from_waker(waker);
4212
4213 assert!(std::future::Future::poll(future.as_mut(), &mut task_cx).is_ready());
4214 }
4215
4216 #[test]
4217 fn request_terminal_future_is_woken_when_finalization_wins() {
4218 use std::sync::atomic::AtomicBool;
4219
4220 struct WakeFlag(AtomicBool);
4221
4222 impl std::task::Wake for WakeFlag {
4223 fn wake(self: Arc<Self>) {
4224 self.0.store(true, Ordering::Release);
4225 }
4226 }
4227
4228 let cancellation = McpRequestCancellation::new();
4229 let wake_flag = Arc::new(WakeFlag(AtomicBool::new(false)));
4230 let waker = std::task::Waker::from(Arc::clone(&wake_flag));
4231 let mut task_cx = std::task::Context::from_waker(&waker);
4232 let mut future = Box::pin(cancellation.terminated());
4233
4234 assert!(!cancellation.is_terminal());
4235 assert!(std::future::Future::poll(future.as_mut(), &mut task_cx).is_pending());
4236 assert!(cancellation.begin_finalization());
4237 assert!(cancellation.is_terminal());
4238 assert!(wake_flag.0.load(Ordering::Acquire));
4239 assert!(std::future::Future::poll(future.as_mut(), &mut task_cx).is_ready());
4240 }
4241
4242 #[test]
4243 fn request_local_cancellation_is_deferred_inside_framework_mask() {
4244 let cancellation = McpRequestCancellation::new();
4245 let ctx =
4246 McpContext::new(Cx::for_testing(), 1).with_request_cancellation(cancellation.clone());
4247
4248 let checkpoint = ctx
4249 .masked(|| {
4250 cancellation.cancel();
4251 ctx.checkpoint()
4252 })
4253 .expect("framework mask should be admitted");
4254
4255 assert!(checkpoint.is_ok());
4256 assert!(ctx.ensure_live().is_err());
4257 }
4258
4259 #[test]
4260 fn request_local_cancellation_stops_state_and_capability_effects() {
4261 let state = SessionState::new();
4262 assert!(state.set("existing", 1_u32));
4263 let cancellation = McpRequestCancellation::new();
4264 let ctx = McpContext::with_state(Cx::for_testing(), 1, state.clone())
4265 .with_sampling(Arc::new(NoOpSamplingSender))
4266 .with_elicitation(Arc::new(NoOpElicitationSender))
4267 .with_request_cancellation(cancellation.clone());
4268
4269 assert!(ctx.can_sample());
4270 assert!(ctx.can_elicit());
4271 assert!(cancellation.cancel());
4272
4273 assert!(!ctx.set_state("late", 2_u32));
4274 assert!(ctx.remove_state("existing").is_none());
4275 assert!(!ctx.disable_tool("late-tool"));
4276 assert!(!ctx.disable_resource("late://resource"));
4277 assert!(!ctx.disable_prompt("late-prompt"));
4278 assert!(!ctx.can_sample());
4279 assert!(!ctx.can_elicit());
4280 assert_eq!(state.get::<u32>("existing"), Some(1));
4281 assert!(!state.contains("late"));
4282 }
4283
4284 #[test]
4285 fn admitted_mask_allows_critical_state_commit_before_cancellation_surfaces() {
4286 let state = SessionState::new();
4287 let cancellation = McpRequestCancellation::new();
4288 let ctx = McpContext::with_state(Cx::for_testing(), 1, state.clone())
4289 .with_request_cancellation(cancellation.clone());
4290
4291 let committed = ctx
4292 .masked(|| {
4293 assert!(cancellation.cancel());
4294 ctx.set_state("critical-commit", true)
4295 })
4296 .expect("mask should be admitted before cancellation");
4297
4298 assert!(committed);
4299 assert_eq!(state.get::<bool>("critical-commit"), Some(true));
4300 assert!(ctx.ensure_live().is_err());
4301 }
4302
4303 #[test]
4304 fn active_request_clone_cannot_replace_cancellation_authority() {
4305 let original = McpRequestCancellation::new();
4306 let replacement = McpRequestCancellation::new();
4307 let root =
4308 McpContext::new(Cx::for_testing(), 1).with_request_cancellation(original.clone());
4309 let (scoped, _guard) = root
4310 .begin_request_scope()
4311 .expect("new context should activate one request lease");
4312 let attempted_escape = scoped
4313 .clone()
4314 .with_request_cancellation(replacement.clone());
4315
4316 assert!(original.cancel());
4317 assert!(attempted_escape.ensure_live().is_err());
4318 assert!(!replacement.is_cancel_requested());
4319 }
4320
4321 #[test]
4322 fn request_finalization_and_cancellation_have_one_atomic_winner() {
4323 let cancellation_wins = McpRequestCancellation::new();
4324 assert!(cancellation_wins.cancel());
4325 assert!(!cancellation_wins.begin_finalization());
4326 assert!(cancellation_wins.is_cancel_requested());
4327 assert!(cancellation_wins.is_terminal());
4328
4329 let finalization_wins = McpRequestCancellation::new();
4330 assert!(finalization_wins.begin_finalization());
4331 assert!(finalization_wins.is_finalizing());
4332 assert!(finalization_wins.is_terminal());
4333 assert!(!finalization_wins.cancel());
4334 assert!(!finalization_wins.is_cancel_requested());
4335 }
4336
4337 #[test]
4338 fn test_mcp_context_checkpoint_budget_exhausted() {
4339 let cx = Cx::for_testing_with_budget(Budget::ZERO);
4340 let ctx = McpContext::new(cx, 1);
4341
4342 assert!(ctx.checkpoint().is_err());
4344 }
4345
4346 #[test]
4347 fn checkpoint_does_not_treat_zero_cost_as_poll_exhaustion() {
4348 let budget = Budget::new().with_poll_quota(2).with_cost_quota(0);
4349 let cx = Cx::for_testing_with_budget(budget);
4350 let ctx = McpContext::new(cx.clone(), 1);
4351
4352 assert!(ctx.checkpoint().is_ok());
4353 assert!(!cx.is_cancel_requested());
4354 assert_eq!(ctx.budget().cost_quota, Some(0));
4355 }
4356
4357 #[test]
4358 fn closed_request_lease_cannot_be_revived_or_use_framework_capabilities() {
4359 let state = SessionState::new();
4360 let root = McpContext::with_state(Cx::for_testing(), 1, state);
4361 let clone_created_before_scope = root.clone();
4362 let (scoped, guard) = root
4363 .begin_request_scope()
4364 .expect("new context should create one request lease");
4365 let escaped = scoped.clone();
4366 drop(guard);
4367
4368 assert!(escaped.ensure_live().is_err());
4369 assert!(escaped.checkpoint().is_err());
4370 assert!(escaped.consume_cost(0).is_err());
4371 assert!(escaped.masked(|| 42).is_err());
4372 assert!(!escaped.set_auth(AuthContext::with_subject("late")));
4373 assert!(!escaped.set_state("late", true));
4374 assert!(escaped.auth().is_none());
4375 assert!(!escaped.can_call_tools());
4376 assert!(!escaped.can_read_resources());
4377 assert!(clone_created_before_scope.ensure_live().is_err());
4378
4379 assert!(clone_created_before_scope.begin_request_scope().is_none());
4380 }
4381
4382 #[test]
4383 fn test_mcp_context_masked_section() {
4384 let cx = Cx::for_testing();
4385 let ctx = McpContext::new(cx, 1);
4386
4387 let result = ctx.masked(|| 42).expect("mask should be admitted");
4389 assert_eq!(result, 42);
4390 }
4391
4392 #[test]
4393 fn test_mcp_context_budget() {
4394 let cx = Cx::for_testing();
4395 let ctx = McpContext::new(cx, 1);
4396
4397 let budget = ctx.budget();
4399 assert!(!budget.is_exhausted());
4401 }
4402
4403 #[test]
4404 fn budget_ceiling_is_monotone_and_visible_to_checkpoints() {
4405 let ambient_deadline = wall_now().saturating_add_nanos(5_000_000_000);
4406 let tighter_deadline = ambient_deadline.saturating_sub_nanos(1_000_000_000);
4407 let later_deadline = ambient_deadline.saturating_add_nanos(1_000_000_000);
4408 let cx = Cx::for_testing_with_budget(Budget::new().with_deadline(ambient_deadline));
4409 let ctx = McpContext::new(cx, 1)
4410 .with_budget_ceiling(Budget::new().with_deadline(tighter_deadline))
4411 .with_budget_ceiling(Budget::new().with_deadline(later_deadline));
4412
4413 assert_eq!(ctx.budget().deadline, Some(tighter_deadline));
4414 assert!(ctx.checkpoint().is_ok());
4415 }
4416
4417 #[test]
4418 fn operation_deadline_tightens_child_without_leaking_to_parent() {
4419 let parent_deadline = wall_now().saturating_add_nanos(5_000_000_000);
4420 let child_deadline = parent_deadline.saturating_sub_nanos(1_000_000_000);
4421 let parent = McpContext::new(Cx::for_testing(), 1)
4422 .with_budget_ceiling(Budget::new().with_deadline(parent_deadline));
4423 let child = parent.clone().with_operation_deadline(Some(child_deadline));
4424 let grandchild = child.clone().with_operation_deadline(None);
4425
4426 assert_eq!(parent.budget().deadline, Some(parent_deadline));
4427 assert_eq!(child.budget().deadline, Some(child_deadline));
4428 assert_eq!(grandchild.budget().deadline, Some(child_deadline));
4429 }
4430
4431 #[test]
4432 fn framework_poll_ceiling_drains_across_clones_at_n_plus_one() {
4433 const LIMIT: u32 = 3;
4434
4435 let ctx = McpContext::new(Cx::for_testing(), 1)
4436 .with_budget_ceiling(Budget::new().with_poll_quota(LIMIT));
4437 let clone = ctx.clone();
4438
4439 for admitted in 0..LIMIT {
4440 let result = if admitted % 2 == 0 {
4441 ctx.checkpoint()
4442 } else {
4443 clone.checkpoint()
4444 };
4445 assert!(result.is_ok(), "checkpoint {} should fit", admitted + 1);
4446 let expected = LIMIT - admitted - 1;
4447 assert_eq!(ctx.budget().poll_quota, expected);
4448 assert_eq!(clone.budget().poll_quota, expected);
4449 }
4450
4451 assert!(clone.checkpoint().is_err(), "checkpoint N+1 must fail");
4452 assert_eq!(ctx.budget().poll_quota, 0);
4453 assert!(!ctx.cx().is_cancel_requested());
4454 }
4455
4456 #[test]
4457 fn ambient_poll_budget_drains_across_clones_without_mutating_cx() {
4458 const LIMIT: u32 = 3;
4459
4460 let cx = Cx::for_testing_with_budget(Budget::new().with_poll_quota(LIMIT));
4461 let ctx = McpContext::new(cx.clone(), 1);
4462 let clone = ctx.clone();
4463
4464 for admitted in 0..LIMIT {
4465 let result = if admitted % 2 == 0 {
4466 ctx.checkpoint()
4467 } else {
4468 clone.checkpoint()
4469 };
4470 assert!(
4471 result.is_ok(),
4472 "ambient checkpoint {} should fit",
4473 admitted + 1
4474 );
4475 assert_eq!(ctx.budget().poll_quota, LIMIT - admitted - 1);
4476 }
4477
4478 let debits_before_rejection = ctx
4479 .budget_state
4480 .lock()
4481 .unwrap_or_else(std::sync::PoisonError::into_inner)
4482 .ambient_poll_debits;
4483 assert!(
4484 clone.checkpoint().is_err(),
4485 "ambient checkpoint N+1 must fail"
4486 );
4487 assert_eq!(
4488 ctx.budget_state
4489 .lock()
4490 .unwrap_or_else(std::sync::PoisonError::into_inner)
4491 .ambient_poll_debits,
4492 debits_before_rejection,
4493 "a rejected checkpoint must not partially debit the ledger"
4494 );
4495 assert_eq!(ctx.budget().poll_quota, 0);
4496 assert_eq!(cx.budget().poll_quota, LIMIT);
4497 assert!(!cx.is_cancel_requested());
4498 assert!(ctx.ensure_live().is_ok());
4499 }
4500
4501 #[test]
4502 fn tighter_ambient_poll_limit_does_not_debit_looser_ceiling_on_rejection() {
4503 let cx = Cx::for_testing_with_budget(Budget::new().with_poll_quota(2));
4504 let ctx =
4505 McpContext::new(cx.clone(), 1).with_budget_ceiling(Budget::new().with_poll_quota(3));
4506
4507 assert!(ctx.checkpoint().is_ok());
4508 assert!(ctx.checkpoint().is_ok());
4509 assert!(ctx.checkpoint().is_err());
4510
4511 let state = *ctx
4512 .budget_state
4513 .lock()
4514 .unwrap_or_else(std::sync::PoisonError::into_inner);
4515 assert_eq!(state.ambient_poll_debits, 2);
4516 assert_eq!(state.ceiling.map(|budget| budget.poll_quota), Some(1));
4517 assert_eq!(cx.budget().poll_quota, 2);
4518 }
4519
4520 #[test]
4521 fn framework_cost_ceiling_drains_across_clones_at_n_plus_one() {
4522 const LIMIT: u64 = 3;
4523
4524 let ctx = McpContext::new(Cx::for_testing(), 1)
4525 .with_budget_ceiling(Budget::new().with_cost_quota(LIMIT));
4526 let clone = ctx.clone();
4527
4528 for admitted in 0..LIMIT {
4529 let result = if admitted % 2 == 0 {
4530 ctx.consume_cost(1)
4531 } else {
4532 clone.consume_cost(1)
4533 };
4534 assert!(result.is_ok(), "cost debit {} should fit", admitted + 1);
4535 let expected = Some(LIMIT - admitted - 1);
4536 assert_eq!(ctx.budget().cost_quota, expected);
4537 assert_eq!(clone.budget().cost_quota, expected);
4538 }
4539
4540 assert!(clone.consume_cost(1).is_err(), "cost debit N+1 must fail");
4541 assert_eq!(ctx.budget().cost_quota, Some(0));
4542 assert!(!ctx.cx().is_cancel_requested());
4543 assert!(
4544 ctx.ensure_live().is_ok(),
4545 "an exactly admitted final debit is not an overrun"
4546 );
4547 }
4548
4549 #[test]
4550 fn framework_poll_and_cost_debits_are_independent() {
4551 let ctx = McpContext::new(Cx::for_testing(), 1)
4552 .with_budget_ceiling(Budget::new().with_poll_quota(2).with_cost_quota(2));
4553
4554 assert!(ctx.checkpoint().is_ok());
4555 assert_eq!(ctx.budget().poll_quota, 1);
4556 assert_eq!(ctx.budget().cost_quota, Some(2));
4557
4558 assert!(ctx.consume_cost(1).is_ok());
4559 assert_eq!(ctx.budget().poll_quota, 1);
4560 assert_eq!(ctx.budget().cost_quota, Some(1));
4561 }
4562
4563 #[test]
4564 fn exact_poll_depletion_is_live_until_the_next_poll_admission() {
4565 let ctx = McpContext::new(Cx::for_testing(), 1)
4566 .with_budget_ceiling(Budget::new().with_poll_quota(1));
4567
4568 assert!(ctx.checkpoint().is_ok());
4569 assert_eq!(ctx.budget().poll_quota, 0);
4570 assert!(ctx.ensure_live().is_ok());
4571 assert!(ctx.checkpoint().is_err());
4572 }
4573
4574 #[test]
4575 fn zero_framework_quotas_fail_without_cancelling_ambient_context() {
4576 let poll_ctx = McpContext::new(Cx::for_testing(), 1)
4577 .with_budget_ceiling(Budget::new().with_poll_quota(0));
4578 let cost_ctx = McpContext::new(Cx::for_testing(), 2)
4579 .with_budget_ceiling(Budget::new().with_cost_quota(0));
4580
4581 assert!(poll_ctx.checkpoint().is_err());
4582 assert_eq!(poll_ctx.budget().poll_quota, 0);
4583 assert!(!poll_ctx.cx().is_cancel_requested());
4584
4585 assert!(cost_ctx.consume_cost(0).is_ok());
4586 assert!(cost_ctx.consume_cost(1).is_err());
4587 assert_eq!(cost_ctx.budget().cost_quota, Some(0));
4588 assert!(!cost_ctx.cx().is_cancel_requested());
4589 }
4590
4591 #[test]
4592 fn oversized_framework_cost_debit_is_atomic() {
4593 let ctx = McpContext::new(Cx::for_testing(), 1)
4594 .with_budget_ceiling(Budget::new().with_cost_quota(2));
4595
4596 assert!(ctx.consume_cost(3).is_err());
4597 assert_eq!(ctx.budget().cost_quota, Some(2));
4598 assert!(ctx.consume_cost(2).is_ok());
4599 assert_eq!(ctx.budget().cost_quota, Some(0));
4600 assert!(ctx.consume_cost(1).is_err());
4601 }
4602
4603 #[test]
4604 fn zero_ambient_cost_quota_prevents_framework_cost_debit() {
4605 let ambient = Budget::new().with_cost_quota(0);
4606 let ctx = McpContext::new(Cx::for_testing_with_budget(ambient), 1)
4607 .with_budget_ceiling(Budget::new().with_cost_quota(3));
4608
4609 assert!(ctx.consume_cost(1).is_err());
4610 assert_eq!(ctx.budget().cost_quota, Some(0));
4611 }
4612
4613 #[test]
4614 fn positive_ambient_cost_quota_drains_cumulatively_across_clones() {
4615 const LIMIT: u64 = 3;
4616 let ambient = Budget::new().with_cost_quota(LIMIT);
4617 let ctx = McpContext::new(Cx::for_testing_with_budget(ambient), 1);
4618 let clone = ctx.clone();
4619
4620 for admitted in 0..LIMIT {
4621 let result = if admitted % 2 == 0 {
4622 ctx.consume_cost(1)
4623 } else {
4624 clone.consume_cost(1)
4625 };
4626 assert!(result.is_ok(), "ambient debit {} should fit", admitted + 1);
4627 assert_eq!(ctx.budget().cost_quota, Some(LIMIT - admitted - 1));
4628 }
4629
4630 assert!(
4631 clone.consume_cost(1).is_err(),
4632 "ambient debit N+1 must fail"
4633 );
4634 assert_eq!(ctx.budget().cost_quota, Some(0));
4635 assert_eq!(
4636 ctx.cx().budget().cost_quota,
4637 Some(LIMIT),
4638 "request-local accounting must not mutate the caller-owned Cx"
4639 );
4640 }
4641
4642 #[test]
4643 fn rejected_cost_debit_does_not_record_an_ambient_checkpoint() {
4644 let cx = Cx::for_testing_with_budget(Budget::new().with_cost_quota(2));
4645 let ctx = McpContext::new(cx, 1);
4646 let before = ctx.cx().checkpoint_state().checkpoint_count;
4647
4648 assert!(ctx.consume_cost(3).is_err());
4649 assert_eq!(ctx.cx().checkpoint_state().checkpoint_count, before);
4650 assert_eq!(ctx.budget().cost_quota, Some(2));
4651 }
4652
4653 #[test]
4654 fn zero_cost_debit_observes_explicit_cancellation() {
4655 let cx = Cx::for_testing();
4656 cx.set_cancel_requested(true);
4657 let ctx = McpContext::new(cx, 1);
4658
4659 assert!(ctx.consume_cost(0).is_err());
4660 }
4661
4662 #[test]
4663 fn expired_request_ceiling_fails_without_cancelling_ambient_context() {
4664 let cx = Cx::for_testing();
4665 let ctx = McpContext::new(cx, 1)
4666 .with_budget_ceiling(Budget::new().with_deadline(asupersync::Time::ZERO));
4667
4668 assert!(ctx.is_cancelled());
4669 assert!(ctx.checkpoint().is_err());
4670 assert!(!ctx.cx().is_cancel_requested());
4671 }
4672
4673 #[test]
4674 fn framework_budget_ceiling_is_deferred_while_masked() {
4675 let ctx = McpContext::new(Cx::for_testing(), 1)
4676 .with_budget_ceiling(Budget::new().with_deadline(asupersync::Time::ZERO));
4677
4678 assert!(
4679 ctx.masked(|| ctx.checkpoint())
4680 .expect("mask should be admitted")
4681 .is_ok()
4682 );
4683 assert!(ctx.checkpoint().is_err());
4684 }
4685
4686 #[test]
4687 fn framework_poll_debits_continue_while_enforcement_is_masked() {
4688 let ctx = McpContext::new(Cx::for_testing(), 1)
4689 .with_budget_ceiling(Budget::new().with_poll_quota(1));
4690
4691 ctx.masked(|| {
4692 assert!(ctx.checkpoint().is_ok());
4693 assert_eq!(ctx.budget().poll_quota, 0);
4694 assert!(ctx.checkpoint().is_ok());
4695 })
4696 .expect("mask should be admitted");
4697
4698 assert!(ctx.checkpoint().is_err());
4699 }
4700
4701 #[test]
4702 fn masked_cost_overage_saturates_framework_ceiling() {
4703 let ctx = McpContext::new(Cx::for_testing(), 1)
4704 .with_budget_ceiling(Budget::new().with_cost_quota(2));
4705
4706 assert!(
4707 ctx.masked(|| ctx.consume_cost(3))
4708 .expect("mask should be admitted")
4709 .is_ok()
4710 );
4711 assert_eq!(ctx.budget().cost_quota, Some(0));
4712 assert!(ctx.ensure_live().is_err());
4713 assert!(ctx.consume_cost(1).is_err());
4714 }
4715
4716 #[test]
4717 fn masked_exact_cost_depletion_does_not_become_a_deferred_overrun() {
4718 let ctx = McpContext::new(Cx::for_testing(), 1)
4719 .with_budget_ceiling(Budget::new().with_cost_quota(2));
4720
4721 assert!(
4722 ctx.masked(|| ctx.consume_cost(2))
4723 .expect("mask should be admitted")
4724 .is_ok()
4725 );
4726 assert_eq!(ctx.budget().cost_quota, Some(0));
4727 assert!(ctx.ensure_live().is_ok());
4728 assert!(ctx.consume_cost(0).is_ok());
4729 assert!(ctx.consume_cost(1).is_err());
4730 }
4731
4732 #[test]
4733 fn masked_cost_overage_saturates_tighter_ambient_quota() {
4734 let ambient = Budget::new().with_cost_quota(2);
4735 let ctx = McpContext::new(Cx::for_testing_with_budget(ambient), 1)
4736 .with_budget_ceiling(Budget::new().with_cost_quota(10));
4737
4738 assert!(
4739 ctx.masked(|| ctx.consume_cost(3))
4740 .expect("mask should be admitted")
4741 .is_ok()
4742 );
4743 assert_eq!(ctx.budget().cost_quota, Some(0));
4744 assert_eq!(
4745 ctx.budget_state
4746 .lock()
4747 .unwrap_or_else(std::sync::PoisonError::into_inner)
4748 .ceiling
4749 .and_then(|budget| budget.cost_quota),
4750 Some(7),
4751 "the looser framework ceiling is still debited independently"
4752 );
4753 assert!(ctx.consume_cost(1).is_err());
4754 }
4755
4756 #[test]
4757 fn framework_mask_is_shared_with_clones_and_restored_after_exit() {
4758 let ctx = McpContext::new(Cx::for_testing(), 1)
4759 .with_budget_ceiling(Budget::new().with_poll_quota(0));
4760 let clone = ctx.clone();
4761
4762 assert!(
4763 ctx.masked(|| clone.checkpoint())
4764 .expect("mask should be admitted")
4765 .is_ok()
4766 );
4767 assert!(clone.checkpoint().is_err());
4768 }
4769
4770 #[test]
4771 fn framework_mask_depth_is_restored_after_unwind() {
4772 let ctx = McpContext::new(Cx::for_testing(), 1)
4773 .with_budget_ceiling(Budget::new().with_deadline(asupersync::Time::ZERO));
4774
4775 let _ = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
4776 let _ = ctx.masked(|| panic!("test-only masked-section panic"));
4777 }));
4778
4779 assert!(ctx.checkpoint().is_err());
4780 assert_eq!(ctx.framework_mask_depth.load(Ordering::SeqCst), 0);
4781 }
4782
4783 #[test]
4784 fn test_cancelled_error_display() {
4785 let err = CancelledError;
4786 assert_eq!(err.to_string(), "request cancelled");
4787 }
4788
4789 #[test]
4790 fn handler_log_respects_client_floor_and_missing_floor() {
4791 let captured = Arc::new(Mutex::new(Vec::new()));
4792 struct CaptureSender(Arc<Mutex<Vec<(McpLogLevel, String)>>>);
4793 impl NotificationSender for CaptureSender {
4794 fn send_progress(&self, _progress: f64, _total: Option<f64>, _message: Option<&str>) {}
4795 fn send_log(&self, level: McpLogLevel, _logger: Option<&str>, data: serde_json::Value) {
4796 self.0
4797 .lock()
4798 .unwrap_or_else(std::sync::PoisonError::into_inner)
4799 .push((level, data.as_str().unwrap_or_default().to_owned()));
4800 }
4801 }
4802
4803 let silent = McpContext::new(Cx::for_testing(), 1)
4804 .with_log_sender(Arc::new(CaptureSender(Arc::clone(&captured))));
4805 silent.info("before-floor");
4806 assert!(captured.lock().expect("lock").is_empty());
4807
4808 let ctx = silent.with_min_log_level(Some(McpLogLevel::Info));
4809 assert_eq!(ctx.min_log_level(), Some(McpLogLevel::Info));
4810 ctx.debug("too-low");
4811 ctx.info("admitted");
4812 ctx.warning("also-admitted");
4813 let emitted = captured.lock().expect("lock").clone();
4814 assert_eq!(
4815 emitted,
4816 vec![
4817 (McpLogLevel::Info, "admitted".to_owned()),
4818 (McpLogLevel::Warning, "also-admitted".to_owned()),
4819 ]
4820 );
4821 }
4822
4823 #[test]
4824 fn catalog_change_emits_only_when_the_disabled_set_mutates() {
4825 let captured = Arc::new(Mutex::new(Vec::new()));
4826 struct CaptureSender(Arc<Mutex<Vec<McpCatalogKind>>>);
4827 impl NotificationSender for CaptureSender {
4828 fn send_progress(&self, _progress: f64, _total: Option<f64>, _message: Option<&str>) {}
4829 fn send_catalog_changed(&self, kind: McpCatalogKind) {
4830 self.0
4831 .lock()
4832 .unwrap_or_else(std::sync::PoisonError::into_inner)
4833 .push(kind);
4834 }
4835 }
4836
4837 let ctx = McpContext::with_state(Cx::for_testing(), 1, SessionState::new())
4838 .with_log_sender(Arc::new(CaptureSender(Arc::clone(&captured))));
4839 assert!(ctx.disable_tool("admin"));
4840 assert!(ctx.disable_tool("admin"));
4841 assert!(ctx.enable_tool("admin"));
4842 assert!(ctx.enable_tool("admin"));
4843 assert!(ctx.disable_resource("file://secret"));
4844 assert!(ctx.disable_prompt("hidden"));
4845 assert_eq!(
4846 *captured.lock().expect("lock"),
4847 vec![
4848 McpCatalogKind::Tools,
4849 McpCatalogKind::Tools,
4850 McpCatalogKind::Resources,
4851 McpCatalogKind::Prompts,
4852 ]
4853 );
4854 }
4855
4856 #[test]
4857 fn catalog_publisher_receives_mutations_even_without_a_session_sender() {
4858 let captured = Arc::new(Mutex::new(Vec::new()));
4859 struct CapturePublisher(Arc<Mutex<Vec<McpCatalogKind>>>);
4860 impl CatalogChangePublisher for CapturePublisher {
4861 fn publish_catalog_changed(&self, kind: McpCatalogKind) -> bool {
4862 self.0
4863 .lock()
4864 .unwrap_or_else(std::sync::PoisonError::into_inner)
4865 .push(kind);
4866 true
4867 }
4868 fn publish_resource_updated(&self, _uri: &str) -> bool {
4869 false
4870 }
4871 }
4872
4873 let ctx = McpContext::with_state(Cx::for_testing(), 1, SessionState::new())
4874 .with_catalog_publisher(Arc::new(CapturePublisher(Arc::clone(&captured))));
4875 assert!(ctx.disable_tool("admin"));
4876 assert!(ctx.disable_tool("admin"));
4877 assert_eq!(*captured.lock().expect("lock"), vec![McpCatalogKind::Tools]);
4878 }
4879
4880 #[test]
4881 fn notify_resource_updated_requires_a_live_subscription() {
4882 let captured = Arc::new(Mutex::new(Vec::new()));
4883 struct CaptureSender(Arc<Mutex<Vec<String>>>);
4884 impl NotificationSender for CaptureSender {
4885 fn send_progress(&self, _progress: f64, _total: Option<f64>, _message: Option<&str>) {}
4886 fn send_resource_updated(&self, uri: &str) {
4887 self.0
4888 .lock()
4889 .unwrap_or_else(std::sync::PoisonError::into_inner)
4890 .push(uri.to_owned());
4891 }
4892 }
4893
4894 let ctx = McpContext::new(Cx::for_testing(), 1)
4895 .with_log_sender(Arc::new(CaptureSender(Arc::clone(&captured))))
4896 .with_resource_subscriptions(["file:///watched.txt"]);
4897 assert!(!ctx.notify_resource_updated("file:///other.txt"));
4898 assert!(ctx.notify_resource_updated("file:///watched.txt"));
4899 assert_eq!(
4900 *captured.lock().expect("lock"),
4901 vec!["file:///watched.txt".to_owned()]
4902 );
4903 }
4904
4905 #[test]
4906 fn test_into_outcome_ok() {
4907 let result: Result<i32, CancelledError> = Ok(42);
4908 let outcome: Outcome<i32, CancelledError> = result.into_outcome();
4909 assert!(matches!(outcome, Outcome::Ok(42)));
4910 }
4911
4912 #[test]
4913 fn test_into_outcome_cancelled() {
4914 let result: Result<i32, CancelledError> = Err(CancelledError);
4915 let outcome: Outcome<i32, ()> = result.into_outcome();
4916 assert!(matches!(outcome, Outcome::Cancelled(_)));
4917 }
4918
4919 #[test]
4920 fn test_mcp_context_no_progress_reporter_by_default() {
4921 let cx = Cx::for_testing();
4922 let ctx = McpContext::new(cx, 1);
4923 assert!(!ctx.has_progress_reporter());
4924 }
4925
4926 #[test]
4927 fn test_mcp_context_with_progress_reporter() {
4928 let cx = Cx::for_testing();
4929 let sender = Arc::new(NoOpNotificationSender);
4930 let reporter = ProgressReporter::new(sender);
4931 let ctx = McpContext::with_progress(cx, 1, reporter);
4932 assert!(ctx.has_progress_reporter());
4933 }
4934
4935 #[test]
4936 fn progress_reporter_builder_preserves_request_accounting_domain() {
4937 let ctx = McpContext::new(Cx::for_testing(), 1)
4938 .with_budget_ceiling(Budget::new().with_cost_quota(5));
4939 let reporter = ProgressReporter::new(Arc::new(NoOpNotificationSender));
4940 let derived = ctx.clone().with_progress_reporter(reporter);
4941
4942 assert!(derived.has_progress_reporter());
4943 assert!(!ctx.has_progress_reporter());
4944 assert!(ctx.consume_cost(3).is_ok());
4945 assert_eq!(derived.budget().cost_quota, Some(2));
4946 }
4947
4948 #[test]
4949 fn isolated_auth_stages_identity_without_handler_capabilities() {
4950 let root = McpContext::with_state(Cx::for_testing(), 1, SessionState::new())
4951 .with_budget_ceiling(Budget::new().with_cost_quota(2))
4952 .with_sampling(Arc::new(NoOpSamplingSender))
4953 .with_elicitation(Arc::new(NoOpElicitationSender))
4954 .with_roots_provider(Arc::new(FixedRootsProvider));
4955 let staged = root.clone().with_isolated_auth();
4956
4957 assert!(staged.auth().is_none());
4958 assert!(!staged.has_session_state());
4959 assert!(!staged.can_sample());
4960 assert!(!staged.can_elicit());
4961 assert!(!staged.can_list_roots());
4962 assert!(!staged.can_read_resources());
4963 assert!(!staged.can_call_tools());
4964 assert!(staged.set_auth(AuthContext::with_subject("tentative")));
4965 assert_eq!(
4966 staged.auth().and_then(|auth| auth.subject),
4967 Some("tentative".to_string())
4968 );
4969 assert_eq!(root.auth().and_then(|auth| auth.subject), None);
4970
4971 assert!(root.set_auth(AuthContext::with_subject("committed")));
4972 let attempted_reisolation = root.clone().with_isolated_auth();
4973 assert_eq!(
4974 attempted_reisolation.auth().and_then(|auth| auth.subject),
4975 Some("committed".to_string())
4976 );
4977
4978 assert!(staged.consume_cost(1).is_ok());
4979 assert_eq!(root.budget().cost_quota, Some(1));
4980 }
4981
4982 #[test]
4983 fn committed_anonymous_auth_is_hidden_and_write_once() {
4984 let ctx = McpContext::with_state(Cx::for_testing(), 1, SessionState::new());
4985
4986 assert!(ctx.commit_anonymous_auth());
4987 assert!(ctx.auth().is_none());
4988 assert!(matches!(ctx.cache_auth_partition(), Some(None)));
4989 assert!(!ctx.set_auth(AuthContext::with_subject("forged")));
4990 assert!(!ctx.commit_anonymous_auth());
4991
4992 let clone = ctx.clone();
4993 assert!(clone.auth().is_none());
4994 assert!(matches!(clone.cache_auth_partition(), Some(None)));
4995 }
4996
4997 #[test]
4998 fn authenticated_cache_partition_contains_committed_facts() {
4999 let ctx = McpContext::new(Cx::for_testing(), 1);
5000 assert!(ctx.set_auth(AuthContext::with_subject("alice")));
5001
5002 let Some(Some(auth)) = ctx.cache_auth_partition() else {
5003 panic!("authenticated admission must expose cache partition facts");
5004 };
5005 assert_eq!(auth.subject.as_deref(), Some("alice"));
5006 }
5007
5008 #[test]
5009 fn test_report_progress_without_reporter() {
5010 let cx = Cx::for_testing();
5011 let ctx = McpContext::new(cx, 1);
5012 ctx.report_progress(0.5, Some("test"));
5014 ctx.report_progress_with_total(5.0, 10.0, None);
5015 }
5016
5017 #[test]
5018 fn test_report_progress_with_reporter() {
5019 use std::sync::atomic::{AtomicU32, Ordering};
5020
5021 struct CountingSender {
5022 count: AtomicU32,
5023 }
5024
5025 impl NotificationSender for CountingSender {
5026 fn send_progress(&self, _progress: f64, _total: Option<f64>, _message: Option<&str>) {
5027 self.count.fetch_add(1, Ordering::SeqCst);
5028 }
5029 }
5030
5031 let cx = Cx::for_testing();
5032 let sender = Arc::new(CountingSender {
5033 count: AtomicU32::new(0),
5034 });
5035 let reporter = ProgressReporter::new(sender.clone());
5036 let ctx = McpContext::with_progress(cx, 1, reporter);
5037
5038 ctx.report_progress(0.25, Some("step 1"));
5039 ctx.report_progress(0.5, None);
5040 ctx.report_progress_with_total(3.0, 4.0, Some("step 3"));
5041
5042 assert_eq!(sender.count.load(Ordering::SeqCst), 3);
5043 }
5044
5045 #[test]
5046 fn request_local_cancellation_suppresses_subsequent_progress() {
5047 use std::sync::atomic::{AtomicU32, Ordering};
5048
5049 struct CountingSender {
5050 count: AtomicU32,
5051 }
5052
5053 impl NotificationSender for CountingSender {
5054 fn send_progress(&self, _progress: f64, _total: Option<f64>, _message: Option<&str>) {
5055 self.count.fetch_add(1, Ordering::SeqCst);
5056 }
5057 }
5058
5059 let sender = Arc::new(CountingSender {
5060 count: AtomicU32::new(0),
5061 });
5062 let cancellation = McpRequestCancellation::new();
5063 let ctx =
5064 McpContext::with_progress(Cx::for_testing(), 1, ProgressReporter::new(sender.clone()))
5065 .with_request_cancellation(cancellation.clone());
5066
5067 ctx.report_progress(0.25, Some("before cancellation"));
5068 assert!(cancellation.cancel());
5069 ctx.report_progress(0.5, Some("after cancellation"));
5070
5071 assert_eq!(sender.count.load(Ordering::SeqCst), 1);
5072 assert!(!ctx.has_progress_reporter());
5073 }
5074
5075 #[test]
5076 fn test_progress_reporter_debug() {
5077 let sender = Arc::new(NoOpNotificationSender);
5078 let reporter = ProgressReporter::new(sender);
5079 let debug = format!("{reporter:?}");
5080 assert!(debug.contains("ProgressReporter"));
5081 }
5082
5083 #[test]
5084 fn test_noop_notification_sender() {
5085 let sender = NoOpNotificationSender;
5086 sender.send_progress(0.5, Some(1.0), Some("test"));
5088 }
5089
5090 #[test]
5092 fn test_mcp_context_no_session_state_by_default() {
5093 let cx = Cx::for_testing();
5094 let ctx = McpContext::new(cx, 1);
5095 assert!(!ctx.has_session_state());
5096 }
5097
5098 #[test]
5099 fn test_mcp_context_with_session_state() {
5100 let cx = Cx::for_testing();
5101 let state = SessionState::new();
5102 let ctx = McpContext::with_state(cx, 1, state);
5103 assert!(ctx.has_session_state());
5104 }
5105
5106 #[test]
5107 fn cache_admission_fails_if_session_state_changes_before_completion() {
5108 let state = SessionState::new();
5109 let ctx = McpContext::with_state(Cx::for_testing(), 1, state.clone());
5110 let admitted = ctx
5111 .begin_session_cache_partition()
5112 .expect("test platform must provide cache-partition entropy");
5113 assert_eq!(ctx.complete_session_cache_partition(), Some(admitted));
5114
5115 assert!(state.set("changed", true));
5116 assert!(ctx.complete_session_cache_partition().is_none());
5117 assert!(ctx.begin_session_cache_partition().is_none());
5118 }
5119
5120 #[test]
5121 fn response_cache_hit_markers_are_middleware_specific() {
5122 let ctx = McpContext::new(Cx::for_testing(), 1);
5123 assert!(ctx.mark_response_cache_hit(10));
5124 assert!(ctx.response_was_cache_hit(10));
5125 assert!(!ctx.response_was_cache_hit(11));
5126 assert!(!ctx.mark_response_cache_hit(0));
5127 }
5128
5129 #[test]
5130 fn test_mcp_context_get_set_state() {
5131 let cx = Cx::for_testing();
5132 let state = SessionState::new();
5133 let ctx = McpContext::with_state(cx, 1, state);
5134
5135 assert!(ctx.set_state("counter", 42));
5137
5138 let value: Option<i32> = ctx.get_state("counter");
5140 assert_eq!(value, Some(42));
5141 }
5142
5143 #[test]
5144 fn test_mcp_context_state_not_available() {
5145 let cx = Cx::for_testing();
5146 let ctx = McpContext::new(cx, 1);
5147
5148 assert!(!ctx.set_state("key", "value"));
5150
5151 let value: Option<String> = ctx.get_state("key");
5153 assert!(value.is_none());
5154 }
5155
5156 #[test]
5157 fn test_mcp_context_has_state() {
5158 let cx = Cx::for_testing();
5159 let state = SessionState::new();
5160 let ctx = McpContext::with_state(cx, 1, state);
5161
5162 assert!(!ctx.has_state("missing"));
5163
5164 ctx.set_state("present", true);
5165 assert!(ctx.has_state("present"));
5166 }
5167
5168 #[test]
5169 fn test_mcp_context_remove_state() {
5170 let cx = Cx::for_testing();
5171 let state = SessionState::new();
5172 let ctx = McpContext::with_state(cx, 1, state);
5173
5174 ctx.set_state("key", "value");
5175 assert!(ctx.has_state("key"));
5176
5177 let removed = ctx.remove_state("key");
5178 assert!(removed.is_some());
5179 assert!(!ctx.has_state("key"));
5180 }
5181
5182 #[test]
5183 fn test_mcp_context_with_state_and_progress() {
5184 let cx = Cx::for_testing();
5185 let state = SessionState::new();
5186 let sender = Arc::new(NoOpNotificationSender);
5187 let reporter = ProgressReporter::new(sender);
5188
5189 let ctx = McpContext::with_state_and_progress(cx, 1, state, reporter);
5190
5191 assert!(ctx.has_session_state());
5192 assert!(ctx.has_progress_reporter());
5193 }
5194
5195 #[test]
5196 fn test_mcp_context_auth_is_request_local() {
5197 let cx = Cx::for_testing();
5198 let state = SessionState::new();
5199 let ctx = McpContext::with_state(cx, 1, state.clone());
5200
5201 assert!(ctx.set_auth(AuthContext::with_subject("alice")));
5202
5203 assert_eq!(
5204 ctx.auth().and_then(|auth| auth.subject),
5205 Some("alice".to_string())
5206 );
5207 assert!(
5208 state.is_empty(),
5209 "request auth must not be persisted into session state"
5210 );
5211 }
5212
5213 #[test]
5214 fn test_mcp_context_clones_share_request_auth() {
5215 let cx = Cx::for_testing();
5216 let ctx = McpContext::new(cx, 1);
5217 let cloned = ctx.clone();
5218
5219 assert!(cloned.set_auth(AuthContext::with_subject("bob")));
5220
5221 assert_eq!(
5222 ctx.auth().and_then(|auth| auth.subject),
5223 Some("bob".to_string())
5224 );
5225 }
5226
5227 #[test]
5228 fn committed_request_auth_is_write_once_across_clones() {
5229 let ctx =
5230 McpContext::new(Cx::for_testing(), 1).with_auth(AuthContext::with_subject("verified"));
5231 let clone = ctx.clone();
5232
5233 assert!(!clone.set_auth(AuthContext::with_subject("replacement")));
5234 assert_eq!(
5235 ctx.auth().and_then(|auth| auth.subject),
5236 Some("verified".to_string())
5237 );
5238 }
5239
5240 #[test]
5241 fn test_new_mcp_contexts_do_not_share_request_auth_even_with_same_cx() {
5242 let cx = Cx::for_testing();
5243 let state = SessionState::new();
5244 let first = McpContext::with_state(cx.clone(), 7, state.clone());
5245 let second = McpContext::with_state(cx, 7, state);
5246
5247 assert!(first.set_auth(AuthContext::with_subject("carol")));
5248
5249 assert!(second.auth().is_none());
5250 }
5251
5252 #[test]
5253 fn test_new_mcp_contexts_do_not_share_request_auth_across_requests() {
5254 let state = SessionState::new();
5255 let first = McpContext::with_state(Cx::for_testing(), 7, state.clone());
5256 let second = McpContext::with_state(Cx::for_testing(), 8, state);
5257
5258 assert!(first.set_auth(AuthContext::with_subject("dave")));
5259
5260 assert_eq!(
5261 first.auth().and_then(|auth| auth.subject),
5262 Some("dave".to_string())
5263 );
5264 assert!(second.auth().is_none());
5265 }
5266
5267 #[test]
5268 fn test_mcp_context_drop_does_not_leak_request_auth() {
5269 let cx = Cx::for_testing();
5270
5271 {
5272 let ctx = McpContext::new(cx.clone(), 9);
5273 assert!(ctx.set_auth(AuthContext::with_subject("erin")));
5274 }
5275
5276 assert!(
5277 McpContext::new(cx, 9).auth().is_none(),
5278 "fresh contexts must start without inherited request auth"
5279 );
5280 }
5281
5282 #[test]
5287 fn test_mcp_context_tools_enabled_by_default() {
5288 let cx = Cx::for_testing();
5289 let state = SessionState::new();
5290 let ctx = McpContext::with_state(cx, 1, state);
5291
5292 assert!(ctx.is_tool_enabled("any_tool"));
5293 assert!(ctx.is_tool_enabled("another_tool"));
5294 }
5295
5296 #[test]
5297 fn test_mcp_context_disable_enable_tool() {
5298 let cx = Cx::for_testing();
5299 let state = SessionState::new();
5300 let ctx = McpContext::with_state(cx, 1, state);
5301
5302 assert!(ctx.is_tool_enabled("my_tool"));
5304
5305 assert!(ctx.disable_tool("my_tool"));
5307 assert!(!ctx.is_tool_enabled("my_tool"));
5308 assert!(ctx.is_tool_enabled("other_tool"));
5309
5310 assert!(ctx.enable_tool("my_tool"));
5312 assert!(ctx.is_tool_enabled("my_tool"));
5313 }
5314
5315 #[test]
5316 fn test_mcp_context_disable_enable_resource() {
5317 let cx = Cx::for_testing();
5318 let state = SessionState::new();
5319 let ctx = McpContext::with_state(cx, 1, state);
5320
5321 assert!(ctx.is_resource_enabled("file://secret"));
5323
5324 assert!(ctx.disable_resource("file://secret"));
5326 assert!(!ctx.is_resource_enabled("file://secret"));
5327 assert!(ctx.is_resource_enabled("file://public"));
5328
5329 assert!(ctx.enable_resource("file://secret"));
5331 assert!(ctx.is_resource_enabled("file://secret"));
5332 }
5333
5334 #[test]
5335 fn test_mcp_context_disable_enable_prompt() {
5336 let cx = Cx::for_testing();
5337 let state = SessionState::new();
5338 let ctx = McpContext::with_state(cx, 1, state);
5339
5340 assert!(ctx.is_prompt_enabled("admin_prompt"));
5342
5343 assert!(ctx.disable_prompt("admin_prompt"));
5345 assert!(!ctx.is_prompt_enabled("admin_prompt"));
5346 assert!(ctx.is_prompt_enabled("user_prompt"));
5347
5348 assert!(ctx.enable_prompt("admin_prompt"));
5350 assert!(ctx.is_prompt_enabled("admin_prompt"));
5351 }
5352
5353 #[test]
5354 fn test_mcp_context_disable_multiple_tools() {
5355 let cx = Cx::for_testing();
5356 let state = SessionState::new();
5357 let ctx = McpContext::with_state(cx, 1, state);
5358
5359 ctx.disable_tool("tool1");
5360 ctx.disable_tool("tool2");
5361 ctx.disable_tool("tool3");
5362
5363 assert!(!ctx.is_tool_enabled("tool1"));
5364 assert!(!ctx.is_tool_enabled("tool2"));
5365 assert!(!ctx.is_tool_enabled("tool3"));
5366 assert!(ctx.is_tool_enabled("tool4"));
5367
5368 let disabled = ctx.disabled_tools();
5369 assert_eq!(disabled.len(), 3);
5370 assert!(disabled.contains("tool1"));
5371 assert!(disabled.contains("tool2"));
5372 assert!(disabled.contains("tool3"));
5373 }
5374
5375 #[test]
5376 fn test_mcp_context_disabled_sets_empty_by_default() {
5377 let cx = Cx::for_testing();
5378 let state = SessionState::new();
5379 let ctx = McpContext::with_state(cx, 1, state);
5380
5381 assert!(ctx.disabled_tools().is_empty());
5382 assert!(ctx.disabled_resources().is_empty());
5383 assert!(ctx.disabled_prompts().is_empty());
5384 }
5385
5386 #[test]
5387 fn test_mcp_context_enable_disable_no_state() {
5388 let cx = Cx::for_testing();
5389 let ctx = McpContext::new(cx, 1);
5390
5391 assert!(!ctx.disable_tool("tool"));
5393 assert!(!ctx.enable_tool("tool"));
5394
5395 assert!(ctx.is_tool_enabled("tool"));
5397 }
5398
5399 #[test]
5400 fn test_mcp_context_disabled_state_persists_across_contexts() {
5401 let state = SessionState::new();
5402
5403 {
5405 let cx = Cx::for_testing();
5406 let ctx = McpContext::with_state(cx, 1, state.clone());
5407 ctx.disable_tool("shared_tool");
5408 }
5409
5410 {
5412 let cx = Cx::for_testing();
5413 let ctx = McpContext::with_state(cx, 2, state.clone());
5414 assert!(!ctx.is_tool_enabled("shared_tool"));
5415 }
5416 }
5417
5418 #[test]
5423 fn test_mcp_context_no_capabilities_by_default() {
5424 let cx = Cx::for_testing();
5425 let ctx = McpContext::new(cx, 1);
5426
5427 assert!(ctx.client_capabilities().is_none());
5428 assert!(ctx.server_capabilities().is_none());
5429 assert!(!ctx.client_supports_sampling());
5430 assert!(!ctx.client_supports_elicitation());
5431 assert!(!ctx.client_supports_roots());
5432 }
5433
5434 #[test]
5435 fn test_mcp_context_with_client_capabilities() {
5436 let cx = Cx::for_testing();
5437 let caps = ClientCapabilityInfo::new()
5438 .with_sampling()
5439 .with_elicitation(true, false)
5440 .with_roots(true);
5441
5442 let ctx = McpContext::new(cx, 1).with_client_capabilities(caps);
5443
5444 assert!(ctx.client_capabilities().is_some());
5445 assert!(ctx.client_supports_sampling());
5446 assert!(ctx.client_supports_elicitation());
5447 assert!(ctx.client_supports_elicitation_form());
5448 assert!(!ctx.client_supports_elicitation_url());
5449 assert!(ctx.client_supports_roots());
5450 }
5451
5452 #[test]
5453 fn test_mcp_context_with_client_implementation() {
5454 let cx = Cx::for_testing();
5455 let mut identity = ClientImplementationInfo::new("e2e-client", "1.0.0");
5456 identity.title = Some("Client Title".to_owned());
5457 let ctx = McpContext::new(cx, 1).with_client_implementation(identity);
5458 let observed = ctx
5459 .client_implementation()
5460 .expect("the attached identity must be retained");
5461 assert_eq!(observed.name, "e2e-client");
5462 assert_eq!(observed.title.as_deref(), Some("Client Title"));
5463 assert!(observed.has_extras());
5464 let bare = McpContext::new(Cx::for_testing(), 2);
5465 assert!(bare.client_implementation().is_none());
5466 }
5467
5468 #[test]
5469 fn test_mcp_context_with_server_capabilities() {
5470 let cx = Cx::for_testing();
5471 let caps = ServerCapabilityInfo::new()
5472 .with_tools()
5473 .with_resources(true)
5474 .with_prompts()
5475 .with_logging();
5476
5477 let ctx = McpContext::new(cx, 1).with_server_capabilities(caps);
5478
5479 let server_caps = ctx.server_capabilities().unwrap();
5480 assert!(server_caps.tools);
5481 assert!(server_caps.resources);
5482 assert!(server_caps.resources_subscribe);
5483 assert!(server_caps.prompts);
5484 assert!(server_caps.logging);
5485 }
5486
5487 #[test]
5488 fn test_client_capability_info_builders() {
5489 let caps = ClientCapabilityInfo::new();
5490 assert!(!caps.sampling);
5491 assert!(!caps.elicitation);
5492 assert!(!caps.roots);
5493
5494 let caps = caps.with_sampling();
5495 assert!(caps.sampling);
5496
5497 let caps = ClientCapabilityInfo::new().with_elicitation(true, true);
5498 assert!(caps.elicitation);
5499 assert!(caps.elicitation_form);
5500 assert!(caps.elicitation_url);
5501
5502 let caps = ClientCapabilityInfo::new().with_roots(false);
5503 assert!(caps.roots);
5504 assert!(!caps.roots_list_changed);
5505 }
5506
5507 #[test]
5508 fn test_server_capability_info_builders() {
5509 let caps = ServerCapabilityInfo::new();
5510 assert!(!caps.tools);
5511 assert!(!caps.resources);
5512 assert!(!caps.prompts);
5513 assert!(!caps.logging);
5514
5515 let caps = caps
5516 .with_tools()
5517 .with_resources(false)
5518 .with_prompts()
5519 .with_logging();
5520 assert!(caps.tools);
5521 assert!(caps.resources);
5522 assert!(!caps.resources_subscribe);
5523 assert!(caps.prompts);
5524 assert!(caps.logging);
5525 }
5526
5527 #[test]
5532 fn test_resource_content_item_text() {
5533 let item = ResourceContentItem::text("test://uri", "hello");
5534 assert_eq!(item.uri, "test://uri");
5535 assert_eq!(item.mime_type.as_deref(), Some("text/plain"));
5536 assert_eq!(item.as_text(), Some("hello"));
5537 assert!(item.as_blob().is_none());
5538 assert!(item.is_text());
5539 assert!(!item.is_blob());
5540 }
5541
5542 #[test]
5543 fn test_resource_content_item_json() {
5544 let item = ResourceContentItem::json("data://config", r#"{"key":"val"}"#);
5545 assert_eq!(item.uri, "data://config");
5546 assert_eq!(item.mime_type.as_deref(), Some("application/json"));
5547 assert_eq!(item.as_text(), Some(r#"{"key":"val"}"#));
5548 assert!(item.is_text());
5549 assert!(!item.is_blob());
5550 }
5551
5552 #[test]
5553 fn test_resource_content_item_blob() {
5554 let item = ResourceContentItem::blob("binary://data", "application/octet-stream", "AQID");
5555 assert_eq!(item.uri, "binary://data");
5556 assert_eq!(item.mime_type.as_deref(), Some("application/octet-stream"));
5557 assert!(item.as_text().is_none());
5558 assert_eq!(item.as_blob(), Some("AQID"));
5559 assert!(!item.is_text());
5560 assert!(item.is_blob());
5561 }
5562
5563 #[test]
5568 fn test_resource_read_result_text() {
5569 let result = ResourceReadResult::text("test://doc", "content");
5570 assert_eq!(result.first_text(), Some("content"));
5571 assert!(result.first_blob().is_none());
5572 assert_eq!(result.contents.len(), 1);
5573 }
5574
5575 #[test]
5576 fn test_resource_read_result_new_multiple() {
5577 let result = ResourceReadResult::new(vec![
5578 ResourceContentItem::text("a://1", "first"),
5579 ResourceContentItem::blob("b://2", "image/png", "base64data"),
5580 ]);
5581 assert_eq!(result.contents.len(), 2);
5582 assert_eq!(result.first_text(), Some("first"));
5584 assert!(result.first_blob().is_none());
5586 }
5587
5588 #[test]
5589 fn test_resource_read_result_empty() {
5590 let result = ResourceReadResult::new(vec![]);
5591 assert!(result.first_text().is_none());
5592 assert!(result.first_blob().is_none());
5593 }
5594
5595 #[test]
5596 fn test_resource_read_result_blob_first() {
5597 let result = ResourceReadResult::new(vec![ResourceContentItem::blob(
5598 "b://1",
5599 "image/png",
5600 "data",
5601 )]);
5602 assert!(result.first_text().is_none());
5603 assert_eq!(result.first_blob(), Some("data"));
5604 }
5605
5606 #[test]
5611 fn test_tool_content_item_text() {
5612 let item = ToolContentItem::text("hello");
5613 assert_eq!(item.as_text(), Some("hello"));
5614 assert!(item.is_text());
5615 }
5616
5617 #[test]
5618 fn test_tool_content_item_image() {
5619 let item = ToolContentItem::Image {
5620 data: "base64img".to_string(),
5621 mime_type: "image/png".to_string(),
5622 };
5623 assert!(item.as_text().is_none());
5624 assert!(!item.is_text());
5625 }
5626
5627 #[test]
5628 fn test_tool_content_item_audio() {
5629 let item = ToolContentItem::Audio {
5630 data: "base64audio".to_string(),
5631 mime_type: "audio/wav".to_string(),
5632 };
5633 assert!(item.as_text().is_none());
5634 assert!(!item.is_text());
5635 }
5636
5637 #[test]
5638 fn test_tool_content_item_resource() {
5639 let item = ToolContentItem::Resource {
5640 uri: "file://test".to_string(),
5641 mime_type: Some("text/plain".to_string()),
5642 text: Some("embedded".to_string()),
5643 blob: None,
5644 };
5645 assert!(item.as_text().is_none());
5646 assert!(!item.is_text());
5647 }
5648
5649 #[test]
5654 fn test_tool_call_result_success() {
5655 let result = ToolCallResult::success(vec![
5656 ToolContentItem::text("item1"),
5657 ToolContentItem::text("item2"),
5658 ]);
5659 assert!(!result.is_error);
5660 assert_eq!(result.content.len(), 2);
5661 assert_eq!(result.first_text(), Some("item1"));
5662 }
5663
5664 #[test]
5665 fn test_tool_call_result_text() {
5666 let result = ToolCallResult::text("simple output");
5667 assert!(!result.is_error);
5668 assert_eq!(result.content.len(), 1);
5669 assert_eq!(result.first_text(), Some("simple output"));
5670 }
5671
5672 #[test]
5673 fn test_tool_call_result_error() {
5674 let result = ToolCallResult::error("something failed");
5675 assert!(result.is_error);
5676 assert_eq!(result.first_text(), Some("something failed"));
5677 }
5678
5679 #[test]
5680 fn test_tool_call_result_empty() {
5681 let result = ToolCallResult::success(vec![]);
5682 assert!(!result.is_error);
5683 assert!(result.first_text().is_none());
5684 }
5685
5686 #[test]
5691 fn test_elicitation_response_accept() {
5692 let mut data = std::collections::HashMap::new();
5693 data.insert("name".to_string(), serde_json::json!("Alice"));
5694 data.insert("age".to_string(), serde_json::json!(30));
5695 data.insert("active".to_string(), serde_json::json!(true));
5696
5697 let resp = ElicitationResponse::accept(data);
5698 assert!(resp.is_accepted());
5699 assert!(!resp.is_declined());
5700 assert!(!resp.is_cancelled());
5701 assert_eq!(resp.get_string("name"), Some("Alice"));
5702 assert_eq!(resp.get_int("age"), Some(30));
5703 assert_eq!(resp.get_bool("active"), Some(true));
5704 }
5705
5706 #[test]
5707 fn test_elicitation_response_accept_url() {
5708 let resp = ElicitationResponse::accept_url();
5709 assert!(resp.is_accepted());
5710 assert!(resp.content.is_none());
5711 assert!(resp.get_string("anything").is_none());
5712 }
5713
5714 #[test]
5715 fn test_elicitation_response_decline() {
5716 let resp = ElicitationResponse::decline();
5717 assert!(!resp.is_accepted());
5718 assert!(resp.is_declined());
5719 assert!(!resp.is_cancelled());
5720 assert!(resp.get_string("key").is_none());
5721 }
5722
5723 #[test]
5724 fn test_elicitation_response_cancel() {
5725 let resp = ElicitationResponse::cancel();
5726 assert!(!resp.is_accepted());
5727 assert!(!resp.is_declined());
5728 assert!(resp.is_cancelled());
5729 }
5730
5731 #[test]
5732 fn test_elicitation_response_missing_key() {
5733 let mut data = std::collections::HashMap::new();
5734 data.insert("exists".to_string(), serde_json::json!("value"));
5735 let resp = ElicitationResponse::accept(data);
5736
5737 assert!(resp.get_string("missing").is_none());
5738 assert!(resp.get_bool("missing").is_none());
5739 assert!(resp.get_int("missing").is_none());
5740 }
5741
5742 #[test]
5743 fn test_elicitation_response_type_mismatch() {
5744 let mut data = std::collections::HashMap::new();
5745 data.insert("num".to_string(), serde_json::json!(42));
5746 let resp = ElicitationResponse::accept(data);
5747
5748 assert!(resp.get_string("num").is_none());
5750 assert!(resp.get_bool("num").is_none());
5752 assert_eq!(resp.get_int("num"), Some(42));
5754 }
5755
5756 #[test]
5761 fn test_can_sample_false_by_default() {
5762 let cx = Cx::for_testing();
5763 let ctx = McpContext::new(cx, 1);
5764 assert!(!ctx.can_sample());
5765 }
5766
5767 #[test]
5768 fn test_can_elicit_false_by_default() {
5769 let cx = Cx::for_testing();
5770 let ctx = McpContext::new(cx, 1);
5771 assert!(!ctx.can_elicit());
5772 }
5773
5774 #[test]
5775 fn test_can_read_resources_false_by_default() {
5776 let cx = Cx::for_testing();
5777 let ctx = McpContext::new(cx, 1);
5778 assert!(!ctx.can_read_resources());
5779 }
5780
5781 #[test]
5782 fn test_can_call_tools_false_by_default() {
5783 let cx = Cx::for_testing();
5784 let ctx = McpContext::new(cx, 1);
5785 assert!(!ctx.can_call_tools());
5786 }
5787
5788 #[test]
5789 fn test_resource_read_depth_default() {
5790 let cx = Cx::for_testing();
5791 let ctx = McpContext::new(cx, 1);
5792 assert_eq!(ctx.resource_read_depth(), 0);
5793 }
5794
5795 #[test]
5796 fn test_tool_call_depth_default() {
5797 let cx = Cx::for_testing();
5798 let ctx = McpContext::new(cx, 1);
5799 assert_eq!(ctx.tool_call_depth(), 0);
5800 }
5801
5802 #[test]
5807 fn sampling_request_builder_chain() {
5808 let req = SamplingRequest::prompt("hello", 100)
5809 .with_system_prompt("You are helpful")
5810 .with_temperature(0.7)
5811 .with_stop_sequences(vec!["STOP".into()])
5812 .with_model_hints(vec!["gpt-4".into()]);
5813
5814 assert_eq!(req.messages.len(), 1);
5815 assert_eq!(req.max_tokens, 100);
5816 assert_eq!(req.system_prompt.as_deref(), Some("You are helpful"));
5817 assert_eq!(req.temperature, Some(0.7));
5818 assert_eq!(req.stop_sequences, vec!["STOP"]);
5819 assert_eq!(req.model_hints, vec!["gpt-4"]);
5820 }
5821
5822 #[test]
5823 fn sampling_request_message_roles() {
5824 let user = SamplingRequestMessage::user("hi");
5825 assert_eq!(user.role, SamplingRole::User);
5826 assert_eq!(user.text, "hi");
5827
5828 let asst = SamplingRequestMessage::assistant("hello");
5829 assert_eq!(asst.role, SamplingRole::Assistant);
5830 assert_eq!(asst.text, "hello");
5831 }
5832
5833 #[test]
5834 fn sampling_response_new_default_stop_reason() {
5835 let resp = SamplingResponse::new("output", "model-1");
5836 assert_eq!(resp.text, "output");
5837 assert_eq!(resp.model, "model-1");
5838 assert_eq!(resp.stop_reason, SamplingStopReason::EndTurn);
5839 assert_eq!(SamplingStopReason::default(), SamplingStopReason::EndTurn);
5840 }
5841
5842 #[test]
5843 fn sampling_stop_reason_round_trips_optional_open_wire_values() {
5844 let absent = SamplingStopReason::from_wire_value(None);
5845 assert_eq!(absent, SamplingStopReason::Unspecified);
5846 assert_eq!(absent.as_wire_value(), None);
5847
5848 let provider =
5849 SamplingStopReason::from_wire_value(Some("provider_safety_limit".to_owned()));
5850 assert_eq!(
5851 provider,
5852 SamplingStopReason::Other("provider_safety_limit".to_owned())
5853 );
5854 assert_eq!(provider.as_wire_value(), Some("provider_safety_limit"));
5855 }
5856
5857 #[test]
5858 fn noop_sampling_sender_returns_error() {
5859 let sender = NoOpSamplingSender;
5860 let req = SamplingRequest::prompt("test", 10);
5861 let result = crate::block_on(sender.create_message(req));
5862 assert!(result.is_err());
5863 }
5864
5865 #[test]
5866 fn noop_elicitation_sender_returns_error() {
5867 let sender = NoOpElicitationSender;
5868 let req = ElicitationRequest::form("msg", serde_json::json!({}));
5869 let result = crate::block_on(sender.elicit(req));
5870 assert!(result.is_err());
5871 }
5872
5873 #[test]
5874 fn elicitation_request_form_constructor() {
5875 let req = ElicitationRequest::form("Enter name", serde_json::json!({"type": "string"}));
5876 assert_eq!(req.mode, ElicitationMode::Form);
5877 assert_eq!(req.message, "Enter name");
5878 assert!(req.schema.is_some());
5879 assert!(req.url.is_none());
5880 assert!(req.elicitation_id.is_none());
5881 }
5882
5883 #[test]
5884 fn elicitation_request_url_constructor() {
5885 let req = ElicitationRequest::url("Login", "https://example.com", "id-1");
5886 assert_eq!(req.mode, ElicitationMode::Url);
5887 assert_eq!(req.message, "Login");
5888 assert_eq!(req.url.as_deref(), Some("https://example.com"));
5889 assert_eq!(req.elicitation_id.as_deref(), Some("id-1"));
5890 assert!(req.schema.is_none());
5891 }
5892
5893 #[test]
5894 fn mcp_context_with_sampling_enables_can_sample() {
5895 let cx = Cx::for_testing();
5896 let sender = Arc::new(NoOpSamplingSender);
5897 let ctx = McpContext::new(cx, 1).with_sampling(sender);
5898 assert!(ctx.can_sample());
5899 }
5900
5901 #[test]
5902 fn mcp_context_with_elicitation_enables_can_elicit() {
5903 let cx = Cx::for_testing();
5904 let sender = Arc::new(NoOpElicitationSender);
5905 let ctx = McpContext::new(cx, 1).with_elicitation(sender);
5906 assert!(ctx.can_elicit());
5907 }
5908
5909 struct FixedRootsProvider;
5910
5911 impl RootsProvider for FixedRootsProvider {
5912 fn list_roots(
5913 &self,
5914 ) -> std::pin::Pin<
5915 Box<dyn std::future::Future<Output = crate::McpResult<Vec<ClientRoot>>> + Send + '_>,
5916 > {
5917 Box::pin(async {
5918 Ok(vec![
5919 ClientRoot::with_name("file:///workspace", "workspace"),
5920 ClientRoot::new("file:///tmp"),
5921 ])
5922 })
5923 }
5924 }
5925
5926 #[test]
5927 fn mcp_context_roots_provider_returns_client_roots() {
5928 let ctx =
5929 McpContext::new(Cx::for_testing(), 1).with_roots_provider(Arc::new(FixedRootsProvider));
5930
5931 assert!(ctx.can_list_roots());
5932 let roots = crate::block_on(ctx.list_roots()).expect("configured roots provider succeeds");
5933 assert_eq!(
5934 roots,
5935 vec![
5936 ClientRoot::with_name("file:///workspace", "workspace"),
5937 ClientRoot::new("file:///tmp"),
5938 ]
5939 );
5940 }
5941
5942 #[test]
5943 fn mcp_context_without_roots_provider_rejects_without_authority() {
5944 let ctx = McpContext::new(Cx::for_testing(), 1);
5945
5946 assert!(!ctx.can_list_roots());
5947 let error = crate::block_on(ctx.list_roots())
5948 .expect_err("without only the roots provider, the context must reject the request");
5949 assert_eq!(error.code, crate::McpErrorCode::InvalidRequest);
5950 assert_eq!(
5951 error.message,
5952 "Roots not available: client does not support roots capability"
5953 );
5954 }
5955
5956 #[test]
5957 fn mcp_context_depth_setters() {
5958 let cx = Cx::for_testing();
5959 let ctx = McpContext::new(cx, 1)
5960 .with_resource_read_depth(3)
5961 .with_tool_call_depth(5);
5962 assert_eq!(ctx.resource_read_depth(), 3);
5963 assert_eq!(ctx.tool_call_depth(), 5);
5964
5965 let attempted_reset = ctx.with_resource_read_depth(0).with_tool_call_depth(0);
5966 assert_eq!(attempted_reset.resource_read_depth(), 3);
5967 assert_eq!(attempted_reset.tool_call_depth(), 5);
5968 }
5969
5970 #[test]
5971 fn mcp_context_debug_includes_request_id() {
5972 let cx = Cx::for_testing();
5973 let ctx = McpContext::new(cx, 99);
5974 let debug = format!("{ctx:?}");
5975 assert!(debug.contains("request_id: 99"));
5976 }
5977
5978 #[test]
5979 fn mcp_context_cx_and_trace() {
5980 let cx = Cx::for_testing();
5981 let ctx = McpContext::new(cx, 1);
5982 let _ = ctx.cx();
5984 ctx.trace("test event");
5986 }
5987
5988 #[test]
5989 fn final_result_outcome_preserves_dual_era_and_terminal_reason() {
5990 use crate::combinator::{DualEraFinalResult, FinalRequestResult};
5991
5992 let context = McpContext::new(Cx::for_testing(), 1);
5993 let modern = context.final_result_outcome(
5994 FinalRequestResult::<u64, String, &'static str>::modern("typed-final", 42),
5995 );
5996 let legacy =
5997 context.final_result_outcome(FinalRequestResult::<u64, String, &'static str>::legacy(
5998 "legacy-final",
5999 "legacy wire result".to_owned(),
6000 ));
6001
6002 let Outcome::Ok(modern) = modern else {
6003 panic!("live context admits the modern final result");
6004 };
6005 assert_eq!(modern.terminal_reason(), &"typed-final");
6006 assert_eq!(modern.result(), &DualEraFinalResult::Modern(42));
6007
6008 let Outcome::Ok(legacy) = legacy else {
6009 panic!("live context admits the legacy final result");
6010 };
6011 assert_eq!(legacy.terminal_reason(), &"legacy-final");
6012 assert_eq!(
6013 legacy.result(),
6014 &DualEraFinalResult::Legacy("legacy wire result".to_owned())
6015 );
6016 }
6017
6018 #[test]
6019 fn final_result_outcome_cancellation_negative_preserves_cx_reason() {
6020 use crate::combinator::FinalRequestResult;
6021 use asupersync::types::CancelKind;
6022
6023 let cx = Cx::for_testing();
6024 cx.cancel_with(CancelKind::Timeout, Some("final-result race"));
6025 let expected_reason = cx
6026 .cancel_reason()
6027 .expect("cancel_with records the caller-owned terminal reason");
6028 let context = McpContext::new(cx, 1);
6029
6030 let outcome = context.final_result_outcome(
6031 FinalRequestResult::<u64, String, &'static str>::modern("typed-final", 42),
6032 );
6033
6034 let Outcome::Cancelled(reason) = outcome else {
6035 panic!("changing only caller cancellation rejects the same final result");
6036 };
6037 assert_eq!(reason, expected_reason);
6038 }
6039
6040 #[test]
6041 fn final_result_outcome_panic_negative_preserves_payload() {
6042 use crate::combinator::FinalRequestResult;
6043 use asupersync::types::{CancelKind, PanicPayload};
6044
6045 type Final = FinalRequestResult<u64, String, &'static str>;
6046
6047 let cx = Cx::for_testing();
6048 cx.cancel_with(CancelKind::Timeout, Some("competing terminal state"));
6049 let context = McpContext::new(cx, 1);
6050 let payload = PanicPayload::new("final typed result panicked");
6051 let source: crate::McpOutcome<Final> = Outcome::Panicked(payload.clone());
6052
6053 let outcome = context.adapt_final_request_outcome(source);
6054
6055 let Outcome::Panicked(actual) = outcome else {
6056 panic!("changing only the source terminal state to panic preserves panic");
6057 };
6058 assert_eq!(actual, payload);
6059 }
6060}