Skip to main content

dbmcp_server/
pagination.rs

1//! Opaque cursor pagination helpers for MCP tool responses.
2//!
3//! Implements the cursor conventions described in the MCP specification
4//! (`cursor` / `nextCursor`, opaque tokens). Cursors are base64url-encoded
5//! JSON objects of the form `{"offset": <offset>}`. The encoding is an
6//! implementation detail — clients MUST treat cursors as opaque strings.
7//!
8//! Tool requests declare their cursor as [`Option<Cursor>`]; the custom
9//! `Serialize`/`Deserialize` impls encode to / decode from the base64url
10//! wire string, so handlers never see the raw cursor string. Decode
11//! errors surface as serde errors, which rmcp reports as a tool result
12//! with `isError: true` when parsing tool arguments.
13
14use std::borrow::Cow;
15
16use base64::Engine;
17use base64::engine::general_purpose::URL_SAFE_NO_PAD;
18use rmcp::schemars::{JsonSchema, Schema, SchemaGenerator, json_schema};
19use serde::de::{Error as DeError, Unexpected};
20use serde::{Deserialize, Deserializer, Serialize, Serializer};
21
22/// Wire-format payload carried inside an encoded [`Cursor`].
23#[derive(Serialize, Deserialize)]
24struct Payload {
25    offset: u64,
26}
27
28/// Opaque pagination cursor carried on paginated tool requests / responses.
29///
30/// On the wire this serialises as a URL-safe base64 string; in Rust code
31/// it is a typed offset. Construct directly (`Cursor { offset: 100 }`) or
32/// receive via serde. The JSON representation is intentionally opaque to
33/// clients.
34#[derive(Debug, Clone, Copy, PartialEq, Eq)]
35pub struct Cursor {
36    /// Zero-based row offset into the backend's sorted item list.
37    pub offset: u64,
38}
39
40impl Serialize for Cursor {
41    fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
42        serializer.serialize_str(&encode_cursor(self.offset))
43    }
44}
45
46impl<'de> Deserialize<'de> for Cursor {
47    fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
48        let raw = <Cow<'de, str>>::deserialize(deserializer)?;
49        decode_cursor(&raw)
50            .map(|offset| Self { offset })
51            .map_err(|msg| D::Error::invalid_value(Unexpected::Str(&raw), &msg))
52    }
53}
54
55impl JsonSchema for Cursor {
56    fn inline_schema() -> bool {
57        true
58    }
59
60    fn schema_name() -> Cow<'static, str> {
61        "Cursor".into()
62    }
63
64    fn json_schema(_: &mut SchemaGenerator) -> Schema {
65        json_schema!({
66            "type": "string",
67            "description": "Opaque pagination cursor. Echo the `nextCursor` from a prior response; do not parse or modify."
68        })
69    }
70}
71
72/// A single page request resolved from an optional cursor and a page size.
73///
74/// Paginated tools follow a fetch-one-extra pattern: query `size + 1` rows,
75/// then call [`Self::paginate`] to trim the extra row and emit a next cursor
76/// when present. Construct with [`Self::new`]; read [`Self::offset`] /
77/// [`Self::limit`] when building the SQL statement.
78#[derive(Debug, Clone, Copy)]
79pub struct Pager {
80    offset: u64,
81    size: u16,
82}
83
84impl Pager {
85    /// Builds a page request from an optional cursor and the configured page size.
86    #[must_use]
87    pub fn new(cursor: Option<Cursor>, size: u16) -> Self {
88        Self {
89            offset: cursor.map_or(0, |c| c.offset),
90            size,
91        }
92    }
93
94    /// Row offset at which this page starts.
95    ///
96    /// Returned as `i64` so the value is directly bindable to sqlx
97    /// LIMIT/OFFSET placeholders across every backend. Saturates at
98    /// [`i64::MAX`] for cursor offsets that exceed the signed range.
99    #[must_use]
100    pub fn offset(&self) -> i64 {
101        i64::try_from(self.offset).unwrap_or(i64::MAX)
102    }
103
104    /// Row count to fetch from the backend (`size + 1`, for lookahead).
105    ///
106    /// Returned as `i64` for direct sqlx binding. Capped at `u16::MAX + 1`
107    /// since `size` is constructed from a `u16`.
108    #[must_use]
109    pub fn limit(&self) -> i64 {
110        i64::from(self.size) + 1
111    }
112
113    /// Trims over-fetched items to `size` and derives the next cursor.
114    ///
115    /// When `items.len()` exceeds `size`, the tail is dropped and a cursor
116    /// pointing at the next offset is returned. Otherwise the items are
117    /// returned unchanged with `None`.
118    #[must_use]
119    pub fn paginate<T>(&self, mut items: Vec<T>) -> (Vec<T>, Option<Cursor>) {
120        let size = usize::from(self.size);
121        if items.len() > size {
122            items.truncate(size);
123            let offset = self.offset + u64::from(self.size);
124            (items, Some(Cursor { offset }))
125        } else {
126            (items, None)
127        }
128    }
129}
130
131/// Encodes a zero-based page offset as an opaque cursor string.
132fn encode_cursor(offset: u64) -> String {
133    let payload = Payload { offset };
134    let json = serde_json::to_vec(&payload).expect("Payload is infallible to serialize");
135    URL_SAFE_NO_PAD.encode(&json)
136}
137
138/// Decodes a cursor string produced by [`encode_cursor`] back to its offset.
139fn decode_cursor(raw: &str) -> Result<u64, &'static str> {
140    let bytes = URL_SAFE_NO_PAD
141        .decode(raw.as_bytes())
142        .map_err(|_| "invalid pagination cursor: not valid base64")?;
143    let payload: Payload =
144        serde_json::from_slice(&bytes).map_err(|_| "invalid pagination cursor: payload is malformed")?;
145    Ok(payload.offset)
146}
147
148#[cfg(test)]
149mod tests {
150    use base64::Engine;
151    use base64::engine::general_purpose::URL_SAFE_NO_PAD;
152    use serde_json::{Value, json};
153
154    use super::{Cursor, Pager, decode_cursor, encode_cursor};
155
156    #[test]
157    fn encode_decode_round_trips_representative_offsets() {
158        for offset in [0u64, 1, 99, 100, 101, 12_345, u64::MAX / 2, u64::MAX] {
159            let cursor = encode_cursor(offset);
160            let decoded = decode_cursor(&cursor).expect("valid cursor should decode");
161            assert_eq!(decoded, offset);
162        }
163    }
164
165    #[test]
166    fn encoded_cursor_is_url_safe_base64() {
167        let cursor = encode_cursor(100);
168        assert!(!cursor.contains('+'));
169        assert!(!cursor.contains('/'));
170        assert!(!cursor.contains('='));
171    }
172
173    #[test]
174    fn cursor_serializes_as_base64_string() {
175        let cursor = Cursor { offset: 100 };
176        let value = serde_json::to_value(cursor).unwrap();
177        let Value::String(s) = value else {
178            panic!("expected string, got {value:?}");
179        };
180        assert_eq!(decode_cursor(&s).unwrap(), 100);
181    }
182
183    #[test]
184    fn cursor_deserializes_from_valid_base64() {
185        let raw = encode_cursor(42);
186        let cursor: Cursor = serde_json::from_value(Value::String(raw)).unwrap();
187        assert_eq!(cursor.offset, 42);
188    }
189
190    #[test]
191    fn cursor_deserialization_round_trips_through_serde_json() {
192        let original = Cursor { offset: 7 };
193        let json = serde_json::to_string(&original).unwrap();
194        let back: Cursor = serde_json::from_str(&json).unwrap();
195        assert_eq!(back, original);
196    }
197
198    #[test]
199    fn cursor_deserialization_rejects_non_base64() {
200        let err = serde_json::from_value::<Cursor>(json!("!!!not-base64")).expect_err("should fail");
201        assert!(err.to_string().contains("base64"), "error: {err}");
202    }
203
204    #[test]
205    fn cursor_deserialization_rejects_base64_of_non_json() {
206        let raw = URL_SAFE_NO_PAD.encode(b"not json");
207        let err = serde_json::from_value::<Cursor>(json!(raw)).expect_err("should fail");
208        assert!(err.to_string().contains("malformed"), "error: {err}");
209    }
210
211    #[test]
212    fn cursor_deserialization_rejects_payload_missing_fields() {
213        let raw = URL_SAFE_NO_PAD.encode(b"{}");
214        let err = serde_json::from_value::<Cursor>(json!(raw)).expect_err("should fail");
215        assert!(err.to_string().contains("malformed"), "error: {err}");
216    }
217
218    #[test]
219    fn cursor_deserialization_rejects_negative_offset() {
220        let raw = URL_SAFE_NO_PAD.encode(b"{\"offset\":-1}");
221        let err = serde_json::from_value::<Cursor>(json!(raw)).expect_err("should fail");
222        assert!(err.to_string().contains("malformed"), "error: {err}");
223    }
224
225    #[test]
226    fn encoded_cursor_payload_uses_offset_key() {
227        let raw = serde_json::to_value(Cursor { offset: 100 }).unwrap();
228        let Value::String(s) = raw else {
229            panic!("expected string cursor, got {raw:?}");
230        };
231        let bytes = URL_SAFE_NO_PAD.decode(s.as_bytes()).unwrap();
232        let payload: Value = serde_json::from_slice(&bytes).unwrap();
233        let obj = payload.as_object().expect("payload should be a JSON object");
234        assert_eq!(
235            obj.get("offset").and_then(Value::as_u64),
236            Some(100),
237            "payload should carry offset under the `offset` key: {obj:?}"
238        );
239    }
240
241    #[test]
242    fn page_defaults_to_offset_zero_without_cursor() {
243        let pager = Pager::new(None, 50);
244        assert_eq!(pager.offset(), 0);
245        assert_eq!(pager.limit(), 51);
246    }
247
248    #[test]
249    fn page_inherits_offset_from_cursor() {
250        let pager = Pager::new(Some(Cursor { offset: 200 }), 50);
251        assert_eq!(pager.offset(), 200);
252        assert_eq!(pager.limit(), 51);
253    }
254
255    #[test]
256    fn page_paginate_emits_next_cursor_when_over_fetched() {
257        let pager = Pager::new(None, 3);
258        let (items, next) = pager.paginate(vec!["a", "b", "c", "d"]);
259        assert_eq!(items, ["a", "b", "c"]);
260        assert_eq!(next, Some(Cursor { offset: 3 }));
261    }
262
263    #[test]
264    fn page_paginate_drops_next_cursor_on_exact_fit() {
265        let pager = Pager::new(None, 3);
266        let (items, next) = pager.paginate(vec!["a", "b", "c"]);
267        assert_eq!(items, ["a", "b", "c"]);
268        assert!(next.is_none());
269    }
270
271    #[test]
272    fn page_paginate_drops_next_cursor_on_short_page() {
273        let pager = Pager::new(None, 3);
274        let (items, next) = pager.paginate(vec!["a"]);
275        assert_eq!(items, ["a"]);
276        assert!(next.is_none());
277    }
278
279    #[test]
280    fn page_paginate_drops_next_cursor_on_empty_result() {
281        let pager = Pager::new(Some(Cursor { offset: 99 }), 3);
282        let (items, next) = pager.paginate(Vec::<&str>::new());
283        assert!(items.is_empty());
284        assert!(next.is_none());
285    }
286
287    #[test]
288    fn page_paginate_advances_offset_by_page_size() {
289        let pager = Pager::new(Some(Cursor { offset: 100 }), 50);
290        let items: Vec<u32> = (0..51).collect();
291        let (_, next) = pager.paginate(items);
292        assert_eq!(next, Some(Cursor { offset: 150 }));
293    }
294}