mcp_utils/client/
mrtr.rs

1use rmcp::model::{DEFAULT_MRTR_MAX_ROUNDS, ElicitResult, ElicitationAction, InputRequests, InputRequiredResult};
2use std::time::Duration;
3
4#[derive(Debug)]
5pub struct MrtrState {
6    timeout: Duration,
7    input_request_rounds: usize,
8    next_backoff: Duration,
9    state_only_waited: Duration,
10    user_cancelled: bool,
11}
12
13#[derive(Debug, PartialEq)]
14pub enum MrtrAction {
15    /// Sleep for the backoff, then retry with the echoed request state.
16    Poll { backoff: Duration, request_state: String },
17    /// Dispatch each input request to the user, recording every response via
18    /// [`MrtrState::record_response`], then retry with the responses.
19    Elicit { input_requests: InputRequests, request_state: Option<String> },
20    /// Fail the tool call.
21    Abort(AbortReason),
22}
23
24#[derive(Debug, PartialEq, Eq)]
25pub enum AbortReason {
26    EmptyInputRequired,
27    PollingBudgetExhausted,
28    RePromptAfterCancel,
29    InputRoundsExceeded,
30}
31
32impl MrtrState {
33    pub fn new(timeout: Duration) -> Self {
34        Self {
35            timeout,
36            input_request_rounds: 0,
37            next_backoff: BASE_BACKOFF,
38            state_only_waited: Duration::ZERO,
39            user_cancelled: false,
40        }
41    }
42
43    pub fn tick(&mut self, input_required: InputRequiredResult) -> MrtrAction {
44        let input_requests = input_required.input_requests.filter(|requests| !requests.is_empty());
45        match (input_requests, input_required.request_state) {
46            (None, None) => MrtrAction::Abort(AbortReason::EmptyInputRequired),
47            (None, Some(request_state)) => {
48                let backoff = self.next_backoff;
49                if self.state_only_waited + backoff > self.timeout {
50                    MrtrAction::Abort(AbortReason::PollingBudgetExhausted)
51                } else {
52                    self.state_only_waited += backoff;
53                    self.next_backoff = (backoff * 2).min(MAX_BACKOFF);
54                    MrtrAction::Poll { backoff, request_state }
55                }
56            }
57            (Some(input_requests), request_state) => {
58                if self.user_cancelled {
59                    MrtrAction::Abort(AbortReason::RePromptAfterCancel)
60                } else if self.input_request_rounds == DEFAULT_MRTR_MAX_ROUNDS {
61                    MrtrAction::Abort(AbortReason::InputRoundsExceeded)
62                } else {
63                    self.input_request_rounds += 1;
64                    self.next_backoff = BASE_BACKOFF;
65                    self.state_only_waited = Duration::ZERO;
66                    MrtrAction::Elicit { input_requests, request_state }
67                }
68            }
69        }
70    }
71
72    /// Record an elicitation response so a user cancellation makes the next input-requesting
73    /// [`MrtrState::tick`] abort instead of re-prompting.
74    pub fn record_response(&mut self, response: &ElicitResult) {
75        self.user_cancelled |= response.action == ElicitationAction::Cancel;
76    }
77}
78
79impl AbortReason {
80    pub fn message(&self, server_name: &str, timeout: Duration) -> String {
81        match self {
82            Self::EmptyInputRequired => {
83                format!("Server '{server_name}' requested input without any input requests or request state")
84            }
85            Self::PollingBudgetExhausted => {
86                format!("Server '{server_name}' did not complete within {}s of state-only polling", timeout.as_secs())
87            }
88            Self::RePromptAfterCancel => {
89                format!("Input requested by server '{server_name}' was cancelled by the user")
90            }
91            Self::InputRoundsExceeded => {
92                format!("Server '{server_name}' did not complete within {DEFAULT_MRTR_MAX_ROUNDS} MRTR input rounds")
93            }
94        }
95    }
96}
97
98const BASE_BACKOFF: Duration = Duration::from_millis(50);
99const MAX_BACKOFF: Duration = Duration::from_millis(1600);
100
101#[cfg(test)]
102mod tests {
103    use super::*;
104    use rmcp::model::{ElicitRequest, ElicitRequestParams, InputRequest};
105
106    const TIMEOUT: Duration = Duration::from_secs(1);
107
108    #[test]
109    fn empty_input_required_aborts() {
110        let mut test = mrtr();
111        let trace = test.run(|_| input_required().build());
112        assert!(matches!(trace.last(), Some(MrtrAction::Abort(AbortReason::EmptyInputRequired))));
113    }
114
115    #[test]
116    fn state_only_polls_with_growing_backoff_until_the_budget_is_spent() {
117        let timeout = Duration::from_secs(4);
118        let mut test = mrtr().with_timeout(timeout);
119        let trace = test.run(|_| input_required().with_request_state("s").build());
120
121        assert_eq!(
122            trace.as_slice(),
123            &[
124                MrtrAction::Poll { backoff: Duration::from_millis(50), request_state: "s".to_string() },
125                MrtrAction::Poll { backoff: Duration::from_millis(100), request_state: "s".to_string() },
126                MrtrAction::Poll { backoff: Duration::from_millis(200), request_state: "s".to_string() },
127                MrtrAction::Poll { backoff: Duration::from_millis(400), request_state: "s".to_string() },
128                MrtrAction::Poll { backoff: Duration::from_millis(800), request_state: "s".to_string() },
129                MrtrAction::Poll { backoff: Duration::from_millis(1600), request_state: "s".to_string() },
130                MrtrAction::Abort(AbortReason::PollingBudgetExhausted),
131            ]
132        );
133    }
134
135    #[test]
136    fn an_input_round_refreshes_the_polling_budget() {
137        let mut test = mrtr();
138        while matches!(test.rounds.tick(input_required().with_request_state("s").build()), MrtrAction::Poll { .. }) {}
139        test.elicit_round();
140
141        let decision = test.rounds.tick(input_required().with_request_state("s").build());
142        assert!(
143            matches!(decision, MrtrAction::Poll { backoff, .. } if backoff == BASE_BACKOFF),
144            "budget and backoff should reset after an input round, got {decision:?}"
145        );
146    }
147
148    #[test]
149    fn input_rounds_abort_at_the_cap() {
150        let mut test = mrtr();
151        let actions = test.run(|_| input_required().with_form().build());
152        assert!(
153            actions.as_slice()[..DEFAULT_MRTR_MAX_ROUNDS]
154                .iter()
155                .all(|decision| matches!(decision, MrtrAction::Elicit { .. }))
156        );
157        assert!(matches!(actions.last(), Some(MrtrAction::Abort(AbortReason::InputRoundsExceeded))));
158    }
159
160    #[test]
161    fn a_cancelled_round_aborts_the_next_prompt_but_not_state_only_polling() {
162        let mut test = mrtr().answering_with(ElicitationAction::Cancel);
163        let actions = test.run(|_| input_required().with_form().build());
164        assert!(matches!(
165            actions.as_slice(),
166            [MrtrAction::Elicit { .. }, MrtrAction::Abort(AbortReason::RePromptAfterCancel)]
167        ));
168
169        let poll = test.rounds.tick(input_required().with_request_state("s").build());
170        assert!(matches!(poll, MrtrAction::Poll { .. }), "the server may still finish up, got {poll:?}");
171    }
172
173    #[test]
174    fn an_accepted_round_allows_the_next_prompt() {
175        let mut test = mrtr().answering_with(ElicitationAction::Accept);
176        let elicitation_count = std::cell::Cell::new(0);
177        let actions = test.run_until(
178            |_| {
179                elicitation_count.set(elicitation_count.get() + 1);
180                input_required().with_form().build()
181            },
182            |decision| matches!(decision, MrtrAction::Elicit { .. }) && elicitation_count.get() == 2,
183        );
184
185        assert!(matches!(actions.as_slice(), [MrtrAction::Elicit { .. }, MrtrAction::Elicit { .. }]));
186    }
187
188    struct MrtrTest {
189        rounds: MrtrState,
190        response: ElicitResult,
191    }
192
193    fn mrtr() -> MrtrTest {
194        MrtrTest::default()
195    }
196
197    impl Default for MrtrTest {
198        fn default() -> Self {
199            Self { rounds: MrtrState::new(TIMEOUT), response: ElicitResult::new(ElicitationAction::Accept) }
200        }
201    }
202
203    impl MrtrTest {
204        fn with_timeout(mut self, timeout: Duration) -> Self {
205            self.rounds = MrtrState::new(timeout);
206            self
207        }
208
209        fn answering_with(mut self, action: ElicitationAction) -> Self {
210            self.response = ElicitResult::new(action);
211            self
212        }
213
214        fn run_until<T, U>(&mut self, mut next: T, mut terminal: U) -> Vec<MrtrAction>
215        where
216            T: FnMut(Option<&MrtrAction>) -> InputRequiredResult,
217            U: FnMut(&MrtrAction) -> bool,
218        {
219            let mut actions = Vec::new();
220            loop {
221                let action = self.rounds.tick(next(actions.last()));
222                if matches!(action, MrtrAction::Elicit { .. }) {
223                    self.rounds.record_response(&self.response);
224                }
225                let is_terminal = terminal(&action);
226                actions.push(action);
227                if is_terminal {
228                    return actions;
229                }
230            }
231        }
232
233        fn run<T>(&mut self, next: T) -> Vec<MrtrAction>
234        where
235            T: FnMut(Option<&MrtrAction>) -> InputRequiredResult,
236        {
237            self.run_until(next, |decision| matches!(decision, MrtrAction::Abort(_)))
238        }
239
240        fn elicit_round(&mut self) {
241            let actions = self.run_until(
242                |_| input_required().with_form().build(),
243                |decision| matches!(decision, MrtrAction::Elicit { .. }),
244            );
245            assert!(matches!(actions.last(), Some(MrtrAction::Elicit { .. })));
246        }
247    }
248
249    fn input_required() -> InputRequiredBuilder {
250        InputRequiredBuilder::default()
251    }
252
253    #[derive(Default)]
254    struct InputRequiredBuilder {
255        input_requests: Option<InputRequests>,
256        request_state: Option<String>,
257    }
258
259    impl InputRequiredBuilder {
260        fn with_form(mut self) -> Self {
261            let params = ElicitRequestParams::FormElicitationParams {
262                meta: None,
263                message: "m".to_string(),
264                requested_schema: serde_json::from_value(serde_json::json!({
265                    "type": "object",
266                    "properties": {}
267                }))
268                .unwrap(),
269            };
270            let mut requests = InputRequests::new();
271            requests.insert("k".to_string(), InputRequest::Elicitation(ElicitRequest::new(params)));
272            self.input_requests = Some(requests);
273            self
274        }
275
276        fn with_request_state(mut self, state: &str) -> Self {
277            self.request_state = Some(state.to_string());
278            self
279        }
280
281        fn build(self) -> InputRequiredResult {
282            InputRequiredResult::new(self.input_requests, self.request_state)
283        }
284    }
285}
mcp_utils/client/mrtr.rs

mcp_utils/client/
mrtr.rs
