Skip to main content

fastmcp_server/
bidirectional.rs

1//! Bidirectional request handling for server-to-client communication.
2//!
3//! This module provides the infrastructure for server-initiated requests to clients,
4//! such as:
5//! - `sampling/createMessage` - Request LLM completion from the client
6//! - `elicitation/create` - Request user input from the client
7//! - `roots/list` - Request filesystem roots from the client
8//!
9//! # Architecture
10//!
11//! The MCP protocol is bidirectional: while clients typically send requests to servers,
12//! servers can also send requests to clients. This creates a challenge because the
13//! server's main loop is typically blocking on `recv()`.
14//!
15//! The solution is a message dispatcher pattern:
16//! 1. A background task continuously reads from the transport
17//! 2. Incoming messages are routed based on whether they're requests or responses
18//! 3. Responses are matched to pending requests via their ID
19//! 4. Requests are dispatched to handlers
20//!
21//! # Usage
22//!
23//! ```ignore
24//! // Send a request and await the response
25//! let response = request_sender.send_request(
26//!     &cx,
27//!     "sampling/createMessage",
28//!     params,
29//! ).await?;
30//! ```
31
32use std::collections::BTreeMap;
33use std::collections::HashMap;
34use std::collections::hash_map::Entry;
35use std::future::{Future, poll_fn};
36use std::sync::{Arc, Mutex};
37use std::task::Poll;
38use std::time::{Duration, Instant};
39
40use asupersync::Cx;
41use asupersync::channel::oneshot;
42use asupersync::channel::oneshot::RecvError;
43use base64::Engine as _;
44use fastmcp_core::{
45    ClientRoot, ElicitationAction, ElicitationMode, ElicitationRequest, ElicitationResponse,
46    ElicitationSender, McpContext, McpError, McpErrorCode, McpRequestCancellation, McpResult,
47    RootsProvider, SamplingRequest, SamplingResponse, SamplingRole, SamplingSender,
48    SamplingStopReason, draw_security_identifier,
49};
50use fastmcp_protocol::protocol_policy::ProtocolEra;
51use fastmcp_protocol::{
52    CorrelationKey, FinalInputResponses, JsonRpcError, JsonRpcMessage, JsonRpcRequest,
53    JsonRpcResponse, RequestId,
54};
55use serde::Serialize;
56use serde::de::DeserializeOwned;
57use serde::ser::{SerializeMap, SerializeStruct};
58
59/// Default maximum number of concurrent server-to-client requests.
60pub const DEFAULT_MAX_IN_FLIGHT_REQUESTS: usize = 1_024;
61
62/// Absolute maximum accepted by [`PendingRequests::with_max_in_flight`].
63pub const HARD_MAX_IN_FLIGHT_REQUESTS: usize = 16_384;
64
65/// Default maximum rounds for a single final MRTR exchange.
66pub const DEFAULT_MAX_MRTR_ROUNDS: u8 = 8;
67
68/// Absolute maximum rounds a server-local MRTR exchange may use.
69pub const HARD_MAX_MRTR_ROUNDS: u8 = 32;
70
71/// Default maximum embedded input requests in one MRTR result.
72pub const DEFAULT_MAX_MRTR_INPUT_REQUESTS_PER_ROUND: usize = 32;
73
74/// Absolute maximum embedded input requests in one MRTR result.
75pub const HARD_MAX_MRTR_INPUT_REQUESTS_PER_ROUND: usize = 128;
76
77/// Default maximum embedded input requests across one complete MRTR exchange.
78pub const DEFAULT_MAX_MRTR_INPUT_REQUESTS_TOTAL: usize = 128;
79
80/// Absolute maximum embedded input requests across one complete MRTR exchange.
81pub const HARD_MAX_MRTR_INPUT_REQUESTS_TOTAL: usize = 512;
82
83/// Default lifetime for an MRTR request-state record.
84pub const DEFAULT_MRTR_REQUEST_STATE_TTL: Duration = Duration::from_mins(15);
85
86/// Absolute maximum lifetime for an MRTR request-state record.
87pub const HARD_MAX_MRTR_REQUEST_STATE_TTL: Duration = Duration::from_hours(1);
88
89/// Default number of retained, process-local MRTR request-state records.
90pub const DEFAULT_MAX_MRTR_REQUEST_STATES: usize = 4_096;
91
92/// Absolute maximum number of retained, process-local MRTR request-state records.
93pub const HARD_MAX_MRTR_REQUEST_STATES: usize = 65_536;
94
95/// Default maximum encoded request-state bytes admitted from an MRTR retry.
96pub const DEFAULT_MAX_MRTR_REQUEST_STATE_BYTES: usize = 64 * 1024;
97
98/// Maximum encoded request-state bytes admitted from an MRTR retry.
99pub const HARD_MAX_MRTR_REQUEST_STATE_BYTES: usize = 256 * 1024;
100
101const FIRST_SERVER_REQUEST_ID: i64 = 1_000_000;
102/// The first exact-legacy ID is exactly representable by a JavaScript `Number`.
103const FIRST_EXACT_LEGACY_SERVER_REQUEST_ID: i64 = -1;
104/// The inclusive lower bound of JavaScript's integer-safe `Number` range.
105const LAST_EXACT_LEGACY_SERVER_REQUEST_ID: i64 = -9_007_199_254_740_991;
106const INVALID_LIMIT_ERROR: &str = "Invalid bidirectional request limit";
107const IN_FLIGHT_LIMIT_ERROR: &str = "Bidirectional request limit reached";
108const REQUEST_ID_EXHAUSTED_ERROR: &str = "Bidirectional request IDs exhausted";
109const INVALID_RESPONSE_ERROR: &str = "Invalid JSON-RPC response";
110const REMOTE_RESPONSE_ERROR: &str = "Client returned an error response";
111const CONNECTION_CLOSED_ERROR: &str = "Bidirectional connection closed";
112const TRANSPORT_SEND_ERROR: &str = "Failed to send bidirectional request";
113const RESPONSE_CHANNEL_ERROR: &str = "Bidirectional response channel closed";
114const RESPONSE_PAYLOAD_ERROR: &str = "Invalid bidirectional response payload";
115const REQUEST_PAYLOAD_ERROR: &str = "Failed to serialize bidirectional request payload";
116const INVALID_ELICITATION_REQUEST_ERROR: &str = "Invalid elicitation request";
117const INVALID_MRTR_LIMIT_ERROR: &str = "Invalid MRTR exchange limit";
118const MRTR_REQUEST_STATE_ERROR: &str = "Invalid or expired MRTR request state";
119const MRTR_REQUEST_STATE_UNAVAILABLE_ERROR: &str = "Unable to create MRTR request state";
120const MRTR_INPUT_MAP_ERROR: &str = "Invalid MRTR input request or response map";
121const MRTR_RESPONSE_KIND_ERROR: &str = "MRTR input response does not match its request";
122const MRTR_ROUND_LIMIT_ERROR: &str = "MRTR exchange limit reached";
123
124/// The immutable request facts a router binds to one opaque MRTR state.
125///
126/// This is deliberately server-local: it is never serialized and prevents a
127/// state minted for one modern operation from resuming another operation that
128/// happens to request the same embedded input kinds.
129#[derive(Debug, Clone, PartialEq, Eq)]
130pub(crate) struct MrtrExchangeBinding {
131    method: &'static str,
132    target: String,
133    arguments_digest: [u8; 32],
134    session_partition: [u8; 32],
135    principal_digest: Option<[u8; 32]>,
136}
137
138impl MrtrExchangeBinding {
139    /// Captures the router-admitted operation identity for a future retry.
140    #[must_use]
141    pub(crate) fn new(
142        method: &'static str,
143        target: String,
144        arguments_digest: [u8; 32],
145        session_partition: [u8; 32],
146        principal_digest: Option<[u8; 32]>,
147    ) -> Self {
148        Self {
149            method,
150            target,
151            arguments_digest,
152            session_partition,
153            principal_digest,
154        }
155    }
156}
157const LEGACY_INPUT_RETRY_ERROR: &str = "MCP 2024-11-05 does not support input retries";
158
159// ============================================================================
160// Pending Request Tracking
161// ============================================================================
162
163/// A bounded, single-use channel for receiving a response.
164type PendingResponse = McpResult<serde_json::Value>;
165type ResponseSender = oneshot::Sender<PendingResponse>;
166type ResponseReceiver = oneshot::Receiver<PendingResponse>;
167
168/// Immutable wire-ID domain assigned to one pending-request tracker.
169///
170/// Exact legacy reverse requests descend from
171/// [`FIRST_EXACT_LEGACY_SERVER_REQUEST_ID`] through JavaScript's negative safe
172/// integer range. A response from the already issued suffix of that range can
173/// therefore be retired without retaining one tombstone per completed request.
174#[derive(Debug, Clone, Copy, PartialEq, Eq)]
175enum PendingIdDomain {
176    Positive,
177    ExactLegacyNegative,
178}
179
180impl PendingIdDomain {
181    const fn first_id(self) -> i64 {
182        match self {
183            Self::Positive => FIRST_SERVER_REQUEST_ID,
184            Self::ExactLegacyNegative => FIRST_EXACT_LEGACY_SERVER_REQUEST_ID,
185        }
186    }
187
188    fn next_id_after(self, candidate: i64) -> Option<i64> {
189        match self {
190            Self::Positive => candidate.checked_add(1),
191            Self::ExactLegacyNegative => candidate
192                .checked_sub(1)
193                .filter(|next| *next >= LAST_EXACT_LEGACY_SERVER_REQUEST_ID),
194        }
195    }
196
197    fn is_issued_negative_suffix(self, next_id: Option<i64>, id: &CorrelationKey) -> bool {
198        let Self::ExactLegacyNegative = self else {
199            return false;
200        };
201        let CorrelationKey::Integer(integer) = id else {
202            return false;
203        };
204        let Ok(id) = integer.parse::<i64>() else {
205            return false;
206        };
207
208        match next_id {
209            // `next_id` itself has not yet been issued, so the issued suffix
210            // is open at its lower end: `(next_id..=-1)`.
211            Some(next_id) => {
212                (LAST_EXACT_LEGACY_SERVER_REQUEST_ID..=FIRST_EXACT_LEGACY_SERVER_REQUEST_ID)
213                    .contains(&next_id)
214                    && (next_id < id)
215                    && (id <= FIRST_EXACT_LEGACY_SERVER_REQUEST_ID)
216            }
217            None => (LAST_EXACT_LEGACY_SERVER_REQUEST_ID..=FIRST_EXACT_LEGACY_SERVER_REQUEST_ID)
218                .contains(&id),
219        }
220    }
221}
222
223/// Result of routing one response through [`PendingRequests`].
224#[derive(Debug, Clone, Copy, PartialEq, Eq)]
225pub(crate) enum PendingResponseDisposition {
226    /// The response reached its live pending request.
227    Delivered,
228    /// The response belongs to an issued exact-legacy negative ID that has
229    /// already left the pending set.
230    RetiredGeneric,
231    /// The response ID was not issued by this tracker or is absent.
232    Unmatched,
233}
234
235#[derive(Debug)]
236struct PendingState {
237    requests: HashMap<CorrelationKey, PendingRequest>,
238    next_id: Option<i64>,
239    closed: bool,
240}
241
242#[derive(Debug)]
243struct PendingRequest {
244    sender: ResponseSender,
245    request_cancellation: Option<McpRequestCancellation>,
246}
247
248/// Tracks pending server-to-client requests.
249///
250/// When the server sends a request to the client, it registers a response sender
251/// here. When a response arrives, the dispatcher routes it to the correct sender.
252#[derive(Debug)]
253pub struct PendingRequests {
254    state: Mutex<PendingState>,
255    id_domain: PendingIdDomain,
256    max_in_flight: usize,
257}
258
259impl PendingRequests {
260    pub(crate) fn validate_max_in_flight(max_in_flight: usize) -> McpResult<()> {
261        if !(1..=HARD_MAX_IN_FLIGHT_REQUESTS).contains(&max_in_flight) {
262            return Err(McpError::new(
263                McpErrorCode::InvalidParams,
264                INVALID_LIMIT_ERROR,
265            ));
266        }
267        Ok(())
268    }
269
270    fn lock_state(&self) -> std::sync::MutexGuard<'_, PendingState> {
271        match self.state.lock() {
272            Ok(guard) => guard,
273            // Prefer availability over panic if another task panicked while holding the lock.
274            Err(poisoned) => poisoned.into_inner(),
275        }
276    }
277
278    fn new_in_domain(id_domain: PendingIdDomain, max_in_flight: usize) -> Self {
279        Self {
280            state: Mutex::new(PendingState {
281                requests: HashMap::new(),
282                next_id: Some(id_domain.first_id()),
283                closed: false,
284            }),
285            id_domain,
286            max_in_flight,
287        }
288    }
289
290    /// Creates a new pending request tracker.
291    #[must_use]
292    pub fn new() -> Self {
293        Self::new_in_domain(PendingIdDomain::Positive, DEFAULT_MAX_IN_FLIGHT_REQUESTS)
294    }
295
296    /// Creates a tracker with a caller-selected finite in-flight limit.
297    ///
298    /// # Errors
299    ///
300    /// Returns `InvalidParams` when `max_in_flight` is zero or exceeds
301    /// [`HARD_MAX_IN_FLIGHT_REQUESTS`].
302    pub fn with_max_in_flight(max_in_flight: usize) -> McpResult<Self> {
303        Self::validate_max_in_flight(max_in_flight)?;
304
305        Ok(Self::new_in_domain(
306            PendingIdDomain::Positive,
307            max_in_flight,
308        ))
309    }
310
311    /// Creates the exact-legacy tracker whose request IDs descend from
312    /// [`FIRST_EXACT_LEGACY_SERVER_REQUEST_ID`] through the JavaScript-safe
313    /// negative domain.
314    ///
315    /// The domain is fixed for the tracker's full lifetime so a response for
316    /// an issued-but-retired negative ID can be classified in O(1) space.
317    pub(crate) fn with_max_in_flight_for_exact_legacy(max_in_flight: usize) -> McpResult<Self> {
318        Self::validate_max_in_flight(max_in_flight)?;
319
320        Ok(Self::new_in_domain(
321            PendingIdDomain::ExactLegacyNegative,
322            max_in_flight,
323        ))
324    }
325
326    /// Returns the configured maximum number of in-flight requests.
327    #[must_use]
328    pub const fn max_in_flight(&self) -> usize {
329        self.max_in_flight
330    }
331
332    /// Returns the current number of in-flight requests.
333    #[must_use]
334    pub fn in_flight_len(&self) -> usize {
335        self.lock_state().requests.len()
336    }
337
338    /// Atomically allocates a collision-free ID and registers its response
339    /// channel. Allocation scans at most `max_in_flight + 1` candidates.
340    fn register(&self) -> McpResult<(RequestId, ResponseReceiver)> {
341        self.register_with_cancellation(None)
342    }
343
344    fn register_with_cancellation(
345        &self,
346        request_cancellation: Option<McpRequestCancellation>,
347    ) -> McpResult<(RequestId, ResponseReceiver)> {
348        let mut state = self.lock_state();
349        if state.closed {
350            return Err(McpError::internal_error(CONNECTION_CLOSED_ERROR));
351        }
352        if state.requests.len() >= self.max_in_flight {
353            return Err(McpError::internal_error(IN_FLIGHT_LIMIT_ERROR));
354        }
355
356        for _ in 0..=self.max_in_flight {
357            let Some(candidate) = state.next_id else {
358                return Err(McpError::internal_error(REQUEST_ID_EXHAUSTED_ERROR));
359            };
360            let id = RequestId::Number(candidate);
361            let key = id
362                .correlation_key()
363                .map_err(|_| McpError::internal_error(REQUEST_ID_EXHAUSTED_ERROR))?;
364            state.next_id = self.id_domain.next_id_after(candidate);
365
366            if let Entry::Vacant(entry) = state.requests.entry(key) {
367                let (sender, receiver) = oneshot::channel();
368                entry.insert(PendingRequest {
369                    sender,
370                    request_cancellation,
371                });
372                return Ok((id, receiver));
373            }
374        }
375
376        Err(McpError::internal_error(REQUEST_ID_EXHAUSTED_ERROR))
377    }
378
379    /// Routes a response to the appropriate pending request.
380    ///
381    /// Returns `true` only when the response was delivered to a live pending
382    /// request, preserving the established public boolean contract.
383    pub fn route_response(&self, response: &JsonRpcResponse) -> bool {
384        matches!(
385            self.route_response_with_disposition(response),
386            PendingResponseDisposition::Delivered
387        )
388    }
389
390    /// Routes a response and reports whether it was delivered, retired from
391    /// the exact-legacy negative ID suffix, or unmatched.
392    pub(crate) fn route_response_with_disposition(
393        &self,
394        response: &JsonRpcResponse,
395    ) -> PendingResponseDisposition {
396        let Some(ref id) = response.id else {
397            return PendingResponseDisposition::Unmatched;
398        };
399        let Ok(key) = id.correlation_key() else {
400            return PendingResponseDisposition::Unmatched;
401        };
402
403        let (pending, retired_generic) = {
404            let mut state = self.lock_state();
405            let pending = state.requests.remove(&key);
406            let retired_generic = pending.is_none()
407                && self
408                    .id_domain
409                    .is_issued_negative_suffix(state.next_id, &key);
410            (pending, retired_generic)
411        };
412
413        if let Some(pending) = pending {
414            // Validate every response invariant before consuming the waiter.
415            // This also rejects manually-constructed values that bypass serde's guards.
416            let validated = ValidatedResponse::from_response(response);
417            let outcome = validated.into_pending_response();
418            // The response path is synchronous, so use the immediate bounded
419            // oneshot bridge. Receiver dropout returns the value and is safe to
420            // ignore after the map entry has been removed.
421            let _ = pending.sender.send_blocking(outcome);
422            PendingResponseDisposition::Delivered
423        } else if retired_generic {
424            PendingResponseDisposition::RetiredGeneric
425        } else {
426            PendingResponseDisposition::Unmatched
427        }
428    }
429
430    /// Removes a pending request (e.g., on timeout or cancellation).
431    pub fn remove(&self, id: &RequestId) {
432        let Ok(key) = id.correlation_key() else {
433            return;
434        };
435        let mut state = self.lock_state();
436        state.requests.remove(&key);
437    }
438
439    /// Wakes pending server-to-client calls whose owning incoming request is
440    /// terminal, without mutating the caller-owned connection context.
441    pub(crate) fn cancel_cancelled(&self) -> usize {
442        let cancelled = {
443            let mut state = self.lock_state();
444            let ids: Vec<CorrelationKey> = state
445                .requests
446                .iter()
447                .filter(|(_, pending)| {
448                    pending
449                        .request_cancellation
450                        .as_ref()
451                        .is_some_and(McpRequestCancellation::is_terminal)
452                })
453                .map(|(id, _)| id.clone())
454                .collect();
455            ids.into_iter()
456                .filter_map(|id| state.requests.remove(&id))
457                .collect::<Vec<_>>()
458        };
459        let count = cancelled.len();
460        for pending in cancelled {
461            let _ = pending
462                .sender
463                .send_blocking(Err(McpError::request_cancelled()));
464        }
465        count
466    }
467
468    /// Permanently closes the tracker and cancels every pending request.
469    ///
470    /// Closing is irreversible: later registration attempts fail with the
471    /// same fixed connection-closed error. This prevents a request racing with
472    /// connection teardown from installing an orphaned waiter after the drain.
473    pub fn cancel_all(&self) {
474        let senders: Vec<PendingRequest> = {
475            let mut state = self.lock_state();
476            state.closed = true;
477            state.requests.drain().map(|(_, pending)| pending).collect()
478        };
479        for pending in senders {
480            let _ = pending
481                .sender
482                .send_blocking(Err(McpError::internal_error(CONNECTION_CLOSED_ERROR)));
483        }
484    }
485
486    #[cfg(test)]
487    fn set_next_id_for_test(&self, next_id: i64) {
488        self.lock_state().next_id = Some(next_id);
489    }
490}
491
492enum ValidatedResponse<'a> {
493    Success(&'a serde_json::Value),
494    Error(&'a JsonRpcError),
495    Invalid,
496}
497
498impl<'a> ValidatedResponse<'a> {
499    fn from_response(response: &'a JsonRpcResponse) -> Self {
500        if response.validate().is_err() {
501            return Self::Invalid;
502        }
503
504        match (&response.result, &response.error) {
505            (Some(result), None) => Self::Success(result),
506            (None, Some(error)) => Self::Error(error),
507            (Some(_), Some(_)) | (None, None) => Self::Invalid,
508        }
509    }
510
511    fn into_pending_response(self) -> PendingResponse {
512        match self {
513            Self::Success(result) => Ok(result.clone()),
514            Self::Error(error) => Err(McpError::new(
515                error
516                    .code
517                    .as_i32()
518                    .map(McpErrorCode::from)
519                    .unwrap_or(McpErrorCode::InternalError),
520                REMOTE_RESPONSE_ERROR,
521            )),
522            Self::Invalid => Err(McpError::internal_error(INVALID_RESPONSE_ERROR)),
523        }
524    }
525}
526
527impl Default for PendingRequests {
528    fn default() -> Self {
529        Self::new()
530    }
531}
532
533/// Owns the local and peer-facing cleanup for one reverse request.
534///
535/// Once the outbound request has been committed to the transport, dropping its
536/// future without a routed response must both free the local slot and tell the
537/// exact-2024 peer to stop work. MCP 2026-07-28 does not permit this reverse
538/// cancellation control. The legacy notification remains best effort because
539/// a closing transport cannot reliably deliver another frame.
540struct PendingRequestGuard {
541    pending: Arc<PendingRequests>,
542    send_fn: TransportSendFn,
543    era: ProtocolEra,
544    id: RequestId,
545    request_sent: bool,
546    finished: bool,
547}
548
549impl PendingRequestGuard {
550    fn mark_request_sent(&mut self) {
551        self.request_sent = true;
552    }
553
554    fn finish(&mut self) {
555        self.finished = true;
556        self.pending.remove(&self.id);
557    }
558
559    fn cancel(&mut self) {
560        self.pending.remove(&self.id);
561        self.send_cancellation_notification();
562        self.finished = true;
563    }
564
565    fn send_cancellation_notification(&self) {
566        if !self.request_sent || self.era != ProtocolEra::Legacy2024 {
567            return;
568        }
569
570        let message = JsonRpcMessage::Request(JsonRpcRequest::notification(
571            "notifications/cancelled",
572            Some(serde_json::json!({ "requestId": self.id.clone() })),
573        ));
574        // A reverse request is already terminal locally. Do not replace that
575        // outcome with a best-effort control-frame transport failure.
576        let _ = (self.send_fn)(&message);
577    }
578}
579
580impl Drop for PendingRequestGuard {
581    fn drop(&mut self) {
582        self.pending.remove(&self.id);
583        if !self.finished {
584            self.send_cancellation_notification();
585        }
586    }
587}
588
589// ============================================================================
590// Transport Request Sender
591// ============================================================================
592
593/// Callback type for sending messages through the transport.
594pub type TransportSendFn = Arc<dyn Fn(&JsonRpcMessage) -> Result<(), String> + Send + Sync>;
595
596/// Sends server-to-client requests through the transport.
597///
598/// This struct provides a way to send requests to the client and await responses.
599/// It works in conjunction with [`PendingRequests`] to track in-flight requests.
600#[derive(Clone)]
601pub struct RequestSender {
602    /// Pending request tracker.
603    pending: Arc<PendingRequests>,
604    /// Transport send callback.
605    send_fn: TransportSendFn,
606    /// Exact protocol era that governs reverse-request cleanup controls.
607    era: ProtocolEra,
608    /// Request-local cancellation domain installed by server dispatch.
609    request_cancellation: Option<McpRequestCancellation>,
610}
611
612impl RequestSender {
613    /// Creates an exact MCP 2024-11-05 request sender.
614    ///
615    /// Use [`Self::new_for_era`] when the negotiated era is available.
616    pub fn new(pending: Arc<PendingRequests>, send_fn: TransportSendFn) -> Self {
617        Self::new_for_era(ProtocolEra::Legacy2024, pending, send_fn)
618    }
619
620    /// Creates a request sender bound to one negotiated protocol era.
621    ///
622    /// Dropped reverse requests emit `notifications/cancelled` only for exact
623    /// MCP 2024-11-05. MCP 2026-07-28 retains local cleanup but emits no
624    /// server cancellation notification.
625    pub fn new_for_era(
626        era: ProtocolEra,
627        pending: Arc<PendingRequests>,
628        send_fn: TransportSendFn,
629    ) -> Self {
630        Self {
631            pending,
632            send_fn,
633            era,
634            request_cancellation: None,
635        }
636    }
637
638    pub(crate) fn for_request(&self, request_cancellation: McpRequestCancellation) -> Self {
639        Self {
640            pending: Arc::clone(&self.pending),
641            send_fn: Arc::clone(&self.send_fn),
642            era: self.era,
643            request_cancellation: Some(request_cancellation),
644        }
645    }
646
647    fn request_is_terminal(&self) -> bool {
648        self.request_cancellation
649            .as_ref()
650            .is_some_and(McpRequestCancellation::is_terminal)
651    }
652
653    /// Sends a request to the client and waits for a response.
654    ///
655    /// # Errors
656    ///
657    /// Returns an error if:
658    /// - The finite in-flight request limit is reached
659    /// - The transport send fails
660    /// - The request times out (based on budget)
661    /// - The client returns an error response
662    /// - The response envelope or typed payload is invalid
663    /// - The connection is closed
664    pub async fn send_request<T: serde::de::DeserializeOwned>(
665        &self,
666        cx: &Cx,
667        method: &str,
668        params: serde_json::Value,
669    ) -> McpResult<T> {
670        if cx.checkpoint().is_err() || self.request_is_terminal() {
671            return Err(McpError::request_cancelled());
672        }
673
674        let (id, mut receiver) = self
675            .pending
676            .register_with_cancellation(self.request_cancellation.clone())?;
677        let mut guard = PendingRequestGuard {
678            pending: Arc::clone(&self.pending),
679            send_fn: Arc::clone(&self.send_fn),
680            era: self.era,
681            id: id.clone(),
682            request_sent: false,
683            finished: false,
684        };
685        if cx.checkpoint().is_err() || self.request_is_terminal() {
686            return Err(McpError::request_cancelled());
687        }
688
689        let request = JsonRpcRequest::new(method.to_string(), Some(params), id.clone());
690        let message = JsonRpcMessage::Request(request);
691
692        // Send the request through the transport
693        if (self.send_fn)(&message).is_err() {
694            return Err(McpError::internal_error(TRANSPORT_SEND_ERROR));
695        }
696        guard.mark_request_sent();
697
698        let response = if let Some(request_cancellation) = &self.request_cancellation {
699            let mut receive = std::pin::pin!(receiver.recv(cx));
700            let mut terminated = std::pin::pin!(request_cancellation.terminated());
701
702            poll_fn(|task_cx| {
703                // Request termination owns ties: check before polling either
704                // source, after arming its waiter, and once more after polling
705                // the response future.
706                if request_cancellation.is_terminal() {
707                    return Poll::Ready(Err(McpError::request_cancelled()));
708                }
709                if terminated.as_mut().poll(task_cx).is_ready() {
710                    return Poll::Ready(Err(McpError::request_cancelled()));
711                }
712
713                let receive_poll = receive.as_mut().poll(task_cx);
714                if request_cancellation.is_terminal() {
715                    return Poll::Ready(Err(McpError::request_cancelled()));
716                }
717                match receive_poll {
718                    Poll::Ready(Ok(response)) => Poll::Ready(response),
719                    Poll::Ready(Err(RecvError::Cancelled)) => {
720                        Poll::Ready(Err(McpError::request_cancelled()))
721                    }
722                    Poll::Ready(Err(RecvError::Closed | RecvError::PolledAfterCompletion)) => {
723                        Poll::Ready(Err(McpError::internal_error(RESPONSE_CHANNEL_ERROR)))
724                    }
725                    Poll::Pending => Poll::Pending,
726                }
727            })
728            .await
729        } else {
730            match receiver.recv(cx).await {
731                Ok(response) => response,
732                Err(RecvError::Cancelled) => Err(McpError::request_cancelled()),
733                Err(RecvError::Closed | RecvError::PolledAfterCompletion) => {
734                    Err(McpError::internal_error(RESPONSE_CHANNEL_ERROR))
735                }
736            }
737        };
738
739        let response = match response {
740            Ok(response) => response,
741            Err(error) => {
742                if error.code == McpErrorCode::RequestCancelled
743                    && (cx.checkpoint().is_err() || self.request_is_terminal())
744                {
745                    guard.cancel();
746                } else {
747                    guard.finish();
748                }
749                return Err(error);
750            }
751        };
752
753        // A response and cancellation may become visible together. Preserve
754        // caller cancellation/budget precedence before decoding peer data.
755        if cx.checkpoint().is_err() || self.request_is_terminal() {
756            guard.cancel();
757            return Err(McpError::request_cancelled());
758        }
759
760        let result = serde_json::from_value(response)
761            .map_err(|_| McpError::internal_error(RESPONSE_PAYLOAD_ERROR));
762        guard.finish();
763        result
764    }
765}
766
767impl std::fmt::Debug for RequestSender {
768    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
769        f.debug_struct("RequestSender")
770            .field("pending", &self.pending)
771            .finish_non_exhaustive()
772    }
773}
774
775// ============================================================================
776// Sampling Sender Implementation
777// ============================================================================
778
779/// Sends sampling requests to the client via the transport.
780#[derive(Clone)]
781pub struct TransportSamplingSender {
782    sender: RequestSender,
783    request_context: McpContext,
784}
785
786impl TransportSamplingSender {
787    /// Creates a sampling sender bound to the originating handler request.
788    pub fn new(sender: RequestSender, request_context: McpContext) -> Self {
789        Self {
790            sender,
791            request_context,
792        }
793    }
794}
795
796impl SamplingSender for TransportSamplingSender {
797    fn create_message(
798        &self,
799        request: SamplingRequest,
800    ) -> std::pin::Pin<Box<dyn std::future::Future<Output = McpResult<SamplingResponse>> + Send + '_>>
801    {
802        Box::pin(async move {
803            // Convert to protocol types
804            let params = fastmcp_protocol::CreateMessageParams {
805                messages: request
806                    .messages
807                    .into_iter()
808                    .map(|m| fastmcp_protocol::SamplingMessage {
809                        role: match m.role {
810                            SamplingRole::User => fastmcp_protocol::Role::User,
811                            SamplingRole::Assistant => fastmcp_protocol::Role::Assistant,
812                        },
813                        content: fastmcp_protocol::SamplingContent::Text { text: m.text },
814                    })
815                    .collect(),
816                max_tokens: fastmcp_protocol::JsonInteger::from(u64::from(request.max_tokens)),
817                system_prompt: request.system_prompt,
818                temperature: request.temperature,
819                stop_sequences: request.stop_sequences,
820                model_preferences: if request.model_hints.is_empty() {
821                    None
822                } else {
823                    Some(fastmcp_protocol::ModelPreferences {
824                        hints: request
825                            .model_hints
826                            .into_iter()
827                            .map(|name| fastmcp_protocol::ModelHint { name: Some(name) })
828                            .collect(),
829                        ..Default::default()
830                    })
831                },
832                include_context: None,
833                metadata: None,
834                meta: None,
835            };
836
837            let params_value = serde_json::to_value(&params)
838                .map_err(|_| McpError::internal_error(REQUEST_PAYLOAD_ERROR))?;
839
840            self.request_context
841                .checkpoint()
842                .map_err(|_| McpError::request_cancelled())?;
843            let result: fastmcp_protocol::CreateMessageResult = self
844                .sender
845                .send_request(
846                    self.request_context.cx(),
847                    "sampling/createMessage",
848                    params_value,
849                )
850                .await?;
851
852            if result.role != fastmcp_protocol::Role::Assistant {
853                return Err(McpError::internal_error(RESPONSE_PAYLOAD_ERROR));
854            }
855
856            Ok(SamplingResponse {
857                text: match result.content {
858                    fastmcp_protocol::SamplingContent::Text { text } => text,
859                    fastmcp_protocol::SamplingContent::Image { data, mime_type } => {
860                        format!("[image: {} bytes, type: {}]", data.len(), mime_type)
861                    }
862                },
863                model: result.model,
864                stop_reason: SamplingStopReason::from_wire_value(result.stop_reason),
865            })
866        })
867    }
868}
869
870// ============================================================================
871// Elicitation Sender Implementation
872// ============================================================================
873
874/// Sends elicitation requests to the client via the transport.
875#[derive(Clone)]
876pub struct TransportElicitationSender {
877    sender: RequestSender,
878    request_context: McpContext,
879}
880
881impl TransportElicitationSender {
882    /// Creates an elicitation sender bound to the originating handler request.
883    pub fn new(sender: RequestSender, request_context: McpContext) -> Self {
884        Self {
885            sender,
886            request_context,
887        }
888    }
889}
890
891impl ElicitationSender for TransportElicitationSender {
892    fn elicit(
893        &self,
894        request: ElicitationRequest,
895    ) -> std::pin::Pin<
896        Box<dyn std::future::Future<Output = McpResult<ElicitationResponse>> + Send + '_>,
897    > {
898        Box::pin(async move {
899            let request_mode = request.mode;
900            let params_value = match request_mode {
901                ElicitationMode::Form => {
902                    let requested_schema = request.schema.ok_or_else(|| {
903                        McpError::invalid_params(INVALID_ELICITATION_REQUEST_ERROR)
904                    })?;
905                    let params = fastmcp_protocol::ElicitRequestFormParams {
906                        mode: fastmcp_protocol::ElicitMode::Form,
907                        message: request.message.clone(),
908                        requested_schema,
909                    };
910                    serde_json::to_value(&params)
911                        .map_err(|_| McpError::internal_error(REQUEST_PAYLOAD_ERROR))?
912                }
913                ElicitationMode::Url => {
914                    let url = request
915                        .url
916                        .filter(|value| !value.is_empty())
917                        .ok_or_else(|| {
918                            McpError::invalid_params(INVALID_ELICITATION_REQUEST_ERROR)
919                        })?;
920                    let elicitation_id = request
921                        .elicitation_id
922                        .filter(|value| !value.is_empty())
923                        .ok_or_else(|| {
924                        McpError::invalid_params(INVALID_ELICITATION_REQUEST_ERROR)
925                    })?;
926                    let params = fastmcp_protocol::ElicitRequestUrlParams {
927                        mode: fastmcp_protocol::ElicitMode::Url,
928                        message: request.message.clone(),
929                        url,
930                        elicitation_id,
931                    };
932                    serde_json::to_value(&params)
933                        .map_err(|_| McpError::internal_error(REQUEST_PAYLOAD_ERROR))?
934                }
935            };
936
937            self.request_context
938                .checkpoint()
939                .map_err(|_| McpError::request_cancelled())?;
940            let result: fastmcp_protocol::ElicitResult = self
941                .sender
942                .send_request(
943                    self.request_context.cx(),
944                    "elicitation/create",
945                    params_value,
946                )
947                .await?;
948
949            let action = match result.action {
950                fastmcp_protocol::ElicitAction::Accept => ElicitationAction::Accept,
951                fastmcp_protocol::ElicitAction::Decline => ElicitationAction::Decline,
952                fastmcp_protocol::ElicitAction::Cancel => ElicitationAction::Cancel,
953            };
954
955            // Decline/cancel content is not accepted form data. It remains a
956            // wire-level SHOULD deviation in the full protocol design, but
957            // this legacy core response has no quarantine slot, so discard it
958            // instead of exposing it to business logic. Accepted URL mode must
959            // never carry in-band form content.
960            let content = match (request_mode, result.action, result.content) {
961                (ElicitationMode::Form, fastmcp_protocol::ElicitAction::Accept, Some(content)) => {
962                    content
963                }
964                (ElicitationMode::Form, fastmcp_protocol::ElicitAction::Accept, None)
965                | (ElicitationMode::Url, fastmcp_protocol::ElicitAction::Accept, Some(_)) => {
966                    return Err(McpError::internal_error(RESPONSE_PAYLOAD_ERROR));
967                }
968                (
969                    _,
970                    fastmcp_protocol::ElicitAction::Decline
971                    | fastmcp_protocol::ElicitAction::Cancel,
972                    _,
973                )
974                | (ElicitationMode::Url, fastmcp_protocol::ElicitAction::Accept, None) => {
975                    return Ok(ElicitationResponse {
976                        action,
977                        content: None,
978                    });
979                }
980            };
981
982            // Convert HashMap<String, ElicitContentValue> to HashMap<String, serde_json::Value>.
983            let content = {
984                let mut map = std::collections::HashMap::new();
985                for (key, value) in content {
986                    let json_value = match value {
987                        fastmcp_protocol::ElicitContentValue::Null => serde_json::Value::Null,
988                        fastmcp_protocol::ElicitContentValue::Bool(b) => serde_json::Value::Bool(b),
989                        fastmcp_protocol::ElicitContentValue::Int(i) => serde_json::to_value(i)
990                            .map_err(|_| McpError::internal_error(RESPONSE_PAYLOAD_ERROR))?,
991                        fastmcp_protocol::ElicitContentValue::Float(f) => {
992                            serde_json::Number::from_f64(f)
993                                .map(serde_json::Value::Number)
994                                .unwrap_or(serde_json::Value::Null)
995                        }
996                        fastmcp_protocol::ElicitContentValue::String(s) => {
997                            serde_json::Value::String(s)
998                        }
999                        fastmcp_protocol::ElicitContentValue::StringArray(arr) => {
1000                            serde_json::Value::Array(
1001                                arr.into_iter().map(serde_json::Value::String).collect(),
1002                            )
1003                        }
1004                    };
1005                    map.insert(key, json_value);
1006                }
1007                Some(map)
1008            };
1009
1010            Ok(ElicitationResponse { action, content })
1011        })
1012    }
1013}
1014
1015// ============================================================================
1016// Roots Provider Implementation
1017// ============================================================================
1018
1019/// Provider for filesystem roots from the client.
1020#[derive(Clone)]
1021pub struct TransportRootsProvider {
1022    sender: RequestSender,
1023    request_context: McpContext,
1024}
1025
1026impl TransportRootsProvider {
1027    /// Creates a roots provider bound to the originating handler request.
1028    ///
1029    /// The provider retains the full framework context, rather than its raw
1030    /// `Cx`, so its reverse `roots/list` request observes the originating
1031    /// request lease, framework budget ceiling, and cancellation domain.
1032    pub fn new(sender: RequestSender, request_context: McpContext) -> Self {
1033        Self {
1034            sender,
1035            request_context,
1036        }
1037    }
1038
1039    /// Lists the filesystem roots from the client.
1040    pub async fn list_roots(&self) -> McpResult<Vec<fastmcp_protocol::Root>> {
1041        self.request_context
1042            .checkpoint()
1043            .map_err(|_| McpError::request_cancelled())?;
1044        let request = self.sender.send_request(
1045            self.request_context.cx(),
1046            "roots/list",
1047            serde_json::json!({}),
1048        );
1049        let result: fastmcp_protocol::ListRootsResult = match self.request_context.budget().deadline
1050        {
1051            Some(deadline) => asupersync::time::timeout_at(deadline, request)
1052                .await
1053                .map_err(|_| McpError::request_cancelled())??,
1054            None => request.await?,
1055        };
1056        self.request_context
1057            .ensure_live()
1058            .map_err(|_| McpError::request_cancelled())?;
1059        Ok(result.roots)
1060    }
1061}
1062
1063impl RootsProvider for TransportRootsProvider {
1064    fn list_roots(
1065        &self,
1066    ) -> std::pin::Pin<Box<dyn std::future::Future<Output = McpResult<Vec<ClientRoot>>> + Send + '_>>
1067    {
1068        Box::pin(async move {
1069            let roots = TransportRootsProvider::list_roots(self).await?;
1070            Ok(roots
1071                .into_iter()
1072                .map(|root| ClientRoot {
1073                    uri: root.uri,
1074                    name: root.name,
1075                })
1076                .collect())
1077        })
1078    }
1079}
1080
1081// ============================================================================
1082// Final MRTR Embedded Input Exchanges
1083// ============================================================================
1084
1085/// The three server-to-client input kinds represented by final MRTR.
1086///
1087/// These are embedded descriptors inside an `inputRequests` map. They are not
1088/// independent JSON-RPC requests and never receive a JSON-RPC ID.
1089#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1090pub enum MrtrInputKind {
1091    /// `elicitation/create` with an [`fastmcp_protocol::ElicitResult`] response.
1092    Elicitation,
1093    /// `sampling/createMessage` with a [`fastmcp_protocol::FinalCreateMessageResult`] response.
1094    Sampling,
1095    /// `roots/list` with a [`fastmcp_protocol::ListRootsResult`] response.
1096    Roots,
1097}
1098
1099impl MrtrInputKind {
1100    /// Returns the exact embedded input request method.
1101    #[must_use]
1102    pub const fn method(self) -> &'static str {
1103        match self {
1104            Self::Elicitation => "elicitation/create",
1105            Self::Sampling => "sampling/createMessage",
1106            Self::Roots => "roots/list",
1107        }
1108    }
1109}
1110
1111/// Elicitation parameters selected by the negotiated protocol era.
1112///
1113/// Final embedded descriptors and exact-2024 JSON-RPC parameters deliberately
1114/// remain separate: the former has no legacy `elicitationId`, while the latter
1115/// requires one.
1116#[derive(Debug, Clone)]
1117pub enum DualEraElicitationParams {
1118    /// Exact-2024 reverse JSON-RPC elicitation parameters.
1119    Legacy2024(fastmcp_protocol::ElicitRequestParams),
1120    /// Final embedded MRTR elicitation parameters.
1121    Modern2026(fastmcp_protocol::FinalEmbeddedElicitationParams),
1122}
1123
1124/// Sampling parameters selected by the negotiated protocol era.
1125///
1126/// Exact-2024 reverse JSON-RPC retains its legacy request model. Final MRTR
1127/// descriptors use the distinct embedded final model so tool declarations,
1128/// tool-choice controls, and final sampling content cannot be lost at the
1129/// era boundary.
1130#[derive(Debug, Clone)]
1131pub enum DualEraSamplingParams {
1132    /// Exact-2024 reverse JSON-RPC sampling parameters.
1133    Legacy2024(fastmcp_protocol::CreateMessageParams),
1134    /// Final embedded MRTR sampling parameters.
1135    Modern2026(fastmcp_protocol::FinalEmbeddedCreateMessageParams),
1136}
1137
1138#[derive(Debug, Clone)]
1139enum MrtrInputParams {
1140    LegacyElicitation(fastmcp_protocol::ElicitRequestParams),
1141    FinalElicitation(fastmcp_protocol::FinalEmbeddedElicitationParams),
1142    FinalSampling(fastmcp_protocol::FinalEmbeddedCreateMessageParams),
1143    Json(serde_json::Value),
1144}
1145
1146/// One final MRTR embedded input descriptor.
1147///
1148/// Its wire representation is exactly `{ "method": ..., "params": ... }`
1149/// when it has parameters, or `{ "method": "roots/list" }` for roots. It
1150/// deliberately has no JSON-RPC envelope, correlation ID, or inherited outer
1151/// request metadata.
1152#[derive(Debug, Clone)]
1153pub struct MrtrInputRequest {
1154    kind: MrtrInputKind,
1155    params: Option<MrtrInputParams>,
1156}
1157
1158impl MrtrInputRequest {
1159    /// Creates an exact-2024 form or URL elicitation request.
1160    ///
1161    #[must_use]
1162    pub fn legacy_elicitation(params: fastmcp_protocol::ElicitRequestParams) -> Self {
1163        Self {
1164            kind: MrtrInputKind::Elicitation,
1165            params: Some(MrtrInputParams::LegacyElicitation(params)),
1166        }
1167    }
1168
1169    /// Creates one final-era elicitation descriptor without translating it
1170    /// through the exact-2024 request shape.
1171    fn final_elicitation(
1172        params: fastmcp_protocol::FinalEmbeddedElicitationParams,
1173    ) -> McpResult<Self> {
1174        serde_json::to_value(&params)
1175            .map_err(|_| McpError::internal_error(REQUEST_PAYLOAD_ERROR))?;
1176        Ok(Self {
1177            kind: MrtrInputKind::Elicitation,
1178            params: Some(MrtrInputParams::FinalElicitation(params)),
1179        })
1180    }
1181
1182    /// Creates a final sampling input descriptor.
1183    ///
1184    /// The embedded descriptor must not inherit the outer request's metadata,
1185    /// so this safe constructor always omits `_meta`.
1186    ///
1187    /// # Errors
1188    ///
1189    /// Returns an internal error only if its protocol parameters cannot be
1190    /// represented as JSON.
1191    pub fn sampling(params: fastmcp_protocol::FinalEmbeddedCreateMessageParams) -> McpResult<Self> {
1192        serde_json::to_value(&params)
1193            .map_err(|_| McpError::internal_error(REQUEST_PAYLOAD_ERROR))?;
1194        Ok(Self {
1195            kind: MrtrInputKind::Sampling,
1196            params: Some(MrtrInputParams::FinalSampling(params)),
1197        })
1198    }
1199
1200    /// Creates a final roots input descriptor with omitted parameters.
1201    #[must_use]
1202    pub const fn roots() -> Self {
1203        Self {
1204            kind: MrtrInputKind::Roots,
1205            params: None,
1206        }
1207    }
1208
1209    /// Returns the input's exact response kind.
1210    #[must_use]
1211    pub const fn kind(&self) -> MrtrInputKind {
1212        self.kind
1213    }
1214
1215    /// Decodes one handler-declared embedded input descriptor.
1216    ///
1217    /// Only the three final MRTR methods are admitted. In particular, a
1218    /// handler cannot smuggle an arbitrary JSON-RPC request or outer metadata
1219    /// through the framework-minted input-required result.
1220    pub(crate) fn from_wire(value: &serde_json::Value) -> McpResult<Self> {
1221        let Some(object) = value.as_object() else {
1222            return Err(McpError::invalid_params(MRTR_INPUT_MAP_ERROR));
1223        };
1224        if object.len() > 2 || object.keys().any(|key| key != "method" && key != "params") {
1225            return Err(McpError::invalid_params(MRTR_INPUT_MAP_ERROR));
1226        }
1227        let Some(method) = object.get("method").and_then(serde_json::Value::as_str) else {
1228            return Err(McpError::invalid_params(MRTR_INPUT_MAP_ERROR));
1229        };
1230        let params = object.get("params");
1231        match method {
1232            "elicitation/create" => {
1233                let params =
1234                    params.ok_or_else(|| McpError::invalid_params(MRTR_INPUT_MAP_ERROR))?;
1235                Self::final_elicitation(
1236                    serde_json::from_value(params.clone())
1237                        .map_err(|_| McpError::invalid_params(MRTR_INPUT_MAP_ERROR))?,
1238                )
1239            }
1240            "sampling/createMessage" => {
1241                let params =
1242                    params.ok_or_else(|| McpError::invalid_params(MRTR_INPUT_MAP_ERROR))?;
1243                Self::sampling(
1244                    serde_json::from_value(params.clone())
1245                        .map_err(|_| McpError::invalid_params(MRTR_INPUT_MAP_ERROR))?,
1246                )
1247            }
1248            "roots/list" if params.is_none() => Ok(Self::roots()),
1249            _ => Err(McpError::invalid_params(MRTR_INPUT_MAP_ERROR)),
1250        }
1251    }
1252
1253    fn with_params<T: Serialize>(kind: MrtrInputKind, params: T) -> McpResult<Self> {
1254        let params = serde_json::to_value(params)
1255            .map_err(|_| McpError::internal_error(REQUEST_PAYLOAD_ERROR))?;
1256        Ok(Self {
1257            kind,
1258            params: Some(MrtrInputParams::Json(params)),
1259        })
1260    }
1261
1262    fn into_legacy_params(self) -> McpResult<serde_json::Value> {
1263        match self.params {
1264            None => Ok(serde_json::json!({})),
1265            Some(MrtrInputParams::LegacyElicitation(params)) => serde_json::to_value(params)
1266                .map_err(|_| McpError::internal_error(REQUEST_PAYLOAD_ERROR)),
1267            Some(MrtrInputParams::FinalElicitation(_)) => {
1268                Err(McpError::invalid_params(INVALID_ELICITATION_REQUEST_ERROR))
1269            }
1270            Some(MrtrInputParams::FinalSampling(_)) => Err(McpError::invalid_params(
1271                "Final sampling cannot be sent as exact-2024 reverse JSON-RPC",
1272            )),
1273            Some(MrtrInputParams::Json(params)) => Ok(params),
1274        }
1275    }
1276}
1277
1278impl Serialize for MrtrInputRequest {
1279    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
1280    where
1281        S: serde::Serializer,
1282    {
1283        let mut descriptor = serializer.serialize_struct(
1284            "MrtrInputRequest",
1285            if self.params.is_some() { 2 } else { 1 },
1286        )?;
1287        descriptor.serialize_field("method", self.kind.method())?;
1288        if let Some(params) = &self.params {
1289            match params {
1290                MrtrInputParams::LegacyElicitation(params) => {
1291                    descriptor.serialize_field("params", params)?;
1292                }
1293                MrtrInputParams::FinalElicitation(params) => {
1294                    descriptor.serialize_field("params", params)?;
1295                }
1296                MrtrInputParams::FinalSampling(params) => {
1297                    descriptor.serialize_field("params", params)?;
1298                }
1299                MrtrInputParams::Json(params) => {
1300                    descriptor.serialize_field("params", params)?;
1301                }
1302            }
1303        }
1304        descriptor.end()
1305    }
1306}
1307
1308/// A typed input response whose value can be accepted only for the matching
1309/// [`MrtrInputKind`] recorded at issuance time.
1310#[derive(Debug, Clone)]
1311pub struct MrtrInputResponse {
1312    kind: MrtrInputKind,
1313    value: serde_json::Value,
1314}
1315
1316impl MrtrInputResponse {
1317    /// Creates an elicitation response value.
1318    ///
1319    /// # Errors
1320    ///
1321    /// Returns an internal error only if its protocol value cannot be
1322    /// represented as JSON.
1323    pub fn elicitation(value: fastmcp_protocol::ElicitResult) -> McpResult<Self> {
1324        Self::with_value(MrtrInputKind::Elicitation, value)
1325    }
1326
1327    /// Creates a sampling response value.
1328    ///
1329    /// # Errors
1330    ///
1331    /// Returns an internal error only if its protocol value cannot be
1332    /// represented as JSON.
1333    pub fn sampling(value: fastmcp_protocol::FinalCreateMessageResult) -> McpResult<Self> {
1334        Self::with_value(MrtrInputKind::Sampling, value)
1335    }
1336
1337    /// Creates a roots response value.
1338    ///
1339    /// # Errors
1340    ///
1341    /// Returns an internal error only if its protocol value cannot be
1342    /// represented as JSON.
1343    pub fn roots(value: fastmcp_protocol::ListRootsResult) -> McpResult<Self> {
1344        Self::with_value(MrtrInputKind::Roots, value)
1345    }
1346
1347    /// Returns the response's exact kind.
1348    #[must_use]
1349    pub const fn kind(&self) -> MrtrInputKind {
1350        self.kind
1351    }
1352
1353    /// Returns this value as the elicitation result it was admitted as.
1354    ///
1355    /// # Errors
1356    ///
1357    /// Returns `InvalidParams` when the caller asks for the wrong response
1358    /// kind. This is a handler-facing resume surface, not a wire decoder.
1359    pub fn elicitation_result(&self) -> McpResult<fastmcp_protocol::ElicitResult> {
1360        if self.kind != MrtrInputKind::Elicitation {
1361            return Err(McpError::invalid_params(MRTR_RESPONSE_KIND_ERROR));
1362        }
1363        serde_json::from_value(self.value.clone())
1364            .map_err(|_| McpError::internal_error(MRTR_RESPONSE_KIND_ERROR))
1365    }
1366
1367    /// Returns this value as the sampling result it was admitted as.
1368    pub fn sampling_result(&self) -> McpResult<fastmcp_protocol::FinalCreateMessageResult> {
1369        if self.kind != MrtrInputKind::Sampling {
1370            return Err(McpError::invalid_params(MRTR_RESPONSE_KIND_ERROR));
1371        }
1372        serde_json::from_value(self.value.clone())
1373            .map_err(|_| McpError::internal_error(MRTR_RESPONSE_KIND_ERROR))
1374    }
1375
1376    /// Returns this value as the roots result it was admitted as.
1377    pub fn roots_result(&self) -> McpResult<fastmcp_protocol::ListRootsResult> {
1378        if self.kind != MrtrInputKind::Roots {
1379            return Err(McpError::invalid_params(MRTR_RESPONSE_KIND_ERROR));
1380        }
1381        serde_json::from_value(self.value.clone())
1382            .map_err(|_| McpError::internal_error(MRTR_RESPONSE_KIND_ERROR))
1383    }
1384
1385    fn from_wire(kind: MrtrInputKind, value: serde_json::Value) -> McpResult<Self> {
1386        let response = match kind {
1387            MrtrInputKind::Elicitation => Self::elicitation(
1388                serde_json::from_value(value)
1389                    .map_err(|_| McpError::invalid_params(MRTR_RESPONSE_KIND_ERROR))?,
1390            )?,
1391            MrtrInputKind::Sampling => Self::sampling(
1392                serde_json::from_value(value)
1393                    .map_err(|_| McpError::invalid_params(MRTR_RESPONSE_KIND_ERROR))?,
1394            )?,
1395            MrtrInputKind::Roots => Self::roots(
1396                serde_json::from_value(value)
1397                    .map_err(|_| McpError::invalid_params(MRTR_RESPONSE_KIND_ERROR))?,
1398            )?,
1399        };
1400        Ok(response)
1401    }
1402
1403    fn with_value<T: Serialize>(kind: MrtrInputKind, value: T) -> McpResult<Self> {
1404        let value = serde_json::to_value(value)
1405            .map_err(|_| McpError::internal_error(RESPONSE_PAYLOAD_ERROR))?;
1406        Ok(Self { kind, value })
1407    }
1408}
1409
1410impl Serialize for MrtrInputResponse {
1411    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
1412    where
1413        S: serde::Serializer,
1414    {
1415        self.value.serialize(serializer)
1416    }
1417}
1418
1419/// A unique, bounded map of final embedded MRTR input requests.
1420#[derive(Debug, Clone, Default)]
1421pub struct MrtrInputRequests {
1422    entries: BTreeMap<String, MrtrInputRequest>,
1423}
1424
1425impl MrtrInputRequests {
1426    /// Creates a unique map of embedded input request descriptors.
1427    ///
1428    /// # Errors
1429    ///
1430    /// Returns `InvalidParams` for an empty or duplicate key, or when more
1431    /// than [`HARD_MAX_MRTR_INPUT_REQUESTS_PER_ROUND`] descriptors are given.
1432    pub fn new(entries: impl IntoIterator<Item = (String, MrtrInputRequest)>) -> McpResult<Self> {
1433        let mut result = Self::default();
1434        for (key, request) in entries {
1435            result.insert(key, request)?;
1436        }
1437        Ok(result)
1438    }
1439
1440    /// Returns whether this map contains no embedded input descriptors.
1441    #[must_use]
1442    pub fn is_empty(&self) -> bool {
1443        self.entries.is_empty()
1444    }
1445
1446    /// Returns the number of embedded input descriptors.
1447    #[must_use]
1448    pub fn len(&self) -> usize {
1449        self.entries.len()
1450    }
1451
1452    /// Looks up one embedded input descriptor by its server-issued key.
1453    #[must_use]
1454    pub fn get(&self, key: &str) -> Option<&MrtrInputRequest> {
1455        self.entries.get(key)
1456    }
1457
1458    /// Iterates over server-issued input keys and their embedded descriptors.
1459    pub fn iter(&self) -> impl Iterator<Item = (&str, &MrtrInputRequest)> {
1460        self.entries
1461            .iter()
1462            .map(|(key, request)| (key.as_str(), request))
1463    }
1464
1465    fn insert(&mut self, key: String, request: MrtrInputRequest) -> McpResult<()> {
1466        if key.is_empty()
1467            || self.entries.len() >= HARD_MAX_MRTR_INPUT_REQUESTS_PER_ROUND
1468            || self.entries.contains_key(&key)
1469        {
1470            return Err(McpError::invalid_params(MRTR_INPUT_MAP_ERROR));
1471        }
1472        self.entries.insert(key, request);
1473        Ok(())
1474    }
1475
1476    fn unresolved_after(&self, responses: &MrtrInputResponses) -> Self {
1477        Self {
1478            entries: self
1479                .entries
1480                .iter()
1481                .filter(|(key, _)| responses.get(key).is_none())
1482                .map(|(key, request)| (key.clone(), request.clone()))
1483                .collect(),
1484        }
1485    }
1486}
1487
1488impl Serialize for MrtrInputRequests {
1489    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
1490    where
1491        S: serde::Serializer,
1492    {
1493        self.entries.serialize(serializer)
1494    }
1495}
1496
1497/// An ordered, unique, bounded collection of final MRTR input response values.
1498///
1499/// The vector preserves the exact accepted retry order all the way through
1500/// [`MrtrCompletedInputs`]. A separate lookup index supports handler key
1501/// access without normalizing the handler-visible response sequence into a
1502/// sorted map.
1503#[derive(Debug, Clone, Default)]
1504pub struct MrtrInputResponses {
1505    entries: Vec<(String, MrtrInputResponse)>,
1506    index: HashMap<String, usize>,
1507}
1508
1509impl MrtrInputResponses {
1510    /// Creates a unique ordered collection of embedded input responses.
1511    ///
1512    /// # Errors
1513    ///
1514    /// Returns `InvalidParams` for an empty or duplicate key, or when more
1515    /// than [`HARD_MAX_MRTR_INPUT_REQUESTS_PER_ROUND`] response values are
1516    /// given.
1517    pub fn new(entries: impl IntoIterator<Item = (String, MrtrInputResponse)>) -> McpResult<Self> {
1518        let mut result = Self::default();
1519        for (key, response) in entries {
1520            result.insert(key, response)?;
1521        }
1522        Ok(result)
1523    }
1524
1525    /// Returns whether this collection contains no input responses.
1526    #[must_use]
1527    pub fn is_empty(&self) -> bool {
1528        self.entries.is_empty()
1529    }
1530
1531    /// Returns the number of input response values.
1532    #[must_use]
1533    pub fn len(&self) -> usize {
1534        self.entries.len()
1535    }
1536
1537    /// Looks up one response by its server-issued input key.
1538    #[must_use]
1539    pub fn get(&self, key: &str) -> Option<&MrtrInputResponse> {
1540        self.index
1541            .get(key)
1542            .and_then(|index| self.entries.get(*index))
1543            .map(|(_, response)| response)
1544    }
1545
1546    /// Iterates over input-response keys and values.
1547    pub fn iter(&self) -> impl Iterator<Item = (&str, &MrtrInputResponse)> {
1548        self.entries
1549            .iter()
1550            .map(|(key, response)| (key.as_str(), response))
1551    }
1552
1553    fn insert(&mut self, key: String, response: MrtrInputResponse) -> McpResult<()> {
1554        if key.is_empty()
1555            || self.entries.len() >= HARD_MAX_MRTR_INPUT_REQUESTS_PER_ROUND
1556            || self.index.contains_key(&key)
1557        {
1558            return Err(McpError::invalid_params(MRTR_INPUT_MAP_ERROR));
1559        }
1560        self.index.insert(key.clone(), self.entries.len());
1561        self.entries.push((key, response));
1562        Ok(())
1563    }
1564
1565    /// Appends one registry-admitted response only if this exchange has not
1566    /// accepted its key already. The registry has separately bounded total
1567    /// exchange growth, so this must not apply the single-round constructor
1568    /// ceiling to accumulated partial retries.
1569    fn append_accepted_if_absent(&mut self, key: &str, response: &MrtrInputResponse) -> bool {
1570        if self.index.contains_key(key) {
1571            return false;
1572        }
1573        self.index.insert(key.to_owned(), self.entries.len());
1574        self.entries.push((key.to_owned(), response.clone()));
1575        true
1576    }
1577}
1578
1579impl Serialize for MrtrInputResponses {
1580    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
1581    where
1582        S: serde::Serializer,
1583    {
1584        let mut map = serializer.serialize_map(Some(self.entries.len()))?;
1585        for (key, response) in &self.entries {
1586            map.serialize_entry(key, response)?;
1587        }
1588        map.end()
1589    }
1590}
1591
1592/// A server-minted final MRTR request state.
1593///
1594/// The wire representation is opaque. Its `Debug` implementation is redacted,
1595/// and only [`MrtrInputRequired`] serializes it into a server response.
1596#[derive(Clone)]
1597pub struct MrtrRequestState(String);
1598
1599impl std::fmt::Debug for MrtrRequestState {
1600    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1601        formatter.write_str("MrtrRequestState([redacted])")
1602    }
1603}
1604
1605impl Serialize for MrtrRequestState {
1606    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
1607    where
1608        S: serde::Serializer,
1609    {
1610        serializer.serialize_str(&self.0)
1611    }
1612}
1613
1614/// A final `input_required` result emitted by the server.
1615///
1616/// Safe server construction always includes a protected request state. An
1617/// empty request map is emitted as a state-only result, which permits an
1618/// immediate retry without manufacturing an empty input map.
1619#[derive(Debug, Clone)]
1620pub struct MrtrInputRequired {
1621    input_requests: Option<MrtrInputRequests>,
1622    request_state: MrtrRequestState,
1623}
1624
1625impl MrtrInputRequired {
1626    /// Returns the embedded request map, if this result needs client input.
1627    #[must_use]
1628    pub fn input_requests(&self) -> Option<&MrtrInputRequests> {
1629        self.input_requests.as_ref()
1630    }
1631}
1632
1633impl Serialize for MrtrInputRequired {
1634    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
1635    where
1636        S: serde::Serializer,
1637    {
1638        let mut result = serializer.serialize_struct(
1639            "MrtrInputRequired",
1640            if self.input_requests.is_some() { 3 } else { 2 },
1641        )?;
1642        result.serialize_field("resultType", "input_required")?;
1643        if let Some(input_requests) = &self.input_requests {
1644            result.serialize_field("inputRequests", input_requests)?;
1645        }
1646        result.serialize_field("requestState", &self.request_state)?;
1647        result.end()
1648    }
1649}
1650
1651/// The outcome of consuming one final MRTR retry.
1652#[derive(Debug, Clone)]
1653pub enum MrtrRetry {
1654    /// More client input is needed; only unsatisfied keys are reissued with a
1655    /// fresh request state.
1656    InputRequired(MrtrInputRequired),
1657    /// All currently requested input values were accepted and type-checked.
1658    Complete(MrtrCompletedInputs),
1659}
1660
1661/// Accumulated, type-bound responses from a completed MRTR input exchange.
1662#[derive(Debug, Clone)]
1663pub struct MrtrCompletedInputs {
1664    responses: MrtrInputResponses,
1665}
1666
1667impl MrtrCompletedInputs {
1668    /// Returns every accepted response, including values accepted in earlier
1669    /// partial retries of this logical exchange.
1670    #[must_use]
1671    pub fn responses(&self) -> &MrtrInputResponses {
1672        &self.responses
1673    }
1674
1675    /// Returns one framework-admitted elicitation response by its issued key.
1676    pub fn elicitation(&self, key: &str) -> McpResult<Option<fastmcp_protocol::ElicitResult>> {
1677        self.responses
1678            .get(key)
1679            .map(MrtrInputResponse::elicitation_result)
1680            .transpose()
1681    }
1682
1683    /// Returns one framework-admitted sampling response by its issued key.
1684    pub fn sampling(
1685        &self,
1686        key: &str,
1687    ) -> McpResult<Option<fastmcp_protocol::FinalCreateMessageResult>> {
1688        self.responses
1689            .get(key)
1690            .map(MrtrInputResponse::sampling_result)
1691            .transpose()
1692    }
1693
1694    /// Returns one framework-admitted roots response by its issued key.
1695    pub fn roots(&self, key: &str) -> McpResult<Option<fastmcp_protocol::ListRootsResult>> {
1696        self.responses
1697            .get(key)
1698            .map(MrtrInputResponse::roots_result)
1699            .transpose()
1700    }
1701}
1702
1703#[derive(Debug, Clone)]
1704struct ExpectedInputLedger {
1705    kinds: BTreeMap<String, MrtrInputKind>,
1706}
1707
1708impl ExpectedInputLedger {
1709    fn from_requests(requests: &MrtrInputRequests) -> Self {
1710        Self {
1711            kinds: requests
1712                .entries
1713                .iter()
1714                .map(|(key, request)| (key.clone(), request.kind()))
1715                .collect(),
1716        }
1717    }
1718
1719    fn get(&self, key: &str) -> Option<MrtrInputKind> {
1720        self.kinds.get(key).copied()
1721    }
1722
1723    fn is_empty(&self) -> bool {
1724        self.kinds.is_empty()
1725    }
1726}
1727
1728#[derive(Debug, Clone)]
1729struct MrtrExchange {
1730    // Only cancellation invalidates a continuation. Normal response
1731    // finalization ends the old JSON-RPC request before the client can send
1732    // its new-ID retry, so treating it as cancellation would make every MRTR
1733    // exchange unusable.
1734    owner_cancellation: McpRequestCancellation,
1735    expires_at: Instant,
1736    round: u8,
1737    total_input_requests: usize,
1738    requests: MrtrInputRequests,
1739    expected: ExpectedInputLedger,
1740    responses: MrtrInputResponses,
1741    binding: Option<MrtrExchangeBinding>,
1742}
1743
1744#[derive(Debug, Default)]
1745struct MrtrExchangeState {
1746    exchanges: HashMap<String, MrtrExchange>,
1747}
1748
1749/// Process-local, one-use final MRTR request-state storage.
1750///
1751/// A record owns the exact key-to-response-kind ledger for the inputs it
1752/// emitted. The opaque state is random, expires, cannot be replayed after a
1753/// successful retry, and is invalidated when its owning request is cancelled.
1754pub struct MrtrExchangeRegistry {
1755    state: Mutex<MrtrExchangeState>,
1756    max_states: usize,
1757    max_rounds: u8,
1758    max_inputs_per_round: usize,
1759    max_total_input_requests: usize,
1760    request_state_ttl: Duration,
1761}
1762
1763impl MrtrExchangeRegistry {
1764    /// Creates a registry with the final protocol's default bounded policy.
1765    #[must_use]
1766    pub fn new() -> Self {
1767        Self {
1768            state: Mutex::new(MrtrExchangeState::default()),
1769            max_states: DEFAULT_MAX_MRTR_REQUEST_STATES,
1770            max_rounds: DEFAULT_MAX_MRTR_ROUNDS,
1771            max_inputs_per_round: DEFAULT_MAX_MRTR_INPUT_REQUESTS_PER_ROUND,
1772            max_total_input_requests: DEFAULT_MAX_MRTR_INPUT_REQUESTS_TOTAL,
1773            request_state_ttl: DEFAULT_MRTR_REQUEST_STATE_TTL,
1774        }
1775    }
1776
1777    /// Creates a registry with caller-selected bounded limits.
1778    ///
1779    /// # Errors
1780    ///
1781    /// Returns `InvalidParams` when a limit is zero or exceeds the final
1782    /// hard ceiling.
1783    pub fn with_limits(
1784        max_states: usize,
1785        max_rounds: u8,
1786        max_inputs_per_round: usize,
1787        max_total_input_requests: usize,
1788        request_state_ttl: Duration,
1789    ) -> McpResult<Self> {
1790        if !(1..=HARD_MAX_MRTR_REQUEST_STATES).contains(&max_states)
1791            || !(1..=HARD_MAX_MRTR_ROUNDS).contains(&max_rounds)
1792            || !(1..=HARD_MAX_MRTR_INPUT_REQUESTS_PER_ROUND).contains(&max_inputs_per_round)
1793            || !(max_inputs_per_round..=HARD_MAX_MRTR_INPUT_REQUESTS_TOTAL)
1794                .contains(&max_total_input_requests)
1795            || request_state_ttl.is_zero()
1796            || request_state_ttl > HARD_MAX_MRTR_REQUEST_STATE_TTL
1797        {
1798            return Err(McpError::invalid_params(INVALID_MRTR_LIMIT_ERROR));
1799        }
1800
1801        Ok(Self {
1802            state: Mutex::new(MrtrExchangeState::default()),
1803            max_states,
1804            max_rounds,
1805            max_inputs_per_round,
1806            max_total_input_requests,
1807            request_state_ttl,
1808        })
1809    }
1810
1811    /// Issues an `input_required` result bound to the owning request.
1812    ///
1813    /// This never sends an independent JSON-RPC request. The returned state
1814    /// retains the issued request map and exact key-to-response-kind ledger for
1815    /// a later retry.
1816    ///
1817    /// # Errors
1818    ///
1819    /// Returns `RequestCancelled` if the owner was already cancelled,
1820    /// `InvalidParams` when the map exceeds the configured round bound, or an
1821    /// internal error if secure state generation fails.
1822    pub fn issue(
1823        &self,
1824        owner_cancellation: McpRequestCancellation,
1825        input_requests: MrtrInputRequests,
1826    ) -> McpResult<MrtrInputRequired> {
1827        self.issue_at(owner_cancellation, None, input_requests, Instant::now())
1828    }
1829
1830    /// Issues an `input_required` result bound to one router-admitted modern
1831    /// operation. Only [`Self::accept_wire_bound`] can consume such a state.
1832    pub(crate) fn issue_bound(
1833        &self,
1834        owner_cancellation: McpRequestCancellation,
1835        binding: MrtrExchangeBinding,
1836        input_requests: MrtrInputRequests,
1837    ) -> McpResult<MrtrInputRequired> {
1838        self.issue_at(
1839            owner_cancellation,
1840            Some(binding),
1841            input_requests,
1842            Instant::now(),
1843        )
1844    }
1845
1846    /// Consumes one client retry's input-response map.
1847    ///
1848    /// Unknown keys are inert and ignored after bounded structural admission.
1849    /// Every recognized key must carry the exact response kind recorded when
1850    /// it was issued. Partial maps are accepted and yield a fresh state with
1851    /// only unsatisfied descriptors.
1852    ///
1853    /// # Errors
1854    ///
1855    /// Returns `InvalidParams` for unknown, expired, oversized, replayed, or
1856    /// wrong-kind state/input combinations, and `RequestCancelled` if the
1857    /// owning request cancellation wins before completion.
1858    pub fn accept(
1859        &self,
1860        request_state: &str,
1861        input_responses: MrtrInputResponses,
1862    ) -> McpResult<MrtrRetry> {
1863        self.accept_at(request_state, None, input_responses, false, Instant::now())
1864    }
1865
1866    /// Decodes and consumes one router-admitted final `inputResponses` map.
1867    ///
1868    /// Final request parameter types retain input responses as JSON values.
1869    /// This boundary keeps router code from selecting a response kind itself:
1870    /// each recognized value is decoded using the kind retained for its
1871    /// server-issued input key. Unknown keys remain inert, and malformed or
1872    /// cross-kind values leave the request state available for a valid retry.
1873    ///
1874    /// # Errors
1875    ///
1876    /// Returns `InvalidParams` if the state is invalid, the response map
1877    /// exceeds this registry's configured bound, or a recognized response
1878    /// cannot be decoded as the kind that was issued for its key.
1879    pub fn accept_wire(
1880        &self,
1881        request_state: &str,
1882        input_responses: &BTreeMap<String, serde_json::Value>,
1883    ) -> McpResult<MrtrRetry> {
1884        self.accept_wire_with_binding(request_state, None, input_responses)
1885    }
1886
1887    /// Decodes and consumes a router retry only when its immutable request
1888    /// facts exactly match the state that issued it.
1889    pub(crate) fn accept_wire_bound(
1890        &self,
1891        request_state: &str,
1892        binding: &MrtrExchangeBinding,
1893        input_responses: &BTreeMap<String, serde_json::Value>,
1894    ) -> McpResult<MrtrRetry> {
1895        self.accept_wire_with_binding(request_state, Some(binding), input_responses)
1896    }
1897
1898    /// Decodes and consumes ordered protocol response entries without first
1899    /// collapsing them into a map.
1900    ///
1901    /// The final protocol decoder has already rejected duplicate keys and
1902    /// retained their wire order. Keeping that representation through this
1903    /// admission boundary ensures a raw retry cannot be normalized before the
1904    /// continuation registry has validated it.
1905    pub(crate) fn accept_final_input_responses_bound(
1906        &self,
1907        request_state: &str,
1908        binding: &MrtrExchangeBinding,
1909        input_responses: &FinalInputResponses,
1910    ) -> McpResult<MrtrRetry> {
1911        let entries = input_responses
1912            .entries()
1913            .iter()
1914            .map(|(key, response)| {
1915                serde_json::to_value(response)
1916                    .map(|value| (key.clone(), value))
1917                    .map_err(|_| {
1918                        McpError::internal_error(
1919                            "final MRTR response could not be encoded for registry admission",
1920                        )
1921                    })
1922            })
1923            .collect::<McpResult<Vec<_>>>()?;
1924        self.accept_wire_entries_with_binding(
1925            request_state,
1926            Some(binding),
1927            entries.len(),
1928            entries.iter().map(|(key, value)| (key, value)),
1929        )
1930    }
1931
1932    /// Consumes a retry whose `inputResponses` member was absent.
1933    ///
1934    /// The distinct entry point preserves the final wire contract: only an
1935    /// absent member can resume a state-only exchange. An explicitly present
1936    /// empty map still flows through [`Self::accept_wire_bound`] and remains
1937    /// invalid rather than silently becoming a state-only retry.
1938    pub(crate) fn accept_state_only_bound(
1939        &self,
1940        request_state: &str,
1941        binding: &MrtrExchangeBinding,
1942    ) -> McpResult<MrtrRetry> {
1943        self.accept_at(
1944            request_state,
1945            Some(binding),
1946            MrtrInputResponses::default(),
1947            true,
1948            Instant::now(),
1949        )
1950    }
1951
1952    /// Returns the configured response-map admission ceiling so the router can
1953    /// reject an oversized raw map before typed request decoding allocates it.
1954    #[must_use]
1955    pub(crate) const fn max_inputs_per_round(&self) -> usize {
1956        self.max_inputs_per_round
1957    }
1958
1959    fn accept_wire_with_binding(
1960        &self,
1961        request_state: &str,
1962        binding: Option<&MrtrExchangeBinding>,
1963        input_responses: &BTreeMap<String, serde_json::Value>,
1964    ) -> McpResult<MrtrRetry> {
1965        self.accept_wire_entries_with_binding(
1966            request_state,
1967            binding,
1968            input_responses.len(),
1969            input_responses.iter(),
1970        )
1971    }
1972
1973    fn accept_wire_entries_with_binding<'a, I>(
1974        &self,
1975        request_state: &str,
1976        binding: Option<&MrtrExchangeBinding>,
1977        input_responses_len: usize,
1978        input_responses: I,
1979    ) -> McpResult<MrtrRetry>
1980    where
1981        I: IntoIterator<Item = (&'a String, &'a serde_json::Value)>,
1982    {
1983        if request_state.len() > DEFAULT_MAX_MRTR_REQUEST_STATE_BYTES {
1984            return Err(McpError::invalid_params(MRTR_REQUEST_STATE_ERROR));
1985        }
1986        if input_responses_len > self.max_inputs_per_round {
1987            return Err(McpError::invalid_params(MRTR_INPUT_MAP_ERROR));
1988        }
1989
1990        let expected = {
1991            let mut state = self.lock_state();
1992            let now = Instant::now();
1993            let exchange = state
1994                .exchanges
1995                .get(request_state)
1996                .cloned()
1997                .ok_or_else(|| McpError::invalid_params(MRTR_REQUEST_STATE_ERROR))?;
1998            if now >= exchange.expires_at || exchange.owner_cancellation.is_cancel_requested() {
1999                state.exchanges.remove(request_state);
2000                return if exchange.owner_cancellation.is_cancel_requested() {
2001                    Err(McpError::request_cancelled())
2002                } else {
2003                    Err(McpError::invalid_params(MRTR_REQUEST_STATE_ERROR))
2004                };
2005            }
2006            if exchange.binding.as_ref() != binding {
2007                return Err(McpError::invalid_params(MRTR_REQUEST_STATE_ERROR));
2008            }
2009            Self::purge_stale(&mut state, now);
2010            exchange.expected
2011        };
2012
2013        let mut typed_responses = MrtrInputResponses::default();
2014        for (key, value) in input_responses {
2015            let Some(kind) = expected.get(key) else {
2016                continue;
2017            };
2018            typed_responses.insert(
2019                key.clone(),
2020                MrtrInputResponse::from_wire(kind, value.clone())?,
2021            )?;
2022        }
2023
2024        // A map that names none of the outstanding inputs is not a partial
2025        // retry. Rotating it would burn a valid continuation without making
2026        // progress, so reject it before the current state can be consumed.
2027        if typed_responses.is_empty() {
2028            return Err(McpError::invalid_params(MRTR_INPUT_MAP_ERROR));
2029        }
2030
2031        self.accept_at(
2032            request_state,
2033            binding,
2034            typed_responses,
2035            false,
2036            Instant::now(),
2037        )
2038    }
2039
2040    /// Returns the number of non-expired, non-cancelled exchanges currently
2041    /// retained by this process-local registry.
2042    #[must_use]
2043    pub fn active_len(&self) -> usize {
2044        let mut state = self.lock_state();
2045        Self::purge_stale(&mut state, Instant::now());
2046        state.exchanges.len()
2047    }
2048
2049    fn issue_at(
2050        &self,
2051        owner_cancellation: McpRequestCancellation,
2052        binding: Option<MrtrExchangeBinding>,
2053        input_requests: MrtrInputRequests,
2054        now: Instant,
2055    ) -> McpResult<MrtrInputRequired> {
2056        if owner_cancellation.is_cancel_requested() {
2057            return Err(McpError::request_cancelled());
2058        }
2059        if input_requests.len() > self.max_inputs_per_round {
2060            return Err(McpError::invalid_params(MRTR_ROUND_LIMIT_ERROR));
2061        }
2062
2063        let expires_at = now
2064            .checked_add(self.request_state_ttl)
2065            .ok_or_else(|| McpError::internal_error(MRTR_REQUEST_STATE_UNAVAILABLE_ERROR))?;
2066        let mut state = self.lock_state();
2067        Self::purge_stale(&mut state, now);
2068        if state.exchanges.len() >= self.max_states {
2069            return Err(McpError::internal_error(MRTR_ROUND_LIMIT_ERROR));
2070        }
2071
2072        let request_state = Self::allocate_request_state(&state)?;
2073        let expected = ExpectedInputLedger::from_requests(&input_requests);
2074        let input_requests = (!input_requests.is_empty()).then_some(input_requests);
2075        state.exchanges.insert(
2076            request_state.0.clone(),
2077            MrtrExchange {
2078                owner_cancellation,
2079                expires_at,
2080                round: 1,
2081                total_input_requests: input_requests.as_ref().map_or(0, MrtrInputRequests::len),
2082                expected,
2083                requests: input_requests.clone().unwrap_or_default(),
2084                responses: MrtrInputResponses::default(),
2085                binding,
2086            },
2087        );
2088        Ok(MrtrInputRequired {
2089            input_requests,
2090            request_state,
2091        })
2092    }
2093
2094    fn accept_at(
2095        &self,
2096        request_state: &str,
2097        binding: Option<&MrtrExchangeBinding>,
2098        input_responses: MrtrInputResponses,
2099        state_only_retry: bool,
2100        now: Instant,
2101    ) -> McpResult<MrtrRetry> {
2102        if request_state.len() > DEFAULT_MAX_MRTR_REQUEST_STATE_BYTES {
2103            return Err(McpError::invalid_params(MRTR_REQUEST_STATE_ERROR));
2104        }
2105        if input_responses.len() > self.max_inputs_per_round {
2106            return Err(McpError::invalid_params(MRTR_INPUT_MAP_ERROR));
2107        }
2108
2109        let mut state = self.lock_state();
2110        let Some(exchange) = state.exchanges.get(request_state).cloned() else {
2111            return Err(McpError::invalid_params(MRTR_REQUEST_STATE_ERROR));
2112        };
2113        if exchange.binding.as_ref() != binding {
2114            return Err(McpError::invalid_params(MRTR_REQUEST_STATE_ERROR));
2115        }
2116        if now >= exchange.expires_at || exchange.owner_cancellation.is_cancel_requested() {
2117            state.exchanges.remove(request_state);
2118            return if exchange.owner_cancellation.is_cancel_requested() {
2119                Err(McpError::request_cancelled())
2120            } else {
2121                Err(McpError::invalid_params(MRTR_REQUEST_STATE_ERROR))
2122            };
2123        }
2124        if state_only_retry && (!exchange.expected.is_empty() || !exchange.requests.is_empty()) {
2125            return Err(McpError::invalid_params(MRTR_INPUT_MAP_ERROR));
2126        }
2127
2128        let mut accepted_responses = exchange.responses.clone();
2129        let mut made_progress = false;
2130        for (key, response) in input_responses.iter() {
2131            if let Some(expected_kind) = exchange.expected.get(key) {
2132                if expected_kind != response.kind() {
2133                    return Err(McpError::invalid_params(MRTR_RESPONSE_KIND_ERROR));
2134                }
2135                if accepted_responses.append_accepted_if_absent(key, response) {
2136                    made_progress = true;
2137                }
2138            }
2139        }
2140
2141        if exchange.owner_cancellation.is_cancel_requested() {
2142            state.exchanges.remove(request_state);
2143            return Err(McpError::request_cancelled());
2144        }
2145
2146        // Unknown-only (or otherwise no-progress) typed retries must be as
2147        // inert as their wire-decoded counterparts. Rotating here would burn
2148        // the caller's valid continuation without accepting any outstanding
2149        // input response.
2150        if !exchange.requests.is_empty() && !made_progress {
2151            return Err(McpError::invalid_params(MRTR_INPUT_MAP_ERROR));
2152        }
2153
2154        let missing_requests = exchange.requests.unresolved_after(&accepted_responses);
2155        if missing_requests.is_empty() {
2156            state.exchanges.remove(request_state);
2157            return Ok(MrtrRetry::Complete(MrtrCompletedInputs {
2158                responses: accepted_responses,
2159            }));
2160        }
2161
2162        let next_round = exchange
2163            .round
2164            .checked_add(1)
2165            .ok_or_else(|| McpError::invalid_params(MRTR_ROUND_LIMIT_ERROR))?;
2166        let next_total = exchange
2167            .total_input_requests
2168            .checked_add(missing_requests.len())
2169            .ok_or_else(|| McpError::invalid_params(MRTR_ROUND_LIMIT_ERROR))?;
2170        if next_round > self.max_rounds
2171            || missing_requests.len() > self.max_inputs_per_round
2172            || next_total > self.max_total_input_requests
2173        {
2174            state.exchanges.remove(request_state);
2175            return Err(McpError::invalid_params(MRTR_ROUND_LIMIT_ERROR));
2176        }
2177
2178        // Generate the successor before consuming the current state. An RNG
2179        // failure therefore leaves the original exchange intact and does not
2180        // create a state/ledger gap.
2181        let next_state = Self::allocate_request_state(&state)?;
2182        let successor = MrtrExchange {
2183            owner_cancellation: exchange.owner_cancellation,
2184            expires_at: exchange.expires_at,
2185            round: next_round,
2186            total_input_requests: next_total,
2187            expected: ExpectedInputLedger::from_requests(&missing_requests),
2188            requests: missing_requests.clone(),
2189            responses: accepted_responses,
2190            binding: exchange.binding,
2191        };
2192        state.exchanges.remove(request_state);
2193        state.exchanges.insert(next_state.0.clone(), successor);
2194
2195        Ok(MrtrRetry::InputRequired(MrtrInputRequired {
2196            input_requests: Some(missing_requests),
2197            request_state: next_state,
2198        }))
2199    }
2200
2201    fn lock_state(&self) -> std::sync::MutexGuard<'_, MrtrExchangeState> {
2202        match self.state.lock() {
2203            Ok(guard) => guard,
2204            Err(poisoned) => poisoned.into_inner(),
2205        }
2206    }
2207
2208    fn purge_stale(state: &mut MrtrExchangeState, now: Instant) {
2209        state.exchanges.retain(|_, exchange| {
2210            now < exchange.expires_at && !exchange.owner_cancellation.is_cancel_requested()
2211        });
2212    }
2213
2214    fn allocate_request_state(state: &MrtrExchangeState) -> McpResult<MrtrRequestState> {
2215        for _ in 0..4 {
2216            let identifier = draw_security_identifier()
2217                .map_err(|_| McpError::internal_error(MRTR_REQUEST_STATE_UNAVAILABLE_ERROR))?;
2218            let encoded =
2219                base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(identifier.as_bytes());
2220            if !state.exchanges.contains_key(&encoded) {
2221                return Ok(MrtrRequestState(encoded));
2222            }
2223        }
2224        Err(McpError::internal_error(
2225            MRTR_REQUEST_STATE_UNAVAILABLE_ERROR,
2226        ))
2227    }
2228}
2229
2230impl std::fmt::Debug for MrtrExchangeRegistry {
2231    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
2232        formatter
2233            .debug_struct("MrtrExchangeRegistry")
2234            .field("active_len", &self.active_len())
2235            .field("max_states", &self.max_states)
2236            .field("max_rounds", &self.max_rounds)
2237            .field("max_inputs_per_round", &self.max_inputs_per_round)
2238            .field("max_total_input_requests", &self.max_total_input_requests)
2239            .finish()
2240    }
2241}
2242
2243impl Default for MrtrExchangeRegistry {
2244    fn default() -> Self {
2245        Self::new()
2246    }
2247}
2248
2249// ============================================================================
2250// Dual-era Server-to-client Boundary
2251// ============================================================================
2252
2253/// The typed outcome of one server-to-client input request.
2254///
2255/// Exact MCP 2024-11-05 completes the reverse JSON-RPC request and returns
2256/// its typed client response. MCP 2026-07-28 instead returns a server result
2257/// carrying an `input_required` exchange for the client to retry.
2258#[derive(Debug, Clone)]
2259pub enum DualEraServerToClientResult<T> {
2260    /// A response to an exact-2024 reverse JSON-RPC request.
2261    Legacy(T),
2262    /// A final-era result that requires the client to retry with input.
2263    InputRequired(MrtrInputRequired),
2264}
2265
2266/// Era-selected server-to-client input boundary.
2267///
2268/// The exact MCP 2024-11-05 variant owns the transport sender and may issue
2269/// only `sampling/createMessage`, `elicitation/create`, and `roots/list`
2270/// reverse JSON-RPC requests. The MCP 2026-07-28 variant deliberately does
2271/// not retain that sender: it can only issue and consume the bounded
2272/// [`MrtrExchangeRegistry`] `input_required` retry flow.
2273#[derive(Clone)]
2274pub enum DualEraServerToClient {
2275    /// Exact MCP 2024-11-05 reverse JSON-RPC support.
2276    Legacy2024 {
2277        /// The connection's reverse-request sender.
2278        sender: RequestSender,
2279    },
2280    /// MCP 2026-07-28 embedded input/retry support.
2281    Modern2026 {
2282        /// The server-local registry for bounded input exchanges.
2283        exchanges: Arc<MrtrExchangeRegistry>,
2284    },
2285}
2286
2287impl DualEraServerToClient {
2288    /// Selects the sole server-to-client mechanism for one negotiated era.
2289    ///
2290    /// The legacy sender is intentionally consumed and discarded for MCP
2291    /// 2026-07-28. That makes reverse JSON-RPC unavailable in the final-era
2292    /// variant even if the connection also has a transport send callback.
2293    #[must_use]
2294    pub fn new(
2295        era: ProtocolEra,
2296        legacy_sender: RequestSender,
2297        exchanges: Arc<MrtrExchangeRegistry>,
2298    ) -> Self {
2299        match era {
2300            ProtocolEra::Legacy2024 => Self::Legacy2024 {
2301                sender: legacy_sender,
2302            },
2303            ProtocolEra::Modern2026 => Self::Modern2026 { exchanges },
2304        }
2305    }
2306
2307    /// Returns the exact negotiated era selected by this boundary.
2308    #[must_use]
2309    pub const fn era(&self) -> ProtocolEra {
2310        match self {
2311            Self::Legacy2024 { .. } => ProtocolEra::Legacy2024,
2312            Self::Modern2026 { .. } => ProtocolEra::Modern2026,
2313        }
2314    }
2315
2316    /// Requests a client sampling completion.
2317    ///
2318    /// In the legacy era this sends `sampling/createMessage` directly. In the
2319    /// final era it returns an `input_required` result whose input descriptor
2320    /// has that exact method and is owned by `owner_cancellation`.
2321    pub async fn sampling_create_message(
2322        &self,
2323        cx: &Cx,
2324        owner_cancellation: McpRequestCancellation,
2325        input_key: impl Into<String>,
2326        params: DualEraSamplingParams,
2327    ) -> McpResult<DualEraServerToClientResult<fastmcp_protocol::CreateMessageResult>> {
2328        let input = match (self, params) {
2329            (Self::Legacy2024 { .. }, DualEraSamplingParams::Legacy2024(params)) => {
2330                MrtrInputRequest::with_params(MrtrInputKind::Sampling, params)?
2331            }
2332            (Self::Modern2026 { .. }, DualEraSamplingParams::Modern2026(params)) => {
2333                MrtrInputRequest::sampling(params)?
2334            }
2335            _ => {
2336                return Err(McpError::invalid_params(
2337                    "Sampling parameters do not match the negotiated protocol era",
2338                ));
2339            }
2340        };
2341        self.dispatch(cx, owner_cancellation, input_key.into(), input)
2342            .await
2343    }
2344
2345    /// Requests client elicitation input.
2346    ///
2347    /// In the legacy era this sends `elicitation/create` directly. In the
2348    /// final era it returns an `input_required` result whose input descriptor
2349    /// has that exact method and is owned by `owner_cancellation`.
2350    pub async fn elicitation_create(
2351        &self,
2352        cx: &Cx,
2353        owner_cancellation: McpRequestCancellation,
2354        input_key: impl Into<String>,
2355        params: DualEraElicitationParams,
2356    ) -> McpResult<DualEraServerToClientResult<fastmcp_protocol::ElicitResult>> {
2357        let input = match (self, params) {
2358            (Self::Legacy2024 { .. }, DualEraElicitationParams::Legacy2024(params)) => {
2359                MrtrInputRequest::legacy_elicitation(params)
2360            }
2361            (Self::Modern2026 { .. }, DualEraElicitationParams::Modern2026(params)) => {
2362                MrtrInputRequest::final_elicitation(params)?
2363            }
2364            _ => return Err(McpError::invalid_params(INVALID_ELICITATION_REQUEST_ERROR)),
2365        };
2366        self.dispatch(cx, owner_cancellation, input_key.into(), input)
2367            .await
2368    }
2369
2370    /// Requests the client's filesystem roots.
2371    ///
2372    /// In the legacy era this sends `roots/list` directly. In the final era
2373    /// it returns an `input_required` result whose input descriptor has that
2374    /// exact method and is owned by `owner_cancellation`.
2375    pub async fn roots_list(
2376        &self,
2377        cx: &Cx,
2378        owner_cancellation: McpRequestCancellation,
2379        input_key: impl Into<String>,
2380    ) -> McpResult<DualEraServerToClientResult<fastmcp_protocol::ListRootsResult>> {
2381        self.dispatch(
2382            cx,
2383            owner_cancellation,
2384            input_key.into(),
2385            MrtrInputRequest::roots(),
2386        )
2387        .await
2388    }
2389
2390    /// Consumes one final-era `inputResponses` retry.
2391    ///
2392    /// This accepts retries only for MCP 2026-07-28. The caller supplies the
2393    /// active [`Cx`] for cancellation/budget authority; the registry retains
2394    /// the request-local cancellation owner that was bound when it issued the
2395    /// corresponding `input_required` result.
2396    pub fn accept_input_retry(
2397        &self,
2398        cx: &Cx,
2399        request_state: &str,
2400        input_responses: MrtrInputResponses,
2401    ) -> McpResult<MrtrRetry> {
2402        if cx.checkpoint().is_err() {
2403            return Err(McpError::request_cancelled());
2404        }
2405
2406        match self {
2407            Self::Legacy2024 { .. } => Err(McpError::invalid_params(LEGACY_INPUT_RETRY_ERROR)),
2408            Self::Modern2026 { exchanges } => exchanges.accept(request_state, input_responses),
2409        }
2410    }
2411
2412    async fn dispatch<T: DeserializeOwned>(
2413        &self,
2414        cx: &Cx,
2415        owner_cancellation: McpRequestCancellation,
2416        input_key: String,
2417        input_request: MrtrInputRequest,
2418    ) -> McpResult<DualEraServerToClientResult<T>> {
2419        if cx.checkpoint().is_err() || owner_cancellation.is_cancel_requested() {
2420            return Err(McpError::request_cancelled());
2421        }
2422
2423        match self {
2424            Self::Legacy2024 { sender } => {
2425                let method = input_request.kind().method();
2426                let params = input_request.into_legacy_params()?;
2427                let response = sender
2428                    .for_request(owner_cancellation)
2429                    .send_request(cx, method, params)
2430                    .await?;
2431                Ok(DualEraServerToClientResult::Legacy(response))
2432            }
2433            Self::Modern2026 { exchanges } => {
2434                let input_requests = MrtrInputRequests::new([(input_key, input_request)])?;
2435                exchanges
2436                    .issue(owner_cancellation, input_requests)
2437                    .map(DualEraServerToClientResult::InputRequired)
2438            }
2439        }
2440    }
2441}
2442
2443// ============================================================================
2444// Tests
2445// ============================================================================
2446
2447#[cfg(test)]
2448mod tests {
2449    use super::*;
2450    use fastmcp_core::block_on;
2451
2452    fn receive_pending(mut receiver: ResponseReceiver) -> PendingResponse {
2453        let cx = Cx::for_testing();
2454        block_on(receiver.recv(&cx)).expect("pending response channel must remain connected")
2455    }
2456
2457    fn mrtr_state_from_wire(result: &MrtrInputRequired) -> String {
2458        serde_json::to_value(result)
2459            .expect("MRTR result must serialize")
2460            .get("requestState")
2461            .and_then(serde_json::Value::as_str)
2462            .expect("MRTR result must contain opaque request state")
2463            .to_owned()
2464    }
2465
2466    fn mrtr_roots_response() -> MrtrInputResponse {
2467        MrtrInputResponse::roots(fastmcp_protocol::ListRootsResult::empty())
2468            .expect("roots response must serialize")
2469    }
2470
2471    fn final_form_elicitation_params() -> fastmcp_protocol::FinalEmbeddedElicitationParams {
2472        serde_json::from_value(serde_json::json!({
2473            "mode": "form",
2474            "message": "Choose a display name",
2475            "requestedSchema": {
2476                "$schema": "https://json-schema.org/draft/2020-12/schema",
2477                "type": "object",
2478                "properties": {"displayName": {"type": "string", "minLength": 1}},
2479                "required": ["displayName"],
2480            },
2481        }))
2482        .expect("final form elicitation parameters must admit")
2483    }
2484
2485    fn final_url_elicitation_params() -> fastmcp_protocol::FinalEmbeddedElicitationParams {
2486        serde_json::from_value(serde_json::json!({
2487            "mode": "url",
2488            "message": "Authorize access",
2489            "url": "https://example.com/authorize",
2490        }))
2491        .expect("final URL elicitation parameters must admit")
2492    }
2493
2494    fn final_sampling_params() -> fastmcp_protocol::FinalEmbeddedCreateMessageParams {
2495        serde_json::from_value(serde_json::json!({
2496            "messages": [{
2497                "role": "assistant",
2498                "content": {
2499                    "type": "tool_use",
2500                    "id": "weather-1",
2501                    "name": "weather",
2502                    "input": {"city": "Boston"},
2503                },
2504            }],
2505            "maxTokens": 16,
2506            "tools": [{
2507                "name": "weather",
2508                "inputSchema": {"type": "object"},
2509            }],
2510            "toolChoice": {"mode": "required"},
2511        }))
2512        .expect("final sampling parameters with tool use must admit")
2513    }
2514
2515    fn final_sampling_result() -> fastmcp_protocol::FinalCreateMessageResult {
2516        serde_json::from_value(serde_json::json!({
2517            "content": {
2518                "type": "tool_use",
2519                "id": "weather-2",
2520                "name": "weather",
2521                "input": {"city": "Cambridge"},
2522            },
2523            "role": "assistant",
2524            "model": "test-model",
2525        }))
2526        .expect("final sampling result with tool use must admit")
2527    }
2528
2529    #[test]
2530    fn mrtr_embeds_exact_input_maps_and_completes_with_bound_responses() {
2531        let registry = MrtrExchangeRegistry::new();
2532        let owner = McpRequestCancellation::new();
2533        let input_requests = MrtrInputRequests::new([
2534            (
2535                "elicit".to_owned(),
2536                MrtrInputRequest::final_elicitation(final_form_elicitation_params())
2537                    .expect("elicitation request must serialize"),
2538            ),
2539            (
2540                "sample".to_owned(),
2541                MrtrInputRequest::sampling(final_sampling_params())
2542                    .expect("sampling request must serialize"),
2543            ),
2544            ("roots".to_owned(), MrtrInputRequest::roots()),
2545        ])
2546        .expect("unique MRTR input map");
2547
2548        let required = registry
2549            .issue(owner.clone(), input_requests)
2550            .expect("MRTR input result must issue");
2551        assert!(
2552            owner.begin_finalization(),
2553            "normal original-response finalization must not cancel MRTR state"
2554        );
2555        let wire = serde_json::to_value(&required).expect("MRTR result must serialize");
2556        assert_eq!(wire["resultType"], "input_required");
2557        assert_eq!(
2558            wire["inputRequests"]["elicit"]["method"],
2559            "elicitation/create"
2560        );
2561        assert_eq!(
2562            wire["inputRequests"]["sample"]["method"],
2563            "sampling/createMessage"
2564        );
2565        assert_eq!(
2566            wire["inputRequests"]["roots"],
2567            serde_json::json!({"method": "roots/list"})
2568        );
2569        for key in ["elicit", "sample", "roots"] {
2570            assert!(wire["inputRequests"][key].get("jsonrpc").is_none());
2571            assert!(wire["inputRequests"][key].get("id").is_none());
2572        }
2573        assert!(
2574            wire["inputRequests"]["sample"]["params"]
2575                .get("_meta")
2576                .is_none()
2577        );
2578        assert_eq!(
2579            wire["inputRequests"]["sample"]["params"]["toolChoice"],
2580            serde_json::json!({"mode": "required"}),
2581            "final sampling tool-choice controls must survive MRTR issuance"
2582        );
2583        assert_eq!(
2584            wire["inputRequests"]["sample"]["params"]["tools"][0]["name"], "weather",
2585            "final sampling tool declarations must survive MRTR issuance"
2586        );
2587
2588        let request_state = mrtr_state_from_wire(&required);
2589        let partial_responses = MrtrInputResponses::new([
2590            (
2591                "elicit".to_owned(),
2592                MrtrInputResponse::elicitation(fastmcp_protocol::ElicitResult::decline())
2593                    .expect("elicitation response must serialize"),
2594            ),
2595            ("inert-unknown-key".to_owned(), mrtr_roots_response()),
2596        ])
2597        .expect("unique MRTR response map");
2598        let response_wire =
2599            serde_json::to_value(&partial_responses).expect("MRTR responses must serialize");
2600        assert_eq!(
2601            response_wire["elicit"],
2602            serde_json::json!({"action": "decline"})
2603        );
2604
2605        let retry = registry
2606            .accept(&request_state, partial_responses)
2607            .expect("partial MRTR response map must reissue only missing inputs");
2608        let MrtrRetry::InputRequired(retry) = retry else {
2609            panic!("partial MRTR response map must not complete the exchange");
2610        };
2611        let retry_wire = serde_json::to_value(&retry).expect("retry result must serialize");
2612        assert!(retry_wire["inputRequests"].get("elicit").is_none());
2613        assert_eq!(
2614            retry_wire["inputRequests"]["sample"]["method"],
2615            "sampling/createMessage"
2616        );
2617        assert_eq!(
2618            retry_wire["inputRequests"]["roots"],
2619            serde_json::json!({"method": "roots/list"})
2620        );
2621        let retry_state = mrtr_state_from_wire(&retry);
2622        assert_ne!(
2623            retry_state, request_state,
2624            "partial retry needs fresh state"
2625        );
2626        let old_state_error = registry
2627            .accept(&request_state, MrtrInputResponses::default())
2628            .expect_err("the predecessor state must not replay after partial acceptance");
2629        assert_eq!(old_state_error.code, McpErrorCode::InvalidParams);
2630
2631        let complete = registry
2632            .accept(
2633                &retry_state,
2634                MrtrInputResponses::new([
2635                    (
2636                        "sample".to_owned(),
2637                        MrtrInputResponse::sampling(final_sampling_result())
2638                            .expect("sampling response must serialize"),
2639                    ),
2640                    ("roots".to_owned(), mrtr_roots_response()),
2641                ])
2642                .expect("unique MRTR response map"),
2643            )
2644            .expect("matching remaining MRTR responses must complete");
2645        let MrtrRetry::Complete(complete) = complete else {
2646            panic!("all matching MRTR responses must complete the exchange");
2647        };
2648        assert_eq!(complete.responses().len(), 3);
2649        assert_eq!(
2650            complete
2651                .responses()
2652                .get("elicit")
2653                .map(MrtrInputResponse::kind),
2654            Some(MrtrInputKind::Elicitation)
2655        );
2656        assert_eq!(
2657            complete
2658                .responses()
2659                .get("sample")
2660                .map(MrtrInputResponse::kind),
2661            Some(MrtrInputKind::Sampling)
2662        );
2663        assert!(matches!(
2664            complete.sampling("sample"),
2665            Ok(Some(fastmcp_protocol::FinalCreateMessageResult {
2666                content: fastmcp_protocol::FinalSamplingMessageContent::Block(
2667                    fastmcp_protocol::FinalSamplingMessageContentBlock::ToolUse { .. }
2668                ),
2669                ..
2670            }))
2671        ));
2672        assert_eq!(
2673            complete
2674                .responses()
2675                .get("roots")
2676                .map(MrtrInputResponse::kind),
2677            Some(MrtrInputKind::Roots)
2678        );
2679        assert_eq!(registry.active_len(), 0);
2680    }
2681
2682    #[test]
2683    fn mrtr_rejects_cross_kind_before_consumption_and_rejects_replay() {
2684        let registry = MrtrExchangeRegistry::new();
2685        let input_requests =
2686            MrtrInputRequests::new([("roots".to_owned(), MrtrInputRequest::roots())])
2687                .expect("unique MRTR input map");
2688        let required = registry
2689            .issue(McpRequestCancellation::new(), input_requests)
2690            .expect("MRTR input result must issue");
2691        let request_state = mrtr_state_from_wire(&required);
2692
2693        let wrong_kind = MrtrInputResponses::new([(
2694            "roots".to_owned(),
2695            MrtrInputResponse::sampling(final_sampling_result())
2696                .expect("sampling response must serialize"),
2697        )])
2698        .expect("unique MRTR response map");
2699        let error = registry
2700            .accept(&request_state, wrong_kind)
2701            .expect_err("a sampling value cannot fulfill a roots request");
2702        assert_eq!(error.code, McpErrorCode::InvalidParams);
2703        assert_eq!(
2704            registry.active_len(),
2705            1,
2706            "wrong-kind input must not consume state"
2707        );
2708
2709        let matching = MrtrInputResponses::new([("roots".to_owned(), mrtr_roots_response())])
2710            .expect("unique MRTR response map");
2711        assert!(matches!(
2712            registry.accept(&request_state, matching),
2713            Ok(MrtrRetry::Complete(_))
2714        ));
2715        assert_eq!(registry.active_len(), 0);
2716
2717        let replay = registry
2718            .accept(
2719                &request_state,
2720                MrtrInputResponses::new([("roots".to_owned(), mrtr_roots_response())])
2721                    .expect("unique MRTR response map"),
2722            )
2723            .expect_err("a consumed MRTR request state must not replay");
2724        assert_eq!(replay.code, McpErrorCode::InvalidParams);
2725        assert_eq!(registry.active_len(), 0, "replay must not restore state");
2726    }
2727
2728    #[test]
2729    fn mrtr_typed_unknown_only_retry_preserves_the_original_continuation() {
2730        let registry = MrtrExchangeRegistry::new();
2731        let required = registry
2732            .issue(
2733                McpRequestCancellation::new(),
2734                MrtrInputRequests::new([("roots".to_owned(), MrtrInputRequest::roots())])
2735                    .expect("unique MRTR input map"),
2736            )
2737            .expect("MRTR input result must issue");
2738        let request_state = mrtr_state_from_wire(&required);
2739
2740        let unknown_only = MrtrInputResponses::new([("inert".to_owned(), mrtr_roots_response())])
2741            .expect("typed response map permits an inert key");
2742        let error = registry
2743            .accept(&request_state, unknown_only)
2744            .expect_err("unknown-only typed input must not rotate a continuation");
2745        assert_eq!(error.code, McpErrorCode::InvalidParams);
2746        assert_eq!(
2747            registry.active_len(),
2748            1,
2749            "the rejected unknown-only retry must retain the original state"
2750        );
2751
2752        let matching = MrtrInputResponses::new([("roots".to_owned(), mrtr_roots_response())])
2753            .expect("unique matching response map");
2754        assert!(matches!(
2755            registry.accept(&request_state, matching),
2756            Ok(MrtrRetry::Complete(_))
2757        ));
2758        assert_eq!(registry.active_len(), 0);
2759    }
2760
2761    #[test]
2762    fn mrtr_accept_wire_decodes_the_issued_kind_before_consuming_state() {
2763        let registry = MrtrExchangeRegistry::new();
2764        let required = registry
2765            .issue(
2766                McpRequestCancellation::new(),
2767                MrtrInputRequests::new([("roots".to_owned(), MrtrInputRequest::roots())])
2768                    .expect("unique MRTR input map"),
2769            )
2770            .expect("MRTR input result must issue");
2771        let request_state = mrtr_state_from_wire(&required);
2772
2773        let wrong_kind = BTreeMap::from([(
2774            "roots".to_owned(),
2775            serde_json::to_value(
2776                MrtrInputResponse::sampling(final_sampling_result())
2777                    .expect("sampling response must serialize"),
2778            )
2779            .expect("sampling response must convert to a wire value"),
2780        )]);
2781        let error = registry
2782            .accept_wire(&request_state, &wrong_kind)
2783            .expect_err("a sampling wire value cannot fulfill a roots request");
2784        assert_eq!(error.code, McpErrorCode::InvalidParams);
2785        assert_eq!(
2786            registry.active_len(),
2787            1,
2788            "wrong-kind wire input must not consume state"
2789        );
2790
2791        let matching = BTreeMap::from([(
2792            "roots".to_owned(),
2793            serde_json::to_value(mrtr_roots_response())
2794                .expect("roots response must convert to a wire value"),
2795        )]);
2796        assert!(matches!(
2797            registry.accept_wire(&request_state, &matching),
2798            Ok(MrtrRetry::Complete(_))
2799        ));
2800        assert_eq!(registry.active_len(), 0);
2801
2802        let replay = registry
2803            .accept_wire(&request_state, &matching)
2804            .expect_err("a consumed wire request state must not replay");
2805        assert_eq!(replay.code, McpErrorCode::InvalidParams);
2806    }
2807
2808    #[test]
2809    fn mrtr_accepts_ordered_final_input_responses_without_map_normalization() {
2810        let registry = MrtrExchangeRegistry::new();
2811        let binding = MrtrExchangeBinding::new(
2812            "tools/call",
2813            "ordered-tool".to_owned(),
2814            [3; 32],
2815            [4; 32],
2816            None,
2817        );
2818        let required = registry
2819            .issue_bound(
2820                McpRequestCancellation::new(),
2821                binding.clone(),
2822                MrtrInputRequests::new([
2823                    ("second".to_owned(), MrtrInputRequest::roots()),
2824                    ("first".to_owned(), MrtrInputRequest::roots()),
2825                ])
2826                .expect("unique MRTR input keys"),
2827            )
2828            .expect("MRTR input result must issue");
2829        let request_state = mrtr_state_from_wire(&required);
2830        let responses: FinalInputResponses =
2831            serde_json::from_str(r#"{"second":{"roots":[]},"first":{"roots":[]}}"#)
2832                .expect("ordered final responses decode");
2833        assert_eq!(
2834            responses
2835                .entries()
2836                .iter()
2837                .map(|(key, _)| key.as_str())
2838                .collect::<Vec<_>>(),
2839            vec!["second", "first"],
2840            "the registry receives the protocol decoder's order-preserving representation"
2841        );
2842        let completed = registry
2843            .accept_final_input_responses_bound(&request_state, &binding, &responses)
2844            .expect("ordered protocol responses complete the exchange");
2845        let MrtrRetry::Complete(completed) = completed else {
2846            panic!("every expected ordered response completes the exchange");
2847        };
2848        assert_eq!(
2849            completed
2850                .responses()
2851                .iter()
2852                .map(|(key, _)| key)
2853                .collect::<Vec<_>>(),
2854            vec!["second", "first"],
2855            "handler delivery preserves the accepted wire order rather than BTreeMap key order"
2856        );
2857        assert_eq!(registry.active_len(), 0);
2858    }
2859
2860    #[test]
2861    fn mrtr_state_only_retry_requires_absent_input_responses() {
2862        let registry = MrtrExchangeRegistry::new();
2863        let binding = MrtrExchangeBinding::new(
2864            "tools/call",
2865            "state-only-tool".to_owned(),
2866            [7; 32],
2867            [9; 32],
2868            None,
2869        );
2870        let required = registry
2871            .issue_bound(
2872                McpRequestCancellation::new(),
2873                binding.clone(),
2874                MrtrInputRequests::default(),
2875            )
2876            .expect("a state-only exchange issues");
2877        let wire = serde_json::to_value(&required).expect("state-only exchange serializes");
2878        assert!(
2879            wire.get("inputRequests").is_none(),
2880            "state-only input_required omits inputRequests"
2881        );
2882        let request_state = mrtr_state_from_wire(&required);
2883
2884        let explicit_empty = registry
2885            .accept_wire_bound(&request_state, &binding, &BTreeMap::new())
2886            .expect_err("an explicit empty inputResponses map is not state-only");
2887        assert_eq!(explicit_empty.code, McpErrorCode::InvalidParams);
2888        assert_eq!(
2889            registry.active_len(),
2890            1,
2891            "the rejected explicit map leaves the state-only exchange available"
2892        );
2893
2894        let completed = registry
2895            .accept_state_only_bound(&request_state, &binding)
2896            .expect("an absent inputResponses member completes the state-only exchange");
2897        let MrtrRetry::Complete(inputs) = completed else {
2898            panic!("state-only retry must complete without manufacturing inputs");
2899        };
2900        assert!(inputs.responses().is_empty());
2901        assert_eq!(registry.active_len(), 0);
2902    }
2903
2904    #[test]
2905    fn mrtr_expiry_and_owning_request_cancellation_prevent_resolution() {
2906        let registry = MrtrExchangeRegistry::with_limits(
2907            16,
2908            DEFAULT_MAX_MRTR_ROUNDS,
2909            DEFAULT_MAX_MRTR_INPUT_REQUESTS_PER_ROUND,
2910            DEFAULT_MAX_MRTR_INPUT_REQUESTS_TOTAL,
2911            Duration::from_millis(1),
2912        )
2913        .expect("bounded MRTR registry");
2914        let input_requests = || {
2915            MrtrInputRequests::new([("roots".to_owned(), MrtrInputRequest::roots())])
2916                .expect("unique MRTR input map")
2917        };
2918
2919        let expired = registry
2920            .issue(McpRequestCancellation::new(), input_requests())
2921            .expect("MRTR input result must issue");
2922        let expired_state = mrtr_state_from_wire(&expired);
2923        let expiry_error = registry
2924            .accept_at(
2925                &expired_state,
2926                None,
2927                MrtrInputResponses::new([("roots".to_owned(), mrtr_roots_response())])
2928                    .expect("unique MRTR response map"),
2929                false,
2930                Instant::now() + Duration::from_millis(1),
2931            )
2932            .expect_err("expired state must fail before resolution");
2933        assert_eq!(expiry_error.code, McpErrorCode::InvalidParams);
2934        assert_eq!(registry.active_len(), 0, "expired state must be removed");
2935
2936        let owner = McpRequestCancellation::new();
2937        let cancelled = registry
2938            .issue(owner.clone(), input_requests())
2939            .expect("MRTR input result must issue");
2940        let cancelled_state = mrtr_state_from_wire(&cancelled);
2941        assert!(owner.cancel());
2942        let cancellation_error = registry
2943            .accept(
2944                &cancelled_state,
2945                MrtrInputResponses::new([("roots".to_owned(), mrtr_roots_response())])
2946                    .expect("unique MRTR response map"),
2947            )
2948            .expect_err("owner cancellation must win before MRTR resolution");
2949        assert_eq!(cancellation_error.code, McpErrorCode::RequestCancelled);
2950        assert_eq!(registry.active_len(), 0, "cancelled state must be removed");
2951    }
2952
2953    fn dual_era_boundary_with_recording_sender(
2954        era: ProtocolEra,
2955        sent_methods: Arc<Mutex<Vec<String>>>,
2956    ) -> DualEraServerToClient {
2957        let pending = Arc::new(PendingRequests::new());
2958        let pending_for_send = Arc::clone(&pending);
2959        let sent_methods_for_send = Arc::clone(&sent_methods);
2960        let send_fn: TransportSendFn = Arc::new(move |message| {
2961            let JsonRpcMessage::Request(request) = message else {
2962                panic!("the server-to-client boundary may only emit requests");
2963            };
2964            sent_methods_for_send
2965                .lock()
2966                .unwrap_or_else(std::sync::PoisonError::into_inner)
2967                .push(request.method.clone());
2968
2969            let result = match request.method.as_str() {
2970                "sampling/createMessage" => serde_json::json!({
2971                    "content": {"type": "text", "text": "legacy completion"},
2972                    "role": "assistant",
2973                    "model": "legacy-model",
2974                    "stopReason": "endTurn"
2975                }),
2976                "elicitation/create" => serde_json::json!({"action": "decline"}),
2977                "roots/list" => serde_json::json!({"roots": []}),
2978                method => panic!("unexpected reverse JSON-RPC method: {method}"),
2979            };
2980            let id = request
2981                .id
2982                .clone()
2983                .expect("server-to-client requests require an ID");
2984            assert!(
2985                pending_for_send.route_response(&JsonRpcResponse::success(id, result)),
2986                "recorded reverse request must retain its response waiter"
2987            );
2988            Ok(())
2989        });
2990
2991        DualEraServerToClient::new(
2992            era,
2993            RequestSender::new(pending, send_fn),
2994            Arc::new(MrtrExchangeRegistry::new()),
2995        )
2996    }
2997
2998    #[test]
2999    fn dual_era_boundary_legacy_round_trips_the_three_exact_reverse_methods() {
3000        let sent_methods = Arc::new(Mutex::new(Vec::new()));
3001        let boundary = dual_era_boundary_with_recording_sender(
3002            ProtocolEra::Legacy2024,
3003            Arc::clone(&sent_methods),
3004        );
3005        let cx = Cx::for_testing();
3006
3007        let sampling = block_on(boundary.sampling_create_message(
3008            &cx,
3009            McpRequestCancellation::new(),
3010            "sample",
3011            DualEraSamplingParams::Legacy2024(fastmcp_protocol::CreateMessageParams::new(
3012                Vec::new(),
3013                fastmcp_protocol::JsonInteger::from(16_i64),
3014            )),
3015        ))
3016        .expect("legacy sampling must await a direct response");
3017        let DualEraServerToClientResult::Legacy(sampling) = sampling else {
3018            panic!("legacy sampling must not create an MRTR retry");
3019        };
3020        assert_eq!(sampling.model, "legacy-model");
3021
3022        let elicitation = block_on(boundary.elicitation_create(
3023            &cx,
3024            McpRequestCancellation::new(),
3025            "elicit",
3026            DualEraElicitationParams::Legacy2024(fastmcp_protocol::ElicitRequestParams::form(
3027                "Continue?",
3028                serde_json::json!({"type": "object"}),
3029            )),
3030        ))
3031        .expect("legacy elicitation must await a direct response");
3032        assert!(matches!(
3033            elicitation,
3034            DualEraServerToClientResult::Legacy(fastmcp_protocol::ElicitResult {
3035                action: fastmcp_protocol::ElicitAction::Decline,
3036                ..
3037            })
3038        ));
3039
3040        let roots = block_on(boundary.roots_list(&cx, McpRequestCancellation::new(), "roots"))
3041            .expect("legacy roots must await a direct response");
3042        let DualEraServerToClientResult::Legacy(roots) = roots else {
3043            panic!("legacy roots must not create an MRTR retry");
3044        };
3045        assert!(roots.roots.is_empty());
3046
3047        assert_eq!(boundary.era(), ProtocolEra::Legacy2024);
3048        assert_eq!(
3049            *sent_methods
3050                .lock()
3051                .unwrap_or_else(std::sync::PoisonError::into_inner),
3052            vec![
3053                "sampling/createMessage".to_owned(),
3054                "elicitation/create".to_owned(),
3055                "roots/list".to_owned(),
3056            ]
3057        );
3058    }
3059
3060    #[test]
3061    fn dual_era_elicitation_keeps_legacy_url_identity_out_of_final_mrtr() {
3062        let pending = Arc::new(PendingRequests::new());
3063        let pending_for_send = Arc::clone(&pending);
3064        let sent_params = Arc::new(Mutex::new(Vec::new()));
3065        let sent_params_for_send = Arc::clone(&sent_params);
3066        let send_fn: TransportSendFn = Arc::new(move |message| {
3067            let JsonRpcMessage::Request(request) = message else {
3068                panic!("legacy isolation boundary may only emit requests");
3069            };
3070            assert_eq!(request.method, "elicitation/create");
3071            sent_params_for_send
3072                .lock()
3073                .unwrap_or_else(std::sync::PoisonError::into_inner)
3074                .push(
3075                    request
3076                        .params
3077                        .clone()
3078                        .expect("legacy elicitation carries params"),
3079                );
3080            let id = request
3081                .id
3082                .clone()
3083                .expect("legacy elicitation requires a JSON-RPC id");
3084            assert!(pending_for_send.route_response(&JsonRpcResponse::success(
3085                id,
3086                serde_json::json!({"action": "decline"}),
3087            )));
3088            Ok(())
3089        });
3090        let legacy = DualEraServerToClient::new(
3091            ProtocolEra::Legacy2024,
3092            RequestSender::new(pending, send_fn),
3093            Arc::new(MrtrExchangeRegistry::new()),
3094        );
3095        let cx = Cx::for_testing();
3096
3097        block_on(legacy.elicitation_create(
3098            &cx,
3099            McpRequestCancellation::new(),
3100            "legacy-url",
3101            DualEraElicitationParams::Legacy2024(fastmcp_protocol::ElicitRequestParams::url(
3102                "Authorize legacy access",
3103                "https://example.com/legacy-authorize",
3104                "legacy-elicitation-id",
3105            )),
3106        ))
3107        .expect("exact-2024 URL elicitation must retain its identity");
3108        assert_eq!(
3109            sent_params
3110                .lock()
3111                .unwrap_or_else(std::sync::PoisonError::into_inner)
3112                .as_slice(),
3113            &[serde_json::json!({
3114                "mode": "url",
3115                "message": "Authorize legacy access",
3116                "url": "https://example.com/legacy-authorize",
3117                "elicitationId": "legacy-elicitation-id",
3118            })],
3119        );
3120
3121        let legacy_error = block_on(legacy.elicitation_create(
3122            &cx,
3123            McpRequestCancellation::new(),
3124            "final-on-legacy",
3125            DualEraElicitationParams::Modern2026(final_url_elicitation_params()),
3126        ))
3127        .expect_err("a final descriptor must not cross into exact-2024 JSON-RPC");
3128        assert_eq!(legacy_error.code, McpErrorCode::InvalidParams);
3129        assert_eq!(
3130            sent_params
3131                .lock()
3132                .unwrap_or_else(std::sync::PoisonError::into_inner)
3133                .len(),
3134            1,
3135            "rejected final input must not reach the legacy sender",
3136        );
3137
3138        let modern_registry = Arc::new(MrtrExchangeRegistry::new());
3139        let modern = DualEraServerToClient::Modern2026 {
3140            exchanges: Arc::clone(&modern_registry),
3141        };
3142        let modern_error = block_on(modern.elicitation_create(
3143            &cx,
3144            McpRequestCancellation::new(),
3145            "legacy-on-final",
3146            DualEraElicitationParams::Legacy2024(fastmcp_protocol::ElicitRequestParams::url(
3147                "Authorize legacy access",
3148                "https://example.com/legacy-authorize",
3149                "legacy-elicitation-id",
3150            )),
3151        ))
3152        .expect_err("a legacy descriptor must not mint final MRTR state");
3153        assert_eq!(modern_error.code, McpErrorCode::InvalidParams);
3154        assert_eq!(modern_registry.active_len(), 0);
3155    }
3156
3157    fn assert_modern_elicitation_issues_and_reconstructs(
3158        input_key: &str,
3159        params: fastmcp_protocol::FinalEmbeddedElicitationParams,
3160        expected_params: serde_json::Value,
3161    ) {
3162        let registry = Arc::new(MrtrExchangeRegistry::new());
3163        let boundary = DualEraServerToClient::Modern2026 {
3164            exchanges: Arc::clone(&registry),
3165        };
3166        let cx = Cx::for_testing();
3167        let result = block_on(boundary.elicitation_create(
3168            &cx,
3169            McpRequestCancellation::new(),
3170            input_key,
3171            DualEraElicitationParams::Modern2026(params),
3172        ))
3173        .expect("final elicitation must issue MRTR state");
3174        let DualEraServerToClientResult::InputRequired(required) = result else {
3175            panic!("modern elicitation must issue a final input_required result");
3176        };
3177        assert_eq!(registry.active_len(), 1);
3178        let wire = serde_json::to_value(&required).expect("final input request serializes");
3179        let descriptor = wire["inputRequests"][input_key].clone();
3180        assert_eq!(descriptor["method"], "elicitation/create");
3181        assert_eq!(descriptor["params"], expected_params);
3182
3183        let reconstructed = MrtrInputRequest::from_wire(&descriptor)
3184            .expect("final emitted descriptor must reconstruct without legacy conversion");
3185        assert_eq!(
3186            serde_json::to_value(reconstructed).expect("reconstructed descriptor serializes"),
3187            descriptor,
3188        );
3189
3190        let complete = boundary
3191            .accept_input_retry(
3192                &cx,
3193                &mrtr_state_from_wire(&required),
3194                MrtrInputResponses::new([(
3195                    input_key.to_owned(),
3196                    MrtrInputResponse::elicitation(fastmcp_protocol::ElicitResult::decline())
3197                        .expect("elicitation response must serialize"),
3198                )])
3199                .expect("one matching final elicitation response"),
3200            )
3201            .expect("final elicitation retry must complete");
3202        let MrtrRetry::Complete(complete) = complete else {
3203            panic!("one matching elicitation response must complete the exchange");
3204        };
3205        assert_eq!(
3206            complete
3207                .responses()
3208                .get(input_key)
3209                .map(MrtrInputResponse::kind),
3210            Some(MrtrInputKind::Elicitation),
3211        );
3212    }
3213
3214    #[test]
3215    fn dual_era_boundary_modern_issues_and_reconstructs_final_form_elicitation() {
3216        assert_modern_elicitation_issues_and_reconstructs(
3217            "form",
3218            final_form_elicitation_params(),
3219            serde_json::json!({
3220                "mode": "form",
3221                "message": "Choose a display name",
3222                "requestedSchema": {
3223                    "$schema": "https://json-schema.org/draft/2020-12/schema",
3224                    "type": "object",
3225                    "properties": {"displayName": {"type": "string", "minLength": 1}},
3226                    "required": ["displayName"],
3227                },
3228            }),
3229        );
3230    }
3231
3232    #[test]
3233    fn dual_era_boundary_modern_issues_and_reconstructs_final_url_elicitation() {
3234        assert_modern_elicitation_issues_and_reconstructs(
3235            "url",
3236            final_url_elicitation_params(),
3237            serde_json::json!({
3238                "mode": "url",
3239                "message": "Authorize access",
3240                "url": "https://example.com/authorize",
3241            }),
3242        );
3243    }
3244
3245    fn legacy_sampling_boundary_with_result(
3246        sampling_result: serde_json::Value,
3247    ) -> DualEraServerToClient {
3248        let pending = Arc::new(PendingRequests::new());
3249        let pending_for_send = Arc::clone(&pending);
3250        let send_fn: TransportSendFn = Arc::new(move |message| {
3251            let JsonRpcMessage::Request(request) = message else {
3252                panic!("the legacy sampling boundary may only emit requests");
3253            };
3254            assert_eq!(request.method, "sampling/createMessage");
3255            let id = request
3256                .id
3257                .clone()
3258                .expect("legacy reverse sampling must retain its JSON-RPC ID");
3259            assert!(
3260                pending_for_send
3261                    .route_response(&JsonRpcResponse::success(id, sampling_result.clone())),
3262                "legacy sampling response must reach its registered waiter"
3263            );
3264            Ok(())
3265        });
3266        DualEraServerToClient::new(
3267            ProtocolEra::Legacy2024,
3268            RequestSender::new(pending, send_fn),
3269            Arc::new(MrtrExchangeRegistry::new()),
3270        )
3271    }
3272
3273    #[test]
3274    fn dual_era_legacy_sampling_round_trips_an_absent_stop_reason() {
3275        let expected = serde_json::json!({
3276            "content": {"type": "text", "text": "legacy completion"},
3277            "role": "assistant",
3278            "model": "legacy-model"
3279        });
3280        let boundary = legacy_sampling_boundary_with_result(expected.clone());
3281        let cx = Cx::for_testing();
3282
3283        let result = block_on(boundary.sampling_create_message(
3284            &cx,
3285            McpRequestCancellation::new(),
3286            "sample",
3287            DualEraSamplingParams::Legacy2024(fastmcp_protocol::CreateMessageParams::new(
3288                Vec::new(),
3289                fastmcp_protocol::JsonInteger::from(16_i64),
3290            )),
3291        ))
3292        .expect("legacy sampling must preserve an absent stopReason");
3293        let DualEraServerToClientResult::Legacy(result) = result else {
3294            panic!("legacy sampling must not create an MRTR retry");
3295        };
3296
3297        assert_eq!(result.stop_reason, None);
3298        assert_eq!(
3299            serde_json::to_value(result).expect("legacy sampling result must re-encode"),
3300            expected
3301        );
3302    }
3303
3304    #[test]
3305    fn dual_era_legacy_sampling_round_trips_an_open_provider_stop_reason() {
3306        let expected = serde_json::json!({
3307            "content": {"type": "text", "text": "legacy completion"},
3308            "role": "assistant",
3309            "model": "legacy-model",
3310            "stopReason": "provider_safety_limit"
3311        });
3312        let boundary = legacy_sampling_boundary_with_result(expected.clone());
3313        let cx = Cx::for_testing();
3314
3315        let result = block_on(boundary.sampling_create_message(
3316            &cx,
3317            McpRequestCancellation::new(),
3318            "sample",
3319            DualEraSamplingParams::Legacy2024(fastmcp_protocol::CreateMessageParams::new(
3320                Vec::new(),
3321                fastmcp_protocol::JsonInteger::from(16_i64),
3322            )),
3323        ))
3324        .expect("legacy sampling must preserve an open provider stopReason");
3325        let DualEraServerToClientResult::Legacy(result) = result else {
3326            panic!("legacy sampling must not create an MRTR retry");
3327        };
3328
3329        assert_eq!(result.stop_reason.as_deref(), Some("provider_safety_limit"));
3330        assert_eq!(
3331            serde_json::to_value(result).expect("legacy sampling result must re-encode"),
3332            expected
3333        );
3334    }
3335
3336    #[test]
3337    fn dual_era_boundary_modern_uses_input_required_retry_flow() {
3338        let sent_methods = Arc::new(Mutex::new(Vec::new()));
3339        let boundary = dual_era_boundary_with_recording_sender(
3340            ProtocolEra::Modern2026,
3341            Arc::clone(&sent_methods),
3342        );
3343        let cx = Cx::for_testing();
3344
3345        let sampling = block_on(boundary.sampling_create_message(
3346            &cx,
3347            McpRequestCancellation::new(),
3348            "sample",
3349            DualEraSamplingParams::Modern2026(final_sampling_params()),
3350        ))
3351        .expect("modern sampling must create an MRTR input result");
3352        let DualEraServerToClientResult::InputRequired(required) = sampling else {
3353            panic!("modern sampling must not send and await reverse JSON-RPC");
3354        };
3355        let wire = serde_json::to_value(&required).expect("MRTR result must serialize");
3356        assert_eq!(wire["resultType"], "input_required");
3357        assert_eq!(
3358            wire["inputRequests"]["sample"]["method"],
3359            "sampling/createMessage"
3360        );
3361        assert!(wire["inputRequests"]["sample"].get("jsonrpc").is_none());
3362        assert!(wire["inputRequests"]["sample"].get("id").is_none());
3363
3364        let complete = boundary
3365            .accept_input_retry(
3366                &cx,
3367                &mrtr_state_from_wire(&required),
3368                MrtrInputResponses::new([(
3369                    "sample".to_owned(),
3370                    MrtrInputResponse::sampling(final_sampling_result())
3371                        .expect("sampling response must serialize"),
3372                )])
3373                .expect("one matching final response"),
3374            )
3375            .expect("modern retry must resolve through the MRTR registry");
3376        let MrtrRetry::Complete(complete) = complete else {
3377            panic!("a matching response must complete the one-input exchange");
3378        };
3379        assert_eq!(complete.responses().len(), 1);
3380        assert_eq!(
3381            complete
3382                .responses()
3383                .get("sample")
3384                .map(MrtrInputResponse::kind),
3385            Some(MrtrInputKind::Sampling)
3386        );
3387        assert_eq!(boundary.era(), ProtocolEra::Modern2026);
3388        assert!(
3389            sent_methods
3390                .lock()
3391                .unwrap_or_else(std::sync::PoisonError::into_inner)
3392                .is_empty(),
3393            "modern input_required must not emit a reverse JSON-RPC request"
3394        );
3395    }
3396
3397    #[test]
3398    fn dual_era_boundary_modern_roots_never_sends_reverse_jsonrpc() {
3399        let sent_methods = Arc::new(Mutex::new(Vec::new()));
3400        let boundary = dual_era_boundary_with_recording_sender(
3401            ProtocolEra::Modern2026,
3402            Arc::clone(&sent_methods),
3403        );
3404        let cx = Cx::for_testing();
3405
3406        let roots = block_on(boundary.roots_list(&cx, McpRequestCancellation::new(), "roots"))
3407            .expect("modern roots must create an MRTR input result");
3408        let DualEraServerToClientResult::InputRequired(required) = roots else {
3409            panic!("changing only the selected era must disable reverse roots/list");
3410        };
3411        let wire = serde_json::to_value(required).expect("MRTR result must serialize");
3412        assert_eq!(
3413            wire["inputRequests"]["roots"],
3414            serde_json::json!({"method": "roots/list"})
3415        );
3416        assert!(
3417            sent_methods
3418                .lock()
3419                .unwrap_or_else(std::sync::PoisonError::into_inner)
3420                .is_empty(),
3421            "MCP 2026-07-28 must not send roots/list as reverse JSON-RPC"
3422        );
3423    }
3424
3425    #[test]
3426    fn test_pending_requests_register_and_route() {
3427        let pending = PendingRequests::new();
3428
3429        // Register a request
3430        let (id, receiver) = pending.register().unwrap();
3431
3432        // Simulate a response
3433        let response = JsonRpcResponse::success(id, serde_json::json!({"result": "ok"}));
3434        assert!(pending.route_response(&response));
3435
3436        // Receive the response
3437        let result = receive_pending(receiver);
3438        assert!(result.is_ok());
3439        assert_eq!(result.unwrap(), serde_json::json!({"result": "ok"}));
3440    }
3441
3442    #[test]
3443    fn test_pending_requests_error_response() {
3444        let pending = PendingRequests::new();
3445
3446        let (id, receiver) = pending.register().unwrap();
3447
3448        // Simulate an error response
3449        let response = JsonRpcResponse::error(
3450            Some(id),
3451            JsonRpcError {
3452                code: (-32600).into(),
3453                message: "Invalid request".to_string(),
3454                data: None,
3455            },
3456        );
3457        assert!(pending.route_response(&response));
3458
3459        // Receive the error
3460        let result = receive_pending(receiver);
3461        assert!(result.is_err());
3462        assert_eq!(result.unwrap_err().message, REMOTE_RESPONSE_ERROR);
3463    }
3464
3465    #[test]
3466    fn test_pending_requests_cancel_all() {
3467        let pending = PendingRequests::new();
3468
3469        let (_, receiver1) = pending.register().unwrap();
3470        let (_, receiver2) = pending.register().unwrap();
3471
3472        // Cancel all
3473        pending.cancel_all();
3474
3475        // Both should receive errors
3476        let result1 = receive_pending(receiver1);
3477        let result2 = receive_pending(receiver2);
3478        assert!(result1.is_err());
3479        assert!(result2.is_err());
3480    }
3481
3482    #[test]
3483    fn request_local_cancellation_wakes_a_pending_bidirectional_wait() {
3484        use std::sync::atomic::{AtomicBool, Ordering};
3485
3486        struct WakeFlag(AtomicBool);
3487
3488        impl std::task::Wake for WakeFlag {
3489            fn wake(self: Arc<Self>) {
3490                self.0.store(true, Ordering::Release);
3491            }
3492        }
3493
3494        let pending = Arc::new(PendingRequests::new());
3495        let sent = Arc::new(AtomicBool::new(false));
3496        let sent_flag = Arc::clone(&sent);
3497        let outbound = Arc::new(Mutex::new(Vec::new()));
3498        let outbound_for_send = Arc::clone(&outbound);
3499        let sender = RequestSender::new(
3500            Arc::clone(&pending),
3501            Arc::new(move |message| {
3502                sent_flag.store(true, Ordering::Release);
3503                outbound_for_send
3504                    .lock()
3505                    .expect("test outbound mutex must not be poisoned")
3506                    .push(message.clone());
3507                Ok(())
3508            }),
3509        );
3510        let cancellation = McpRequestCancellation::new();
3511        let scoped = sender.for_request(cancellation.clone());
3512        let cx = Cx::for_testing();
3513        let mut future = Box::pin(scoped.send_request::<serde_json::Value>(
3514            &cx,
3515            "test/request-local-cancellation",
3516            serde_json::json!({}),
3517        ));
3518        let wake_flag = Arc::new(WakeFlag(AtomicBool::new(false)));
3519        let waker = std::task::Waker::from(Arc::clone(&wake_flag));
3520        let mut task_cx = std::task::Context::from_waker(&waker);
3521
3522        assert!(std::future::Future::poll(future.as_mut(), &mut task_cx).is_pending());
3523        assert!(sent.load(Ordering::Acquire));
3524        assert_eq!(pending.in_flight_len(), 1);
3525
3526        assert!(cancellation.cancel());
3527        assert!(wake_flag.0.load(Ordering::Acquire));
3528        let error = block_on(future).unwrap_err();
3529        assert_eq!(error.code, McpErrorCode::RequestCancelled);
3530        assert_eq!(pending.in_flight_len(), 0);
3531        let outbound = outbound
3532            .lock()
3533            .expect("test outbound mutex must not be poisoned");
3534        assert_eq!(outbound.len(), 2);
3535        let JsonRpcMessage::Request(cancelled) = &outbound[1] else {
3536            panic!("cancelled reverse request must notify the peer");
3537        };
3538        assert_eq!(cancelled.method, "notifications/cancelled");
3539        assert_eq!(
3540            cancelled.params,
3541            Some(serde_json::json!({ "requestId": FIRST_SERVER_REQUEST_ID }))
3542        );
3543    }
3544
3545    #[test]
3546    fn request_finalization_wakes_and_removes_a_pending_bidirectional_wait() {
3547        use std::sync::atomic::{AtomicBool, Ordering};
3548
3549        struct WakeFlag(AtomicBool);
3550
3551        impl std::task::Wake for WakeFlag {
3552            fn wake(self: Arc<Self>) {
3553                self.0.store(true, Ordering::Release);
3554            }
3555        }
3556
3557        let pending = Arc::new(PendingRequests::new());
3558        let sender = RequestSender::new(Arc::clone(&pending), Arc::new(|_| Ok(())));
3559        let cancellation = McpRequestCancellation::new();
3560        let scoped = sender.for_request(cancellation.clone());
3561        let cx = Cx::for_testing();
3562        let mut future = Box::pin(scoped.send_request::<serde_json::Value>(
3563            &cx,
3564            "test/request-finalization",
3565            serde_json::json!({}),
3566        ));
3567        let wake_flag = Arc::new(WakeFlag(AtomicBool::new(false)));
3568        let waker = std::task::Waker::from(Arc::clone(&wake_flag));
3569        let mut task_cx = std::task::Context::from_waker(&waker);
3570
3571        assert!(std::future::Future::poll(future.as_mut(), &mut task_cx).is_pending());
3572        assert_eq!(pending.in_flight_len(), 1);
3573        assert!(cancellation.begin_finalization());
3574        assert!(wake_flag.0.load(Ordering::Acquire));
3575
3576        let error = block_on(future).expect_err("finalization must terminate the retained wait");
3577        assert_eq!(error.code, McpErrorCode::RequestCancelled);
3578        assert_eq!(pending.in_flight_len(), 0);
3579    }
3580
3581    #[test]
3582    fn request_local_cancellation_wins_when_response_is_already_ready() {
3583        let pending = Arc::new(PendingRequests::new());
3584        let pending_for_send = Arc::clone(&pending);
3585        let cancellation = McpRequestCancellation::new();
3586        let cancellation_for_send = cancellation.clone();
3587        let sender = RequestSender::new(
3588            Arc::clone(&pending),
3589            Arc::new(move |message| {
3590                let JsonRpcMessage::Request(request) = message else {
3591                    return Err("expected request".to_string());
3592                };
3593                let id = request
3594                    .id
3595                    .clone()
3596                    .ok_or_else(|| "expected request id".to_string())?;
3597                let response = JsonRpcResponse::success(id, serde_json::json!({"ready": true}));
3598                if !pending_for_send.route_response(&response) {
3599                    return Err("response was not routed".to_string());
3600                }
3601                let _ = cancellation_for_send.cancel();
3602                Ok(())
3603            }),
3604        )
3605        .for_request(cancellation);
3606        let cx = Cx::for_testing();
3607
3608        let error = block_on(sender.send_request::<serde_json::Value>(
3609            &cx,
3610            "test/cancellation-precedence",
3611            serde_json::json!({}),
3612        ))
3613        .expect_err("request-local cancellation must own an observable tie");
3614
3615        assert_eq!(error.code, McpErrorCode::RequestCancelled);
3616        assert_eq!(pending.in_flight_len(), 0);
3617        assert!(!cx.is_cancel_requested());
3618    }
3619
3620    #[test]
3621    fn test_route_unknown_response() {
3622        let pending = PendingRequests::new();
3623
3624        // Route a response with unknown ID
3625        let response = JsonRpcResponse::success(
3626            RequestId::Number(999999),
3627            serde_json::json!({"result": "ok"}),
3628        );
3629        assert!(!pending.route_response(&response));
3630    }
3631
3632    #[test]
3633    fn exact_legacy_negative_response_disposition_delivers_issued_waiter() {
3634        let pending = PendingRequests::with_max_in_flight_for_exact_legacy(1).unwrap();
3635        let (id, receiver) = pending.register().unwrap();
3636        assert_eq!(id, RequestId::Number(-1));
3637
3638        let RequestId::Number(first_emitted_id) = id.clone() else {
3639            panic!("exact-legacy IDs must be numeric");
3640        };
3641        #[allow(clippy::cast_precision_loss, clippy::cast_possible_truncation)]
3642        let f64_round_trip = (first_emitted_id as f64) as i64;
3643        assert_eq!(f64_round_trip, first_emitted_id);
3644
3645        let response = JsonRpcResponse::success(id, serde_json::json!({"result": "ok"}));
3646        assert_eq!(
3647            pending.route_response_with_disposition(&response),
3648            PendingResponseDisposition::Delivered
3649        );
3650        assert_eq!(
3651            receive_pending(receiver).unwrap(),
3652            serde_json::json!({"result": "ok"})
3653        );
3654    }
3655
3656    #[test]
3657    fn exact_legacy_negative_ids_descend_from_minus_one() {
3658        let pending = PendingRequests::with_max_in_flight_for_exact_legacy(2).unwrap();
3659
3660        let (first_id, _first_receiver) = pending.register().unwrap();
3661        let (next_id, _next_receiver) = pending.register().unwrap();
3662
3663        assert_eq!(first_id, RequestId::Number(-1));
3664        assert_eq!(next_id, RequestId::Number(-2));
3665    }
3666
3667    #[test]
3668    fn exact_legacy_negative_response_disposition_retires_issued_removed_id() {
3669        let pending = PendingRequests::with_max_in_flight_for_exact_legacy(1).unwrap();
3670        let (id, _receiver) = pending.register().unwrap();
3671        pending.remove(&id);
3672
3673        let response = JsonRpcResponse::success(id, serde_json::json!(null));
3674        assert_eq!(
3675            pending.route_response_with_disposition(&response),
3676            PendingResponseDisposition::RetiredGeneric
3677        );
3678        assert!(
3679            !pending.route_response(&response),
3680            "the public bool wrapper remains false for a retired response"
3681        );
3682    }
3683
3684    #[test]
3685    fn exact_legacy_negative_response_disposition_rejects_unissued_nearby_id() {
3686        let pending = PendingRequests::with_max_in_flight_for_exact_legacy(1).unwrap();
3687        let (issued_id, _receiver) = pending.register().unwrap();
3688        assert_eq!(issued_id, RequestId::Number(-1));
3689
3690        let unissued_id = RequestId::Number(-2);
3691        let response = JsonRpcResponse::success(unissued_id, serde_json::json!(null));
3692        assert_eq!(
3693            pending.route_response_with_disposition(&response),
3694            PendingResponseDisposition::Unmatched
3695        );
3696        assert!(
3697            !pending.route_response(&response),
3698            "the public bool wrapper remains false for an unissued response"
3699        );
3700    }
3701
3702    #[test]
3703    fn pending_requests_deliver_equivalent_numeric_response_spelling() {
3704        let pending = PendingRequests::new();
3705        let (id, receiver) = pending.register().unwrap();
3706        assert_eq!(id, RequestId::Number(FIRST_SERVER_REQUEST_ID));
3707
3708        let response = JsonRpcResponse::success(
3709            RequestId::Integer(format!("{FIRST_SERVER_REQUEST_ID}e0")),
3710            serde_json::json!({"result": "canonical"}),
3711        );
3712        assert_eq!(
3713            pending.route_response_with_disposition(&response),
3714            PendingResponseDisposition::Delivered
3715        );
3716        assert_eq!(
3717            receive_pending(receiver).unwrap(),
3718            serde_json::json!({"result": "canonical"})
3719        );
3720    }
3721
3722    #[test]
3723    fn exact_legacy_retires_equivalent_numeric_response_spelling() {
3724        let pending = PendingRequests::with_max_in_flight_for_exact_legacy(1).unwrap();
3725        let (id, _receiver) = pending.register().unwrap();
3726        assert_eq!(id, RequestId::Number(-1));
3727        pending.remove(&id);
3728
3729        let response = JsonRpcResponse::success(
3730            RequestId::Integer(format!("{FIRST_EXACT_LEGACY_SERVER_REQUEST_ID}e0")),
3731            serde_json::json!(null),
3732        );
3733        assert_eq!(
3734            pending.route_response_with_disposition(&response),
3735            PendingResponseDisposition::RetiredGeneric
3736        );
3737    }
3738
3739    #[test]
3740    fn removed_positive_id_remains_unmatched() {
3741        let pending = PendingRequests::new();
3742        let (id, _receiver) = pending.register().unwrap();
3743        pending.remove(&id);
3744
3745        let response = JsonRpcResponse::success(id, serde_json::json!(null));
3746        assert_eq!(
3747            pending.route_response_with_disposition(&response),
3748            PendingResponseDisposition::Unmatched
3749        );
3750        assert!(!pending.route_response(&response));
3751    }
3752
3753    #[test]
3754    fn exact_legacy_negative_ids_exhaust_at_js_safe_boundary_and_remain_retired() {
3755        let pending = PendingRequests::with_max_in_flight_for_exact_legacy(1).unwrap();
3756        pending.set_next_id_for_test(LAST_EXACT_LEGACY_SERVER_REQUEST_ID);
3757        let (last_id, _receiver) = pending.register().unwrap();
3758        assert_eq!(
3759            last_id,
3760            RequestId::Number(LAST_EXACT_LEGACY_SERVER_REQUEST_ID)
3761        );
3762        pending.remove(&last_id);
3763
3764        let exhausted = pending
3765            .register()
3766            .expect_err("the exact-legacy negative ID domain ends at the JavaScript safe boundary");
3767        assert_eq!(exhausted.message, REQUEST_ID_EXHAUSTED_ERROR);
3768
3769        let response = JsonRpcResponse::success(last_id.clone(), serde_json::json!(null));
3770        assert_eq!(
3771            pending.route_response_with_disposition(&response),
3772            PendingResponseDisposition::RetiredGeneric
3773        );
3774
3775        let first_response =
3776            JsonRpcResponse::success(RequestId::Number(-1), serde_json::json!(null));
3777        assert_eq!(
3778            pending.route_response_with_disposition(&first_response),
3779            PendingResponseDisposition::RetiredGeneric,
3780            "after exhaustion the entire JavaScript-safe negative range is retired"
3781        );
3782
3783        let out_of_range_response = JsonRpcResponse::success(
3784            RequestId::Number(LAST_EXACT_LEGACY_SERVER_REQUEST_ID - 1),
3785            serde_json::json!(null),
3786        );
3787        assert_eq!(
3788            pending.route_response_with_disposition(&out_of_range_response),
3789            PendingResponseDisposition::Unmatched,
3790            "a negative ID outside the JavaScript-safe range is never retired"
3791        );
3792
3793        let permanently_exhausted = pending
3794            .register()
3795            .expect_err("retiring the final negative ID must not permit reuse");
3796        assert_eq!(permanently_exhausted.message, REQUEST_ID_EXHAUSTED_ERROR);
3797    }
3798
3799    // ── PendingRequests additional coverage ───────────────────────────
3800
3801    #[test]
3802    fn pending_requests_default_is_same_as_new() {
3803        let pr = PendingRequests::default();
3804        let (id, _receiver) = pr.register().unwrap();
3805        // IDs start at 1_000_000
3806        assert_eq!(id, RequestId::Number(1_000_000));
3807        assert_eq!(pr.max_in_flight(), DEFAULT_MAX_IN_FLIGHT_REQUESTS);
3808    }
3809
3810    #[test]
3811    fn pending_requests_ids_are_sequential() {
3812        let pr = PendingRequests::new();
3813        let (id1, _receiver1) = pr.register().unwrap();
3814        let (id2, _receiver2) = pr.register().unwrap();
3815        let (id3, _receiver3) = pr.register().unwrap();
3816        assert_eq!(id1, RequestId::Number(1_000_000));
3817        assert_eq!(id2, RequestId::Number(1_000_001));
3818        assert_eq!(id3, RequestId::Number(1_000_002));
3819    }
3820
3821    #[test]
3822    fn pending_requests_limit_configuration_has_exact_hard_boundary() {
3823        let at_hard_limit =
3824            PendingRequests::with_max_in_flight(HARD_MAX_IN_FLIGHT_REQUESTS).unwrap();
3825        assert_eq!(at_hard_limit.max_in_flight(), HARD_MAX_IN_FLIGHT_REQUESTS);
3826
3827        for invalid in [0, HARD_MAX_IN_FLIGHT_REQUESTS + 1] {
3828            let error = PendingRequests::with_max_in_flight(invalid).unwrap_err();
3829            assert_eq!(error.code, McpErrorCode::InvalidParams);
3830            assert_eq!(error.message, INVALID_LIMIT_ERROR);
3831        }
3832    }
3833
3834    #[test]
3835    fn pending_requests_enforces_exact_in_flight_boundary_and_recovers_capacity() {
3836        let pr = PendingRequests::with_max_in_flight(2).unwrap();
3837        let (id1, receiver1) = pr.register().unwrap();
3838        let (_id2, _receiver2) = pr.register().unwrap();
3839        assert_eq!(pr.in_flight_len(), 2);
3840
3841        let error = pr.register().unwrap_err();
3842        assert_eq!(error.code, McpErrorCode::InternalError);
3843        assert_eq!(error.message, IN_FLIGHT_LIMIT_ERROR);
3844
3845        let response = JsonRpcResponse::success(id1, serde_json::json!(1));
3846        assert!(pr.route_response(&response));
3847        assert_eq!(receive_pending(receiver1).unwrap(), serde_json::json!(1));
3848        assert_eq!(pr.in_flight_len(), 1);
3849
3850        let (_id3, _receiver3) = pr.register().unwrap();
3851        assert_eq!(pr.in_flight_len(), 2);
3852    }
3853
3854    #[test]
3855    fn pending_request_ids_fail_closed_before_wrap_or_reuse() {
3856        let pr = PendingRequests::with_max_in_flight(4).unwrap();
3857        pr.set_next_id_for_test(i64::MAX);
3858        let (max_id, _max_receiver) = pr.register().unwrap();
3859        assert_eq!(max_id, RequestId::Number(i64::MAX));
3860
3861        let exhausted = pr
3862            .register()
3863            .expect_err("request IDs must never wrap back to an earlier value");
3864        assert_eq!(exhausted.message, REQUEST_ID_EXHAUSTED_ERROR);
3865        assert_eq!(pr.in_flight_len(), 1);
3866
3867        pr.remove(&max_id);
3868        let still_exhausted = pr
3869            .register()
3870            .expect_err("exhaustion must remain permanent after the last waiter leaves");
3871        assert_eq!(still_exhausted.message, REQUEST_ID_EXHAUSTED_ERROR);
3872    }
3873
3874    #[test]
3875    fn pending_requests_remove_prevents_routing() {
3876        let pr = PendingRequests::new();
3877        let (id, _receiver) = pr.register().unwrap();
3878
3879        // Remove the pending request
3880        pr.remove(&id);
3881
3882        // Routing should fail now
3883        let response = JsonRpcResponse::success(id, serde_json::json!(null));
3884        assert_eq!(
3885            pr.route_response_with_disposition(&response),
3886            PendingResponseDisposition::Unmatched
3887        );
3888        assert!(!pr.route_response(&response));
3889    }
3890
3891    #[test]
3892    fn pending_requests_route_response_without_id_returns_false() {
3893        let pr = PendingRequests::new();
3894        let (id, receiver) = pr.register().unwrap();
3895        // A response with no id
3896        let response = JsonRpcResponse {
3897            jsonrpc: std::borrow::Cow::Borrowed("2.0"),
3898            id: None,
3899            result: Some(serde_json::json!(null)),
3900            error: None,
3901        };
3902        assert!(!pr.route_response(&response));
3903        assert_eq!(pr.in_flight_len(), 1);
3904
3905        let response = JsonRpcResponse::success(id, serde_json::json!(42));
3906        assert!(pr.route_response(&response));
3907        assert_eq!(receive_pending(receiver).unwrap(), serde_json::json!(42));
3908    }
3909
3910    #[test]
3911    fn pending_requests_route_response_with_explicit_null_result() {
3912        let pr = PendingRequests::new();
3913        let (id, receiver) = pr.register().unwrap();
3914
3915        // An explicit JSON null is a present and valid success result.
3916        let response = JsonRpcResponse {
3917            jsonrpc: std::borrow::Cow::Borrowed("2.0"),
3918            id: Some(id),
3919            result: Some(serde_json::Value::Null),
3920            error: None,
3921        };
3922        assert!(pr.route_response(&response));
3923
3924        let result = receive_pending(receiver).unwrap();
3925        assert_eq!(result, serde_json::Value::Null);
3926    }
3927
3928    #[test]
3929    fn pending_requests_rejects_invalid_response_shapes_and_versions() {
3930        let cases = [
3931            (
3932                Some(serde_json::Value::Null),
3933                Some(JsonRpcError {
3934                    code: (-32_603).into(),
3935                    message: "secret both-member detail".to_string(),
3936                    data: Some(serde_json::json!({"secret": true})),
3937                }),
3938                "2.0",
3939            ),
3940            (None, None, "2.0"),
3941            (Some(serde_json::Value::Null), None, "1.0"),
3942        ];
3943
3944        for (result, error, version) in cases {
3945            let pr = PendingRequests::new();
3946            let (id, receiver) = pr.register().unwrap();
3947            let response = JsonRpcResponse {
3948                jsonrpc: std::borrow::Cow::Borrowed(version),
3949                result,
3950                error,
3951                id: Some(id),
3952            };
3953
3954            assert!(pr.route_response(&response));
3955            let error = receive_pending(receiver).unwrap_err();
3956            assert_eq!(error.code, McpErrorCode::InternalError);
3957            assert_eq!(error.message, INVALID_RESPONSE_ERROR);
3958            assert!(error.data.is_none());
3959            assert_eq!(pr.in_flight_len(), 0);
3960        }
3961    }
3962
3963    #[test]
3964    fn pending_requests_route_after_receiver_dropped_does_not_panic() {
3965        let pr = PendingRequests::new();
3966        let (id, receiver) = pr.register().unwrap();
3967
3968        // Drop the receiver
3969        drop(receiver);
3970
3971        // Routing should still succeed (sender.send returns Err but is ignored)
3972        let response = JsonRpcResponse::success(id, serde_json::json!(42));
3973        assert!(pr.route_response(&response));
3974    }
3975
3976    #[test]
3977    fn pending_requests_cancel_all_clears_pending() {
3978        let pr = PendingRequests::new();
3979        let (id, _receiver) = pr.register().unwrap();
3980
3981        pr.cancel_all();
3982
3983        // No more pending requests to route to
3984        let response = JsonRpcResponse::success(id, serde_json::json!(null));
3985        assert!(!pr.route_response(&response));
3986    }
3987
3988    #[test]
3989    fn pending_requests_cancel_all_empty_is_noop() {
3990        let pr = PendingRequests::new();
3991        // Should not panic on empty
3992        pr.cancel_all();
3993    }
3994
3995    #[test]
3996    fn pending_requests_cancel_all_permanently_rejects_new_waiters() {
3997        use std::sync::atomic::{AtomicBool, Ordering};
3998
3999        let pending = Arc::new(PendingRequests::new());
4000        let (_, receiver) = pending.register().unwrap();
4001
4002        pending.cancel_all();
4003        pending.cancel_all();
4004
4005        let cancelled = receive_pending(receiver).unwrap_err();
4006        assert_eq!(cancelled.code, McpErrorCode::InternalError);
4007        assert_eq!(cancelled.message, CONNECTION_CLOSED_ERROR);
4008        assert_eq!(pending.in_flight_len(), 0);
4009
4010        let registration_error = pending.register().unwrap_err();
4011        assert_eq!(registration_error.code, McpErrorCode::InternalError);
4012        assert_eq!(registration_error.message, CONNECTION_CLOSED_ERROR);
4013
4014        let send_called = Arc::new(AtomicBool::new(false));
4015        let send_called_for_callback = Arc::clone(&send_called);
4016        let sender = RequestSender::new(
4017            Arc::clone(&pending),
4018            Arc::new(move |_| {
4019                send_called_for_callback.store(true, Ordering::Release);
4020                Ok(())
4021            }),
4022        );
4023        let cx = Cx::for_testing();
4024        let error = block_on(sender.send_request::<serde_json::Value>(
4025            &cx,
4026            "test/after-close",
4027            serde_json::json!({}),
4028        ))
4029        .unwrap_err();
4030        assert_eq!(error.code, McpErrorCode::InternalError);
4031        assert_eq!(error.message, CONNECTION_CLOSED_ERROR);
4032        assert!(!send_called.load(Ordering::Acquire));
4033    }
4034
4035    #[test]
4036    fn pending_requests_debug_format() {
4037        let pr = PendingRequests::new();
4038        let debug = format!("{:?}", pr);
4039        assert!(debug.contains("PendingRequests"));
4040    }
4041
4042    // ── RequestSender ────────────────────────────────────────────────
4043
4044    #[test]
4045    fn request_sender_debug_format() {
4046        let pending = Arc::new(PendingRequests::new());
4047        let send_fn: TransportSendFn = Arc::new(|_| Ok(()));
4048        let sender = RequestSender::new(pending, send_fn);
4049        let debug = format!("{:?}", sender);
4050        assert!(debug.contains("RequestSender"));
4051    }
4052
4053    #[test]
4054    fn request_sender_transport_failure_returns_error() {
4055        let pending = Arc::new(PendingRequests::new());
4056        let send_fn: TransportSendFn = Arc::new(|_| Err("transport down".to_string()));
4057        let sender = RequestSender::new(pending, send_fn);
4058
4059        let cx = Cx::for_testing();
4060        let result: McpResult<serde_json::Value> =
4061            block_on(sender.send_request(&cx, "test/method", serde_json::json!({})));
4062        let err = result.unwrap_err();
4063        assert_eq!(err.message, TRANSPORT_SEND_ERROR);
4064        assert!(!err.message.contains("transport down"));
4065    }
4066
4067    #[test]
4068    fn request_sender_transport_failure_cleans_up_pending() {
4069        let pending = Arc::new(PendingRequests::new());
4070        let send_fn: TransportSendFn = Arc::new(|_| Err("fail".to_string()));
4071        let sender = RequestSender::new(Arc::clone(&pending), send_fn);
4072
4073        let cx = Cx::for_testing();
4074        let _error: McpResult<serde_json::Value> =
4075            block_on(sender.send_request(&cx, "test/method", serde_json::json!({})));
4076
4077        // The pending request should have been cleaned up
4078        let id = RequestId::Number(1_000_000); // first ID
4079        let response = JsonRpcResponse::success(id, serde_json::json!(null));
4080        assert!(!pending.route_response(&response));
4081    }
4082
4083    #[test]
4084    fn request_sender_clone() {
4085        let pending = Arc::new(PendingRequests::new());
4086        let send_fn: TransportSendFn = Arc::new(|_| Ok(()));
4087        let sender = RequestSender::new(pending, send_fn);
4088        let cloned = sender.clone();
4089        let debug = format!("{:?}", cloned);
4090        assert!(debug.contains("RequestSender"));
4091    }
4092
4093    #[test]
4094    fn dropping_request_future_releases_in_flight_capacity() {
4095        let pending = Arc::new(PendingRequests::with_max_in_flight(1).unwrap());
4096        let send_fn: TransportSendFn = Arc::new(|_| Ok(()));
4097        let sender = RequestSender::new(Arc::clone(&pending), send_fn);
4098        let cx = Cx::for_testing();
4099        {
4100            let mut future = Box::pin(sender.send_request::<serde_json::Value>(
4101                &cx,
4102                "test/method",
4103                serde_json::json!({}),
4104            ));
4105            let waker = std::task::Waker::noop();
4106            let mut task_cx = std::task::Context::from_waker(waker);
4107
4108            assert!(std::future::Future::poll(future.as_mut(), &mut task_cx).is_pending());
4109            assert_eq!(pending.in_flight_len(), 1);
4110        }
4111        assert_eq!(pending.in_flight_len(), 0);
4112    }
4113
4114    // ── RequestSender send_request paths ─────────────────────────────
4115
4116    #[test]
4117    fn reverse_request_routes_matching_response_without_cancellation_cleanup() {
4118        let pending = Arc::new(PendingRequests::new());
4119        let pending_clone = Arc::clone(&pending);
4120        let outbound = Arc::new(Mutex::new(Vec::new()));
4121        let outbound_for_send = Arc::clone(&outbound);
4122        let send_fn: TransportSendFn = Arc::new(move |msg| {
4123            if let JsonRpcMessage::Request(req) = msg {
4124                outbound_for_send
4125                    .lock()
4126                    .expect("test outbound mutex must not be poisoned")
4127                    .push(msg.clone());
4128                let id = req.id.clone().unwrap();
4129                let response = JsonRpcResponse::success(id, serde_json::json!({"answer": 42}));
4130                assert!(pending_clone.route_response(&response));
4131            }
4132            Ok(())
4133        });
4134        let sender = RequestSender::new(Arc::clone(&pending), send_fn);
4135        let cx = Cx::for_testing();
4136        let result: McpResult<serde_json::Value> =
4137            block_on(sender.send_request(&cx, "test/method", serde_json::json!({})));
4138        let value = result.unwrap();
4139        assert_eq!(value["answer"], 42);
4140        assert_eq!(pending.in_flight_len(), 0);
4141        assert_eq!(
4142            outbound
4143                .lock()
4144                .expect("test outbound mutex must not be poisoned")
4145                .len(),
4146            1
4147        );
4148    }
4149
4150    fn dropped_reverse_request_outbound(era: ProtocolEra) -> Vec<JsonRpcMessage> {
4151        let pending = Arc::new(PendingRequests::new());
4152        let pending_clone = Arc::clone(&pending);
4153        let outbound = Arc::new(Mutex::new(Vec::new()));
4154        let outbound_for_send = Arc::clone(&outbound);
4155        let send_fn: TransportSendFn = Arc::new(move |msg| {
4156            if let JsonRpcMessage::Request(req) = msg {
4157                outbound_for_send
4158                    .lock()
4159                    .expect("test outbound mutex must not be poisoned")
4160                    .push(msg.clone());
4161                if let Some(RequestId::Number(id)) = req.id.as_ref() {
4162                    // RH-5 planted negative: only the response correlation ID
4163                    // differs from the successful reverse-request path above.
4164                    let response = JsonRpcResponse::success(
4165                        RequestId::Number(*id + 1),
4166                        serde_json::json!({"answer": 42}),
4167                    );
4168                    assert!(!pending_clone.route_response(&response));
4169                }
4170            }
4171            Ok(())
4172        });
4173        let sender = RequestSender::new_for_era(era, Arc::clone(&pending), send_fn);
4174        let cx = Cx::for_testing();
4175
4176        {
4177            let mut future = Box::pin(sender.send_request::<serde_json::Value>(
4178                &cx,
4179                "test/method",
4180                serde_json::json!({}),
4181            ));
4182            let waker = std::task::Waker::noop();
4183            let mut task_cx = std::task::Context::from_waker(waker);
4184
4185            assert!(std::future::Future::poll(future.as_mut(), &mut task_cx).is_pending());
4186            assert_eq!(pending.in_flight_len(), 1);
4187        }
4188
4189        assert_eq!(pending.in_flight_len(), 0);
4190        outbound
4191            .lock()
4192            .expect("test outbound mutex must not be poisoned")
4193            .clone()
4194    }
4195
4196    #[test]
4197    fn dropped_legacy_reverse_request_emits_cancellation_control() {
4198        let outbound = dropped_reverse_request_outbound(ProtocolEra::Legacy2024);
4199        assert_eq!(outbound.len(), 2);
4200        let JsonRpcMessage::Request(cancelled) = &outbound[1] else {
4201            panic!("dropped exact-2024 reverse request must notify the peer");
4202        };
4203        assert_eq!(cancelled.id, None);
4204        assert_eq!(cancelled.method, "notifications/cancelled");
4205        assert_eq!(
4206            cancelled.params,
4207            Some(serde_json::json!({ "requestId": FIRST_SERVER_REQUEST_ID }))
4208        );
4209    }
4210
4211    #[test]
4212    fn dropped_modern_reverse_request_omits_cancellation_control() {
4213        // RH-5 planted negative: only the selected protocol era differs from
4214        // the legacy positive above.
4215        let outbound = dropped_reverse_request_outbound(ProtocolEra::Modern2026);
4216        assert_eq!(outbound.len(), 1);
4217        let JsonRpcMessage::Request(request) = &outbound[0] else {
4218            panic!("dropped modern reverse request must retain its initial request frame");
4219        };
4220        assert_eq!(request.method, "test/method");
4221        assert_eq!(request.id, Some(RequestId::Number(FIRST_SERVER_REQUEST_ID)));
4222    }
4223
4224    #[test]
4225    fn request_sender_error_response_path() {
4226        let pending = Arc::new(PendingRequests::new());
4227        let pending_clone = Arc::clone(&pending);
4228        let send_fn: TransportSendFn = Arc::new(move |msg| {
4229            if let JsonRpcMessage::Request(req) = msg {
4230                let id = req.id.clone().unwrap();
4231                let response = JsonRpcResponse::error(
4232                    Some(id),
4233                    JsonRpcError {
4234                        code: (-32600).into(),
4235                        message: "bad request".to_string(),
4236                        data: None,
4237                    },
4238                );
4239                pending_clone.route_response(&response);
4240            }
4241            Ok(())
4242        });
4243        let sender = RequestSender::new(Arc::clone(&pending), send_fn);
4244        let cx = Cx::for_testing();
4245        let result: McpResult<serde_json::Value> =
4246            block_on(sender.send_request(&cx, "test/method", serde_json::json!({})));
4247        let err = result.unwrap_err();
4248        assert_eq!(err.message, REMOTE_RESPONSE_ERROR);
4249        assert!(!err.message.contains("bad request"));
4250    }
4251
4252    #[test]
4253    fn request_sender_disconnected_path() {
4254        let pending = Arc::new(PendingRequests::new());
4255        let pending_clone = Arc::clone(&pending);
4256        let send_fn: TransportSendFn = Arc::new(move |msg| {
4257            if let JsonRpcMessage::Request(req) = msg {
4258                let id = req.id.clone().unwrap();
4259                // Remove the pending entry so tx is dropped, causing Disconnected
4260                pending_clone.remove(&id);
4261            }
4262            Ok(())
4263        });
4264        let sender = RequestSender::new(Arc::clone(&pending), send_fn);
4265        let cx = Cx::for_testing();
4266        let result: McpResult<serde_json::Value> =
4267            block_on(sender.send_request(&cx, "test/method", serde_json::json!({})));
4268        let err = result.unwrap_err();
4269        assert_eq!(err.message, RESPONSE_CHANNEL_ERROR);
4270    }
4271
4272    #[test]
4273    fn request_sender_deserialization_error() {
4274        let pending = Arc::new(PendingRequests::new());
4275        let pending_clone = Arc::clone(&pending);
4276        let send_fn: TransportSendFn = Arc::new(move |msg| {
4277            if let JsonRpcMessage::Request(req) = msg {
4278                let id = req.id.clone().unwrap();
4279                // Return a string value, which won't deserialize to Vec<String>
4280                let response =
4281                    JsonRpcResponse::success(id, serde_json::json!("not a vec of strings"));
4282                pending_clone.route_response(&response);
4283            }
4284            Ok(())
4285        });
4286        let sender = RequestSender::new(Arc::clone(&pending), send_fn);
4287        let cx = Cx::for_testing();
4288        let result: McpResult<Vec<String>> =
4289            block_on(sender.send_request(&cx, "test/method", serde_json::json!({})));
4290        let err = result.unwrap_err();
4291        assert_eq!(err.message, RESPONSE_PAYLOAD_ERROR);
4292        assert!(!err.message.contains("expected"));
4293    }
4294
4295    // ── cancel_all error details ─────────────────────────────────────
4296
4297    #[test]
4298    fn cancel_all_sends_connection_closed_error() {
4299        let pr = PendingRequests::new();
4300        let (_, receiver) = pr.register().unwrap();
4301        pr.cancel_all();
4302        let result = receive_pending(receiver);
4303        let err = result.unwrap_err();
4304        assert_eq!(err.code, McpErrorCode::InternalError);
4305        assert_eq!(err.message, CONNECTION_CLOSED_ERROR);
4306        assert!(err.data.is_none());
4307    }
4308
4309    // ── route_response with error containing data ────────────────────
4310
4311    #[test]
4312    fn route_response_error_with_data() {
4313        let pr = PendingRequests::new();
4314        let (id, receiver) = pr.register().unwrap();
4315        let response = JsonRpcResponse::error(
4316            Some(id),
4317            JsonRpcError {
4318                code: (-32001).into(),
4319                message: "custom error".to_string(),
4320                data: Some(serde_json::json!({"detail": "extra info"})),
4321            },
4322        );
4323        assert!(pr.route_response(&response));
4324        let result = receive_pending(receiver);
4325        let err = result.unwrap_err();
4326        assert_eq!(err.code, McpErrorCode::ResourceNotFound);
4327        assert_eq!(err.message, REMOTE_RESPONSE_ERROR);
4328        assert!(err.data.is_none());
4329    }
4330
4331    // ── Multiple concurrent register/route ───────────────────────────
4332
4333    #[test]
4334    fn pending_requests_multiple_register_and_route_independently() {
4335        let pr = PendingRequests::new();
4336        let (id1, rx1) = pr.register().unwrap();
4337        let (id2, rx2) = pr.register().unwrap();
4338        let (id3, rx3) = pr.register().unwrap();
4339
4340        // Route them out of order
4341        let r2 = JsonRpcResponse::success(id2.clone(), serde_json::json!("second"));
4342        let r3 = JsonRpcResponse::success(id3.clone(), serde_json::json!("third"));
4343        let r1 = JsonRpcResponse::success(id1.clone(), serde_json::json!("first"));
4344        assert!(pr.route_response(&r2));
4345        assert!(pr.route_response(&r3));
4346        assert!(pr.route_response(&r1));
4347
4348        assert_eq!(receive_pending(rx1).unwrap(), serde_json::json!("first"));
4349        assert_eq!(receive_pending(rx2).unwrap(), serde_json::json!("second"));
4350        assert_eq!(receive_pending(rx3).unwrap(), serde_json::json!("third"));
4351    }
4352
4353    #[test]
4354    fn pending_request_trackers_isolate_identical_wire_ids() {
4355        let connection_a = PendingRequests::new();
4356        let connection_b = PendingRequests::new();
4357        let (id_a, receiver_a) = connection_a.register().unwrap();
4358        let (id_b, receiver_b) = connection_b.register().unwrap();
4359        assert_eq!(id_a, id_b);
4360
4361        let response = JsonRpcResponse::success(id_b, serde_json::json!("connection-b"));
4362        assert!(connection_b.route_response(&response));
4363        assert_eq!(
4364            receive_pending(receiver_b).unwrap(),
4365            serde_json::json!("connection-b")
4366        );
4367        assert_eq!(connection_b.in_flight_len(), 0);
4368
4369        // Routing on B cannot consume A's same-numbered waiter because the
4370        // registry itself is the immutable connection ownership boundary.
4371        assert_eq!(connection_a.in_flight_len(), 1);
4372        connection_a.cancel_all();
4373        let error = receive_pending(receiver_a).unwrap_err();
4374        assert_eq!(error.message, CONNECTION_CLOSED_ERROR);
4375    }
4376
4377    // ── Transport sender constructors ────────────────────────────────
4378
4379    #[test]
4380    fn transport_sampling_sender_new_and_clone() {
4381        let pending = Arc::new(PendingRequests::new());
4382        let send_fn: TransportSendFn = Arc::new(|_| Ok(()));
4383        let sender = RequestSender::new(pending, send_fn);
4384        let sampling = TransportSamplingSender::new(sender, McpContext::new(Cx::for_testing(), 0));
4385        let _cloned = sampling.clone();
4386    }
4387
4388    #[test]
4389    fn transport_elicitation_sender_new_and_clone() {
4390        let pending = Arc::new(PendingRequests::new());
4391        let send_fn: TransportSendFn = Arc::new(|_| Ok(()));
4392        let sender = RequestSender::new(pending, send_fn);
4393        let elicitation =
4394            TransportElicitationSender::new(sender, McpContext::new(Cx::for_testing(), 0));
4395        let _cloned = elicitation.clone();
4396    }
4397
4398    #[test]
4399    fn transport_roots_provider_new_and_clone() {
4400        let pending = Arc::new(PendingRequests::new());
4401        let send_fn: TransportSendFn = Arc::new(|_| Ok(()));
4402        let sender = RequestSender::new(pending, send_fn);
4403        let roots = TransportRootsProvider::new(sender, McpContext::new(Cx::for_testing(), 0));
4404        let _cloned = roots.clone();
4405    }
4406
4407    // ── lock_state with poisoned mutex ───────────────────────────────
4408
4409    #[test]
4410    fn pending_requests_lock_state_recovers_from_poison() {
4411        let pr = Arc::new(PendingRequests::new());
4412        let (id, receiver) = pr.register().unwrap();
4413
4414        // Poison the mutex by panicking while holding the lock
4415        let pr2 = Arc::clone(&pr);
4416        let _ = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
4417            let _guard = pr2.state.lock().unwrap();
4418            panic!("intentional poison");
4419        }));
4420
4421        // lock_state should recover from poison (into_inner)
4422        // Routing should still work
4423        let response = JsonRpcResponse::success(id, serde_json::json!("recovered"));
4424        assert!(pr.route_response(&response));
4425        let result = receive_pending(receiver).unwrap();
4426        assert_eq!(result, serde_json::json!("recovered"));
4427    }
4428
4429    // ── TransportSamplingSender — create_message ─────────────────────
4430
4431    fn make_sender_with_responder(
4432        responder: impl Fn(&JsonRpcRequest) -> serde_json::Value + Send + Sync + 'static,
4433    ) -> RequestSender {
4434        let pending = Arc::new(PendingRequests::new());
4435        let pending_clone = Arc::clone(&pending);
4436        let send_fn: TransportSendFn = Arc::new(move |msg| {
4437            if let JsonRpcMessage::Request(req) = msg {
4438                let id = req.id.clone().unwrap();
4439                let result = responder(req);
4440                let response = JsonRpcResponse::success(id, result);
4441                pending_clone.route_response(&response);
4442            }
4443            Ok(())
4444        });
4445        RequestSender::new(pending, send_fn)
4446    }
4447
4448    #[test]
4449    fn transport_sampling_sender_create_message_text() {
4450        let sender = make_sender_with_responder(|request| {
4451            let params = request
4452                .params
4453                .as_ref()
4454                .expect("sampling request must retain parameters");
4455            assert!(
4456                params.get("metadata").is_none(),
4457                "the transport must omit unspecified provider metadata"
4458            );
4459            serde_json::json!({
4460                "content": {"type": "text", "text": "Hello world"},
4461                "role": "assistant",
4462                "model": "test-model",
4463                "stopReason": "endTurn"
4464            })
4465        });
4466        let sampling = TransportSamplingSender::new(sender, McpContext::new(Cx::for_testing(), 0));
4467
4468        let request = SamplingRequest {
4469            messages: vec![fastmcp_core::SamplingRequestMessage {
4470                role: SamplingRole::User,
4471                text: "Hi".to_string(),
4472            }],
4473            max_tokens: 100,
4474            system_prompt: Some("Be helpful".to_string()),
4475            temperature: Some(0.7),
4476            stop_sequences: vec!["STOP".to_string()],
4477            model_hints: vec![],
4478        };
4479
4480        let future = SamplingSender::create_message(&sampling, request);
4481        let result = fastmcp_core::block_on(future).unwrap();
4482        assert_eq!(result.text, "Hello world");
4483        assert_eq!(result.model, "test-model");
4484        assert!(matches!(result.stop_reason, SamplingStopReason::EndTurn));
4485    }
4486
4487    #[test]
4488    fn transport_sampling_sender_round_trips_open_legacy_stop_reason_through_callback() {
4489        let expected = serde_json::json!({
4490            "content": {"type": "text", "text": "legacy completion"},
4491            "role": "assistant",
4492            "model": "legacy-model",
4493            "stopReason": "provider_safety_limit"
4494        });
4495        let reply = expected.clone();
4496        let sender = make_sender_with_responder(move |request| {
4497            assert_eq!(request.method, "sampling/createMessage");
4498            reply.clone()
4499        });
4500        let sampling = TransportSamplingSender::new(sender, McpContext::new(Cx::for_testing(), 0));
4501
4502        let callback_response = fastmcp_core::block_on(SamplingSender::create_message(
4503            &sampling,
4504            SamplingRequest::prompt("Hi", 10),
4505        ))
4506        .expect("legacy sampling callback must retain an open stopReason");
4507        assert_eq!(
4508            callback_response.stop_reason,
4509            SamplingStopReason::Other("provider_safety_limit".to_owned())
4510        );
4511
4512        let emitted = fastmcp_protocol::CreateMessageResult {
4513            content: fastmcp_protocol::SamplingContent::Text {
4514                text: callback_response.text,
4515            },
4516            role: fastmcp_protocol::Role::Assistant,
4517            model: callback_response.model,
4518            stop_reason: callback_response
4519                .stop_reason
4520                .as_wire_value()
4521                .map(str::to_owned),
4522            meta: None,
4523        };
4524        let emitted = serde_json::to_value(emitted)
4525            .expect("legacy sampling callback response must serialize");
4526        assert_eq!(emitted, expected);
4527        assert!(emitted.get("resultType").is_none());
4528    }
4529
4530    #[test]
4531    fn transport_sampling_sender_create_message_image() {
4532        let sender = make_sender_with_responder(|_| {
4533            serde_json::json!({
4534                "content": {"type": "image", "data": "aW1hZ2VkYXRh", "mimeType": "image/png"},
4535                "role": "assistant",
4536                "model": "vision-model",
4537                "stopReason": "maxTokens"
4538            })
4539        });
4540        let sampling = TransportSamplingSender::new(sender, McpContext::new(Cx::for_testing(), 0));
4541
4542        let request = SamplingRequest {
4543            messages: vec![fastmcp_core::SamplingRequestMessage {
4544                role: SamplingRole::User,
4545                text: "Describe image".to_string(),
4546            }],
4547            max_tokens: 50,
4548            system_prompt: None,
4549            temperature: None,
4550            stop_sequences: vec![],
4551            model_hints: vec![],
4552        };
4553
4554        let future = SamplingSender::create_message(&sampling, request);
4555        let result = fastmcp_core::block_on(future).unwrap();
4556        // Image content is formatted as "[image: N bytes, type: ...]"
4557        assert!(result.text.contains("image"));
4558        assert!(result.text.contains("image/png"));
4559        assert_eq!(result.model, "vision-model");
4560        assert!(matches!(result.stop_reason, SamplingStopReason::MaxTokens));
4561    }
4562
4563    #[test]
4564    fn transport_sampling_sender_create_message_with_model_hints() {
4565        let sender = make_sender_with_responder(|req| {
4566            // Verify model_preferences was sent
4567            let params: serde_json::Value =
4568                serde_json::from_value(req.params.clone().unwrap()).unwrap();
4569            assert!(params["modelPreferences"]["hints"].is_array());
4570            serde_json::json!({
4571                "content": {"type": "text", "text": "ok"},
4572                "role": "assistant",
4573                "model": "preferred",
4574                "stopReason": "stopSequence"
4575            })
4576        });
4577        let sampling = TransportSamplingSender::new(sender, McpContext::new(Cx::for_testing(), 0));
4578
4579        let request = SamplingRequest {
4580            messages: vec![fastmcp_core::SamplingRequestMessage {
4581                role: SamplingRole::User,
4582                text: "Hi".to_string(),
4583            }],
4584            max_tokens: 10,
4585            system_prompt: None,
4586            temperature: None,
4587            stop_sequences: vec![],
4588            model_hints: vec!["claude-3".to_string()],
4589        };
4590
4591        let future = SamplingSender::create_message(&sampling, request);
4592        let result = fastmcp_core::block_on(future).unwrap();
4593        assert!(matches!(
4594            result.stop_reason,
4595            SamplingStopReason::StopSequence
4596        ));
4597    }
4598
4599    #[test]
4600    fn transport_sampling_sender_create_message_assistant_role() {
4601        let sender = make_sender_with_responder(|req| {
4602            let params: serde_json::Value =
4603                serde_json::from_value(req.params.clone().unwrap()).unwrap();
4604            assert_eq!(params["messages"][0]["role"], "assistant");
4605            serde_json::json!({
4606                "content": {"type": "text", "text": "continued"},
4607                "role": "assistant",
4608                "model": "m",
4609                "stopReason": "endTurn"
4610            })
4611        });
4612        let sampling = TransportSamplingSender::new(sender, McpContext::new(Cx::for_testing(), 0));
4613
4614        let request = SamplingRequest {
4615            messages: vec![fastmcp_core::SamplingRequestMessage {
4616                role: SamplingRole::Assistant,
4617                text: "Previous response".to_string(),
4618            }],
4619            max_tokens: 10,
4620            system_prompt: None,
4621            temperature: None,
4622            stop_sequences: vec![],
4623            model_hints: vec![],
4624        };
4625
4626        let future = SamplingSender::create_message(&sampling, request);
4627        let result = fastmcp_core::block_on(future).unwrap();
4628        assert_eq!(result.text, "continued");
4629    }
4630
4631    #[test]
4632    fn transport_sampling_sender_rejects_non_assistant_result_role() {
4633        let sender = make_sender_with_responder(|_| {
4634            serde_json::json!({
4635                "content": {"type": "text", "text": "not authoritative"},
4636                "role": "user",
4637                "model": "m",
4638                "stopReason": "endTurn"
4639            })
4640        });
4641        let sampling = TransportSamplingSender::new(sender, McpContext::new(Cx::for_testing(), 0));
4642        let request = SamplingRequest::prompt("Hi", 10);
4643
4644        let error = fastmcp_core::block_on(SamplingSender::create_message(&sampling, request))
4645            .expect_err("sampling results must retain the documented assistant role");
4646
4647        assert_eq!(error.message, RESPONSE_PAYLOAD_ERROR);
4648    }
4649
4650    // ── TransportElicitationSender — elicit ──────────────────────────
4651
4652    #[test]
4653    fn transport_elicitation_sender_form_accept_with_content() {
4654        let sender = make_sender_with_responder(|req| {
4655            let params: serde_json::Value =
4656                serde_json::from_value(req.params.clone().unwrap()).unwrap();
4657            assert_eq!(params["mode"], "form");
4658            serde_json::json!({
4659                "action": "accept",
4660                "content": {
4661                    "name": "Alice",
4662                    "age": 30,
4663                    "active": true,
4664                    "score": 9.5,
4665                    "tags": ["a", "b"],
4666                    "empty": null
4667                }
4668            })
4669        });
4670        let elicitation =
4671            TransportElicitationSender::new(sender, McpContext::new(Cx::for_testing(), 0));
4672
4673        let request = ElicitationRequest {
4674            message: "Fill the form".to_string(),
4675            mode: ElicitationMode::Form,
4676            schema: Some(serde_json::json!({"type": "object"})),
4677            url: None,
4678            elicitation_id: None,
4679        };
4680
4681        let future = ElicitationSender::elicit(&elicitation, request);
4682        let result = fastmcp_core::block_on(future).unwrap();
4683        assert!(matches!(result.action, ElicitationAction::Accept));
4684        let content = result.content.unwrap();
4685        assert_eq!(content["name"], serde_json::json!("Alice"));
4686        assert_eq!(content["age"], serde_json::json!(30));
4687        assert_eq!(content["active"], serde_json::json!(true));
4688        assert_eq!(content["score"], serde_json::json!(9.5));
4689        assert_eq!(content["tags"], serde_json::json!(["a", "b"]));
4690        assert_eq!(content["empty"], serde_json::Value::Null);
4691    }
4692
4693    #[test]
4694    fn transport_elicitation_sender_form_decline() {
4695        let sender = make_sender_with_responder(|_| {
4696            serde_json::json!({
4697                "action": "decline"
4698            })
4699        });
4700        let elicitation =
4701            TransportElicitationSender::new(sender, McpContext::new(Cx::for_testing(), 0));
4702
4703        let request = ElicitationRequest {
4704            message: "Confirm?".to_string(),
4705            mode: ElicitationMode::Form,
4706            schema: Some(serde_json::json!({"type": "object"})),
4707            url: None,
4708            elicitation_id: None,
4709        };
4710
4711        let future = ElicitationSender::elicit(&elicitation, request);
4712        let result = fastmcp_core::block_on(future).unwrap();
4713        assert!(matches!(result.action, ElicitationAction::Decline));
4714        assert!(result.content.is_none());
4715    }
4716
4717    #[test]
4718    fn transport_elicitation_sender_url_mode() {
4719        let sender = make_sender_with_responder(|req| {
4720            let params: serde_json::Value =
4721                serde_json::from_value(req.params.clone().unwrap()).unwrap();
4722            assert_eq!(params["mode"], "url");
4723            assert_eq!(params["url"], "https://example.com/auth");
4724            serde_json::json!({
4725                "action": "cancel"
4726            })
4727        });
4728        let elicitation =
4729            TransportElicitationSender::new(sender, McpContext::new(Cx::for_testing(), 0));
4730
4731        let request = ElicitationRequest {
4732            message: "Please authenticate".to_string(),
4733            mode: ElicitationMode::Url,
4734            schema: None,
4735            url: Some("https://example.com/auth".to_string()),
4736            elicitation_id: Some("eid-123".to_string()),
4737        };
4738
4739        let future = ElicitationSender::elicit(&elicitation, request);
4740        let result = fastmcp_core::block_on(future).unwrap();
4741        assert!(matches!(result.action, ElicitationAction::Cancel));
4742    }
4743
4744    // ── TransportRootsProvider — list_roots ──────────────────────────
4745
4746    #[test]
4747    fn transport_roots_provider_list_roots() {
4748        let sender = make_sender_with_responder(|_| {
4749            serde_json::json!({
4750                "roots": [
4751                    {"uri": "file:///home/user/project", "name": "Project"},
4752                    {"uri": "file:///tmp"}
4753                ]
4754            })
4755        });
4756        let roots = TransportRootsProvider::new(sender, McpContext::new(Cx::for_testing(), 0));
4757        let result = block_on(roots.list_roots()).unwrap();
4758        assert_eq!(result.len(), 2);
4759        assert_eq!(result[0].uri, "file:///home/user/project");
4760        assert_eq!(result[0].name, Some("Project".to_string()));
4761        assert_eq!(result[1].uri, "file:///tmp");
4762        assert!(result[1].name.is_none());
4763    }
4764
4765    #[test]
4766    fn transport_roots_provider_maps_wire_roots_to_core_roots() {
4767        let sender = make_sender_with_responder(|_| {
4768            serde_json::json!({
4769                "roots": [{"uri": "file:///workspace", "name": "workspace"}]
4770            })
4771        });
4772        let roots = TransportRootsProvider::new(sender, McpContext::new(Cx::for_testing(), 0));
4773
4774        let result = fastmcp_core::block_on(fastmcp_core::RootsProvider::list_roots(&roots))
4775            .expect("transport roots map into the core context type");
4776        assert_eq!(
4777            result,
4778            vec![ClientRoot::with_name("file:///workspace", "workspace")]
4779        );
4780    }
4781
4782    #[test]
4783    fn transport_roots_provider_empty_roots() {
4784        let sender = make_sender_with_responder(|_| serde_json::json!({ "roots": [] }));
4785        let roots = TransportRootsProvider::new(sender, McpContext::new(Cx::for_testing(), 0));
4786        let result = block_on(roots.list_roots()).unwrap();
4787        assert!(result.is_empty());
4788    }
4789
4790    #[test]
4791    fn transport_roots_provider_trait_preserves_originating_deadline() {
4792        use std::sync::atomic::{AtomicBool, Ordering};
4793
4794        let sent = Arc::new(AtomicBool::new(false));
4795        let sent_for_transport = Arc::clone(&sent);
4796        let sender = RequestSender::new(
4797            Arc::new(PendingRequests::new()),
4798            Arc::new(move |_| {
4799                sent_for_transport.store(true, Ordering::Release);
4800                Ok(())
4801            }),
4802        );
4803        let roots = TransportRootsProvider::new(
4804            sender,
4805            McpContext::new(Cx::for_testing(), 0).with_budget_ceiling(
4806                asupersync::Budget::new().with_deadline(asupersync::Time::ZERO),
4807            ),
4808        );
4809
4810        let error = block_on(fastmcp_core::RootsProvider::list_roots(&roots))
4811            .expect_err("an expired originating request must not issue roots/list");
4812        assert_eq!(error.code, McpErrorCode::RequestCancelled);
4813        assert!(
4814            !sent.load(Ordering::Acquire),
4815            "the raw Cx must not relax the originating framework deadline ceiling"
4816        );
4817    }
4818
4819    #[test]
4820    fn transport_roots_provider_trait_preserves_originating_cancellation() {
4821        use std::sync::atomic::{AtomicBool, Ordering};
4822
4823        let sent = Arc::new(AtomicBool::new(false));
4824        let sent_for_transport = Arc::clone(&sent);
4825        let sender = RequestSender::new(
4826            Arc::new(PendingRequests::new()),
4827            Arc::new(move |_| {
4828                sent_for_transport.store(true, Ordering::Release);
4829                Ok(())
4830            }),
4831        );
4832        let cancellation = McpRequestCancellation::new();
4833        cancellation.cancel();
4834        let roots = TransportRootsProvider::new(
4835            sender,
4836            McpContext::new(Cx::for_testing(), 0).with_request_cancellation(cancellation),
4837        );
4838
4839        let error = block_on(fastmcp_core::RootsProvider::list_roots(&roots))
4840            .expect_err("a cancelled originating request must not issue roots/list");
4841        assert_eq!(error.code, McpErrorCode::RequestCancelled);
4842        assert!(
4843            !sent.load(Ordering::Acquire),
4844            "the raw Cx must not relax the originating framework cancellation"
4845        );
4846    }
4847
4848    // ── RequestSender ID cleanup after success ───────────────────────
4849
4850    // ── RequestSender — cancelled cx path ──────────────────────────
4851
4852    #[test]
4853    fn request_sender_cancelled_cx_returns_cancelled_error() {
4854        let pending = Arc::new(PendingRequests::new());
4855        // Transport succeeds but never sends a response
4856        let send_fn: TransportSendFn = Arc::new(|_| Ok(()));
4857        let sender = RequestSender::new(Arc::clone(&pending), send_fn);
4858
4859        let cx = Cx::for_testing();
4860        cx.set_cancel_requested(true);
4861
4862        let result: McpResult<serde_json::Value> =
4863            block_on(sender.send_request(&cx, "test/cancel", serde_json::json!({})));
4864        let err = result.unwrap_err();
4865        assert_eq!(err.code, McpErrorCode::RequestCancelled);
4866    }
4867
4868    // ── Elicitation request/response validation ─────────────────
4869
4870    #[test]
4871    fn transport_elicitation_sender_rejects_missing_url_fields_before_send() {
4872        let sender = make_sender_with_responder(|_| {
4873            panic!("an invalid URL elicitation must not reach the transport")
4874        });
4875        let elicitation =
4876            TransportElicitationSender::new(sender, McpContext::new(Cx::for_testing(), 0));
4877
4878        let request = ElicitationRequest {
4879            message: "Auth".to_string(),
4880            mode: ElicitationMode::Url,
4881            schema: None,
4882            url: None,
4883            elicitation_id: None,
4884        };
4885
4886        let future = ElicitationSender::elicit(&elicitation, request);
4887        let error = fastmcp_core::block_on(future)
4888            .expect_err("missing URL fields must be a local input error");
4889        assert_eq!(error.code, McpErrorCode::InvalidParams);
4890        assert_eq!(error.message, INVALID_ELICITATION_REQUEST_ERROR);
4891    }
4892
4893    #[test]
4894    fn transport_elicitation_sender_rejects_accepted_form_without_content() {
4895        let sender = make_sender_with_responder(|_| serde_json::json!({ "action": "accept" }));
4896        let elicitation =
4897            TransportElicitationSender::new(sender, McpContext::new(Cx::for_testing(), 0));
4898        let request = ElicitationRequest::form(
4899            "Fill the form",
4900            serde_json::json!({
4901                "type": "object"
4902            }),
4903        );
4904
4905        let error = fastmcp_core::block_on(ElicitationSender::elicit(&elicitation, request))
4906            .expect_err("accepted form mode must carry form content");
4907        assert_eq!(error.message, RESPONSE_PAYLOAD_ERROR);
4908    }
4909
4910    #[test]
4911    fn transport_elicitation_sender_rejects_accepted_url_content() {
4912        let sender = make_sender_with_responder(|_| {
4913            serde_json::json!({
4914                "action": "accept",
4915                "content": {"credential": "must-not-be-exposed"}
4916            })
4917        });
4918        let elicitation =
4919            TransportElicitationSender::new(sender, McpContext::new(Cx::for_testing(), 0));
4920        let request = ElicitationRequest::url("Authenticate", "https://example.com", "eid-1");
4921
4922        let error = fastmcp_core::block_on(ElicitationSender::elicit(&elicitation, request))
4923            .expect_err("accepted URL mode must not expose in-band content");
4924        assert_eq!(error.message, RESPONSE_PAYLOAD_ERROR);
4925        assert!(!error.message.contains("credential"));
4926    }
4927
4928    #[test]
4929    fn transport_elicitation_sender_does_not_expose_non_accept_content() {
4930        let sender = make_sender_with_responder(|_| {
4931            serde_json::json!({
4932                "action": "decline",
4933                "content": {"credential": "must-not-be-exposed"}
4934            })
4935        });
4936        let elicitation =
4937            TransportElicitationSender::new(sender, McpContext::new(Cx::for_testing(), 0));
4938        let request = ElicitationRequest::form(
4939            "Fill the form",
4940            serde_json::json!({
4941                "type": "object"
4942            }),
4943        );
4944
4945        let result = fastmcp_core::block_on(ElicitationSender::elicit(&elicitation, request))
4946            .expect("decline content is a SHOULD deviation, not accepted data");
4947        assert_eq!(result.action, ElicitationAction::Decline);
4948        assert!(result.content.is_none());
4949    }
4950
4951    // ── TransportRootsProvider — transport failure ───────────────
4952
4953    #[test]
4954    fn transport_roots_provider_transport_failure() {
4955        let pending = Arc::new(PendingRequests::new());
4956        let send_fn: TransportSendFn = Arc::new(|_| Err("network error".to_string()));
4957        let sender = RequestSender::new(pending, send_fn);
4958        let roots = TransportRootsProvider::new(sender, McpContext::new(Cx::for_testing(), 0));
4959        let result = block_on(roots.list_roots());
4960        assert_eq!(result.unwrap_err().message, TRANSPORT_SEND_ERROR);
4961    }
4962
4963    #[test]
4964    fn transport_roots_provider_core_trait_preserves_transport_failure() {
4965        let pending = Arc::new(PendingRequests::new());
4966        let send_fn: TransportSendFn = Arc::new(|_| Err("network error".to_string()));
4967        let roots = TransportRootsProvider::new(
4968            RequestSender::new(pending, send_fn),
4969            McpContext::new(Cx::for_testing(), 0),
4970        );
4971
4972        let error = fastmcp_core::block_on(fastmcp_core::RootsProvider::list_roots(&roots))
4973            .expect_err("the same transport failure must cross the core provider seam");
4974        assert_eq!(error.message, TRANSPORT_SEND_ERROR);
4975    }
4976
4977    // ── SamplingSender — transport failure ───────────────────────
4978
4979    #[test]
4980    fn transport_sampling_sender_transport_failure() {
4981        let pending = Arc::new(PendingRequests::new());
4982        let send_fn: TransportSendFn = Arc::new(|_| Err("connection reset".to_string()));
4983        let sender = RequestSender::new(pending, send_fn);
4984        let sampling = TransportSamplingSender::new(sender, McpContext::new(Cx::for_testing(), 0));
4985
4986        let request = SamplingRequest {
4987            messages: vec![fastmcp_core::SamplingRequestMessage {
4988                role: SamplingRole::User,
4989                text: "Hi".to_string(),
4990            }],
4991            max_tokens: 10,
4992            system_prompt: None,
4993            temperature: None,
4994            stop_sequences: vec![],
4995            model_hints: vec![],
4996        };
4997
4998        let future = SamplingSender::create_message(&sampling, request);
4999        let result = fastmcp_core::block_on(future);
5000        assert_eq!(result.unwrap_err().message, TRANSPORT_SEND_ERROR);
5001    }
5002
5003    // ── SamplingSender — multiple messages ───────────────────────
5004
5005    #[test]
5006    fn transport_sampling_sender_multiple_messages() {
5007        let sender = make_sender_with_responder(|req| {
5008            let params: serde_json::Value =
5009                serde_json::from_value(req.params.clone().unwrap()).unwrap();
5010            let messages = params["messages"].as_array().unwrap();
5011            assert_eq!(messages.len(), 3);
5012            assert_eq!(messages[0]["role"], "user");
5013            assert_eq!(messages[1]["role"], "assistant");
5014            assert_eq!(messages[2]["role"], "user");
5015            serde_json::json!({
5016                "content": {"type": "text", "text": "done"},
5017                "role": "assistant",
5018                "model": "m",
5019                "stopReason": "endTurn"
5020            })
5021        });
5022        let sampling = TransportSamplingSender::new(sender, McpContext::new(Cx::for_testing(), 0));
5023
5024        let request = SamplingRequest {
5025            messages: vec![
5026                fastmcp_core::SamplingRequestMessage {
5027                    role: SamplingRole::User,
5028                    text: "Hello".to_string(),
5029                },
5030                fastmcp_core::SamplingRequestMessage {
5031                    role: SamplingRole::Assistant,
5032                    text: "Hi".to_string(),
5033                },
5034                fastmcp_core::SamplingRequestMessage {
5035                    role: SamplingRole::User,
5036                    text: "Follow up".to_string(),
5037                },
5038            ],
5039            max_tokens: 100,
5040            system_prompt: None,
5041            temperature: None,
5042            stop_sequences: vec![],
5043            model_hints: vec![],
5044        };
5045
5046        let future = SamplingSender::create_message(&sampling, request);
5047        let result = fastmcp_core::block_on(future).unwrap();
5048        assert_eq!(result.text, "done");
5049    }
5050
5051    // ── RequestSender — ID cleanup after success ────────────────
5052
5053    #[test]
5054    fn request_sender_id_cleaned_from_pending_after_success() {
5055        let pending = Arc::new(PendingRequests::new());
5056        let pending_clone = Arc::clone(&pending);
5057        let send_fn: TransportSendFn = Arc::new(move |msg| {
5058            if let JsonRpcMessage::Request(req) = msg {
5059                let id = req.id.clone().unwrap();
5060                let response = JsonRpcResponse::success(id, serde_json::json!(null));
5061                pending_clone.route_response(&response);
5062            }
5063            Ok(())
5064        });
5065        let sender = RequestSender::new(Arc::clone(&pending), send_fn);
5066        let cx = Cx::for_testing();
5067        let _: serde_json::Value =
5068            block_on(sender.send_request(&cx, "test/method", serde_json::json!({}))).unwrap();
5069
5070        // The pending request should have been consumed by route_response
5071        let first_id = RequestId::Number(1_000_000);
5072        let response = JsonRpcResponse::success(first_id, serde_json::json!(null));
5073        assert!(!pending.route_response(&response));
5074    }
5075}