1use 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#[derive(Serialize, Deserialize)]
24struct Payload {
25 offset: u64,
26}
27
28#[derive(Debug, Clone, Copy, PartialEq, Eq)]
35pub struct Cursor {
36 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#[derive(Debug, Clone, Copy)]
79pub struct Pager {
80 offset: u64,
81 size: u16,
82}
83
84impl Pager {
85 #[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 #[must_use]
100 pub fn offset(&self) -> i64 {
101 i64::try_from(self.offset).unwrap_or(i64::MAX)
102 }
103
104 #[must_use]
109 pub fn limit(&self) -> i64 {
110 i64::from(self.size) + 1
111 }
112
113 #[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
131fn 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
138fn 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}