1use std::{
2 collections::HashMap,
3 sync::{
4 atomic::{AtomicU64, Ordering},
5 Arc, Mutex,
6 },
7 time::{Duration, Instant},
8};
9
10use rmcp::{
11 model::{CallToolResult, ContentBlock},
12 ErrorData,
13};
14use serde_json::{json, Map, Value};
15
16pub const RESPONSE_OFFSET_PARAM: &str = "_response_offset";
17pub const RESPONSE_PAGE_BYTES_PARAM: &str = "_response_page_bytes";
18pub const RESPONSE_CURSOR_PARAM: &str = "_response_cursor";
19pub const DEFAULT_ACTION_DISCRIMINATOR_FIELD: &str = "_action";
20pub const DEFAULT_MAX_RESPONSE_BYTES: usize = 40_000;
21pub const DEFAULT_RESPONSE_PAGE_BYTES: usize = 16_000;
22pub const MAX_RESPONSE_PAGE_BYTES: usize = 16_000;
23pub const MAX_RESPONSE_CURSOR_BYTES: usize = 256;
24
25#[derive(Clone, Copy, Debug, PartialEq, Eq)]
26pub struct ResponsePagingOptions {
27 pub max_response_bytes: usize,
28 pub action_discriminator_field: &'static str,
29}
30
31impl Default for ResponsePagingOptions {
32 fn default() -> Self {
33 Self {
34 max_response_bytes: DEFAULT_MAX_RESPONSE_BYTES,
35 action_discriminator_field: DEFAULT_ACTION_DISCRIMINATOR_FIELD,
36 }
37 }
38}
39
40#[derive(Clone, Default)]
41pub struct ResponsePageStore {
42 inner: Arc<ResponsePageStoreInner>,
43}
44
45#[derive(Default)]
46struct ResponsePageStoreInner {
47 counter: AtomicU64,
48 entries: Mutex<HashMap<String, CachedResponsePage>>,
49}
50
51struct CachedResponsePage {
52 serialized: String,
53 expires_at: Instant,
54}
55
56impl ResponsePageStore {
57 const TTL: Duration = Duration::from_secs(300);
58
59 pub fn insert(&self, serialized: String) -> String {
60 self.prune_expired();
61 let id = self.inner.counter.fetch_add(1, Ordering::Relaxed) + 1;
62 let cursor = format!("rsp_{id:x}");
63 let entry = CachedResponsePage {
64 serialized,
65 expires_at: Instant::now() + Self::TTL,
66 };
67 self.inner
68 .entries
69 .lock()
70 .expect("response page store mutex should not be poisoned")
71 .insert(cursor.clone(), entry);
72 cursor
73 }
74
75 pub fn get(&self, cursor: &str) -> Option<String> {
76 self.prune_expired();
77 self.inner
78 .entries
79 .lock()
80 .expect("response page store mutex should not be poisoned")
81 .get(cursor)
82 .map(|entry| entry.serialized.clone())
83 }
84
85 fn prune_expired(&self) {
86 let now = Instant::now();
87 self.inner
88 .entries
89 .lock()
90 .expect("response page store mutex should not be poisoned")
91 .retain(|_, entry| entry.expires_at > now);
92 }
93}
94
95#[derive(Clone, Debug, PartialEq, Eq)]
96pub struct ResponsePageRequest {
97 pub cursor: Option<String>,
98 pub offset: usize,
99 pub page_bytes: usize,
100}
101
102impl Default for ResponsePageRequest {
103 fn default() -> Self {
104 Self {
105 cursor: None,
106 offset: 0,
107 page_bytes: DEFAULT_RESPONSE_PAGE_BYTES,
108 }
109 }
110}
111
112impl ResponsePageRequest {
113 pub fn cursor(&self) -> Option<&str> {
114 self.cursor.as_deref()
115 }
116}
117
118pub fn response_page_request(
119 args: Option<&Map<String, Value>>,
120) -> Result<ResponsePageRequest, ErrorData> {
121 let Some(args) = args else {
122 return Ok(ResponsePageRequest::default());
123 };
124 let cursor = optional_string_arg(args, RESPONSE_CURSOR_PARAM)?;
125 let offset = optional_usize_arg(args, RESPONSE_OFFSET_PARAM)?.unwrap_or(0);
126 let page_bytes = optional_usize_arg(args, RESPONSE_PAGE_BYTES_PARAM)?
127 .unwrap_or(DEFAULT_RESPONSE_PAGE_BYTES)
128 .min(MAX_RESPONSE_PAGE_BYTES);
129 if page_bytes == 0 {
130 return Err(ErrorData::invalid_params(
131 format!("{RESPONSE_PAGE_BYTES_PARAM} must be greater than zero"),
132 Some(json!({
133 "kind": "mcp_protocol_error",
134 "schema_version": 1,
135 "code": "invalid_response_page_bytes",
136 "field": RESPONSE_PAGE_BYTES_PARAM,
137 "retryable": true,
138 "remediation": format!("Omit {RESPONSE_PAGE_BYTES_PARAM} or pass an integer from 1 to {MAX_RESPONSE_PAGE_BYTES}."),
139 })),
140 ));
141 }
142 if offset > 0 && cursor.is_none() {
143 return Err(ErrorData::invalid_params(
144 format!("{RESPONSE_CURSOR_PARAM} is required when {RESPONSE_OFFSET_PARAM} is set"),
145 Some(json!({
146 "kind": "mcp_protocol_error",
147 "schema_version": 1,
148 "code": "missing_response_cursor",
149 "field": RESPONSE_CURSOR_PARAM,
150 "retryable": true,
151 "remediation": format!("Use the {RESPONSE_CURSOR_PARAM} value returned by the previous mcp_response_page continuation."),
152 })),
153 ));
154 }
155 Ok(ResponsePageRequest {
156 cursor,
157 offset,
158 page_bytes,
159 })
160}
161
162fn optional_string_arg(
163 args: &Map<String, Value>,
164 field: &str,
165) -> Result<Option<String>, ErrorData> {
166 let Some(value) = args.get(field) else {
167 return Ok(None);
168 };
169 let Some(value) = value.as_str() else {
170 return Err(ErrorData::invalid_params(
171 format!("{field} must be a string"),
172 Some(json!({
173 "kind": "mcp_protocol_error",
174 "schema_version": 1,
175 "code": "invalid_response_cursor",
176 "field": field,
177 "retryable": true,
178 "remediation": format!("Pass {field} exactly as returned by the previous mcp_response_page continuation."),
179 })),
180 ));
181 };
182 if field == RESPONSE_CURSOR_PARAM && value.len() > MAX_RESPONSE_CURSOR_BYTES {
183 return Err(ErrorData::invalid_params(
184 format!("{field} exceeded {MAX_RESPONSE_CURSOR_BYTES} bytes"),
185 Some(json!({
186 "kind": "mcp_protocol_error",
187 "schema_version": 1,
188 "code": "response_cursor_too_long",
189 "field": field,
190 "retryable": true,
191 "remediation": format!("Pass {field} exactly as returned by the previous mcp_response_page continuation."),
192 })),
193 ));
194 }
195 Ok(Some(value.to_owned()))
196}
197
198fn optional_usize_arg(args: &Map<String, Value>, field: &str) -> Result<Option<usize>, ErrorData> {
199 let Some(value) = args.get(field) else {
200 return Ok(None);
201 };
202 let Some(value) = value.as_u64() else {
203 return Err(ErrorData::invalid_params(
204 format!("{field} must be an unsigned integer"),
205 Some(json!({
206 "kind": "mcp_protocol_error",
207 "schema_version": 1,
208 "code": "invalid_response_page_arg",
209 "field": field,
210 "retryable": true,
211 "remediation": format!("Pass {field} as a non-negative integer."),
212 })),
213 ));
214 };
215 usize::try_from(value).map(Some).map_err(|_| {
216 ErrorData::invalid_params(
217 format!("{field} is too large"),
218 Some(json!({
219 "kind": "mcp_protocol_error",
220 "schema_version": 1,
221 "code": "response_page_arg_too_large",
222 "field": field,
223 "retryable": true,
224 "remediation": format!("Pass a smaller {field} value."),
225 })),
226 )
227 })
228}
229
230pub fn strip_response_page_params(arguments: &mut Value) {
231 let Some(arguments) = arguments.as_object_mut() else {
232 return;
233 };
234 arguments.remove(RESPONSE_OFFSET_PARAM);
235 arguments.remove(RESPONSE_PAGE_BYTES_PARAM);
236 arguments.remove(RESPONSE_CURSOR_PARAM);
237}
238
239pub fn tool_result_from_json(
240 mut value: Value,
241 response_pages: &ResponsePageStore,
242 page_request: ResponsePageRequest,
243 options: ResponsePagingOptions,
244 tool: &str,
245 action: Option<&str>,
246 continuation_args: Option<&Map<String, Value>>,
247) -> Result<CallToolResult, ErrorData> {
248 if let Some(cursor) = page_request.cursor.clone() {
249 return tool_result_from_cached_page(
250 response_pages,
251 &cursor,
252 page_request,
253 options,
254 tool,
255 action,
256 );
257 }
258 add_action_discriminator(&mut value, action, options.action_discriminator_field);
259
260 let text = serde_json::to_string(&value)
262 .map_err(|e| ErrorData::internal_error(format!("serialization error: {e}"), None))?;
263 if text.len() <= options.max_response_bytes && page_request.offset == 0 {
264 let mut result = CallToolResult::structured(value);
265 result.content = vec![ContentBlock::text(text)];
266 return Ok(result);
267 }
268
269 let cursor = response_pages.insert(text.clone());
270 let payload = response_page_payload(
271 &text,
272 page_request,
273 options,
274 tool,
275 action,
276 continuation_args,
277 Some(&cursor),
278 );
279 let text = serde_json::to_string(&payload)
280 .map_err(|e| ErrorData::internal_error(format!("serialization error: {e}"), None))?;
281 let mut result = CallToolResult::structured(payload);
282 result.content = vec![ContentBlock::text(text)];
283 Ok(result)
284}
285
286fn add_action_discriminator(
287 value: &mut Value,
288 action: Option<&str>,
289 action_discriminator_field: &str,
290) {
291 let (Some(action), Some(object)) = (action, value.as_object_mut()) else {
292 return;
293 };
294 object.insert(
295 action_discriminator_field.to_owned(),
296 Value::String(action.to_owned()),
297 );
298}
299
300pub fn tool_result_from_cached_page(
301 response_pages: &ResponsePageStore,
302 cursor: &str,
303 page_request: ResponsePageRequest,
304 options: ResponsePagingOptions,
305 tool: &str,
306 action: Option<&str>,
307) -> Result<CallToolResult, ErrorData> {
308 let Some(serialized) = response_pages.get(cursor) else {
309 return Err(ErrorData::invalid_params(
310 "response cursor not found or expired",
311 Some(json!({
312 "kind": "mcp_protocol_error",
313 "schema_version": 1,
314 "code": "response_cursor_not_found",
315 "field": RESPONSE_CURSOR_PARAM,
316 "retryable": true,
317 "remediation": "Re-run the original tool call to create a fresh response cursor.",
318 })),
319 ));
320 };
321 let payload = response_page_payload(
322 &serialized,
323 page_request,
324 options,
325 tool,
326 action,
327 None,
328 Some(cursor),
329 );
330 let text = serde_json::to_string(&payload)
331 .map_err(|e| ErrorData::internal_error(format!("serialization error: {e}"), None))?;
332 let mut result = CallToolResult::structured(payload);
333 result.content = vec![ContentBlock::text(text)];
334 Ok(result)
335}
336
337fn response_page_payload(
338 serialized: &str,
339 page_request: ResponsePageRequest,
340 options: ResponsePagingOptions,
341 tool: &str,
342 action: Option<&str>,
343 continuation_args: Option<&Map<String, Value>>,
344 cursor: Option<&str>,
345) -> Value {
346 let (offset, content, next_offset, has_more) =
347 response_page_slice(serialized, page_request.offset, page_request.page_bytes);
348 let continuation = has_more.then(|| {
349 let arguments = continuation_arguments_with_page(
350 continuation_args,
351 action,
352 next_offset,
353 page_request.page_bytes,
354 cursor,
355 );
356 json!({
357 "tool": tool,
358 "arguments": arguments,
359 "note": "Call the same tool with the same original arguments plus these reserved continuation arguments.",
360 })
361 });
362
363 json!({
364 "kind": "mcp_response_page",
365 "schema_version": 1,
366 "code": "response_page",
367 "message": "Tool response was returned as a scrollable serialized JSON page.",
368 "truncated": false,
369 "serialized_bytes": serialized.len(),
370 "max_response_bytes": options.max_response_bytes,
371 "content_format": "application/json-fragment",
372 "content": content,
373 "page": {
374 "offset": offset,
375 "page_bytes": page_request.page_bytes,
376 "next_offset": next_offset,
377 "has_more": has_more,
378 },
379 "continuation": continuation,
380 })
381}
382
383fn continuation_arguments_with_page(
384 arguments: Option<&Map<String, Value>>,
385 action: Option<&str>,
386 next_offset: usize,
387 page_bytes: usize,
388 cursor: Option<&str>,
389) -> Value {
390 let mut output = arguments.cloned().unwrap_or_default();
391 if !output.contains_key("action") {
392 output.insert(
393 "action".to_owned(),
394 action.map(Value::from).unwrap_or(Value::Null),
395 );
396 }
397 if let Some(cursor) = cursor {
398 output.insert(RESPONSE_CURSOR_PARAM.to_owned(), json!(cursor));
399 }
400 output.insert(RESPONSE_OFFSET_PARAM.to_owned(), json!(next_offset));
401 output.insert(RESPONSE_PAGE_BYTES_PARAM.to_owned(), json!(page_bytes));
402 Value::Object(output)
403}
404
405fn response_page_slice(
406 serialized: &str,
407 requested_offset: usize,
408 page_bytes: usize,
409) -> (usize, &str, usize, bool) {
410 let mut offset = requested_offset.min(serialized.len());
411 while offset < serialized.len() && !serialized.is_char_boundary(offset) {
412 offset += 1;
413 }
414
415 let mut end = offset.saturating_add(page_bytes).min(serialized.len());
416 while end > offset && !serialized.is_char_boundary(end) {
417 end -= 1;
418 }
419
420 (
421 offset,
422 &serialized[offset..end],
423 end,
424 end < serialized.len(),
425 )
426}
427
428#[cfg(test)]
429#[path = "response_paging_tests.rs"]
430mod tests;