Skip to main content

contextvm_sdk/transport/oversized_transfer/
codec.rs

1//! Pure framing codec: digest, UTF-8 char-aware splitting, and frame building.
2//!
3//! Ports `sdk/src/transport/oversized-transfer/sender.ts`. No transport or I/O —
4//! given a serialized JSON-RPC string this produces the ordered sequence of
5//! `notifications/progress` frames a sender publishes.
6
7use sha2::{Digest, Sha256};
8
9use crate::core::types::JsonRpcNotification;
10
11use super::constants::{ACCEPT_PROGRESS, DEFAULT_CHUNK_SIZE, DIGEST_PREFIX, START_PROGRESS};
12use super::errors::OversizedTransferError;
13use super::frame::{CompletionMode, OversizedFrame};
14
15/// Options controlling how a payload is split into frames.
16#[derive(Debug, Clone)]
17pub struct OversizedSenderOptions {
18    /// The MCP `progressToken` stamped on every frame in the transfer.
19    pub progress_token: String,
20    /// Per-chunk byte size. Defaults to [`DEFAULT_CHUNK_SIZE`].
21    pub chunk_size_bytes: usize,
22    /// Whether to reserve the `accept` slot (progress 2) for a handshake, so the
23    /// first `chunk` starts at progress 3. When `false`, chunks start at 2.
24    pub needs_accept_handshake: bool,
25}
26
27impl OversizedSenderOptions {
28    /// Create options for `progress_token` with default chunk size and no handshake.
29    pub fn new(progress_token: impl Into<String>) -> Self {
30        Self {
31            progress_token: progress_token.into(),
32            chunk_size_bytes: DEFAULT_CHUNK_SIZE,
33            needs_accept_handshake: false,
34        }
35    }
36
37    /// Override the per-chunk byte size.
38    pub fn with_chunk_size(mut self, chunk_size_bytes: usize) -> Self {
39        self.chunk_size_bytes = chunk_size_bytes;
40        self
41    }
42
43    /// Set whether the `accept` handshake slot is reserved.
44    pub fn with_accept_handshake(mut self, needs_accept_handshake: bool) -> Self {
45        self.needs_accept_handshake = needs_accept_handshake;
46        self
47    }
48}
49
50/// The ordered frames produced from a serialized payload.
51#[derive(Debug, Clone)]
52pub struct BuiltOversizedFrames {
53    /// The opening `start` frame (progress [`START_PROGRESS`]).
54    pub start: JsonRpcNotification,
55    /// The `chunk` frames in ascending `progress` order.
56    pub chunks: Vec<JsonRpcNotification>,
57    /// The closing `end` frame.
58    pub end: JsonRpcNotification,
59}
60
61impl BuiltOversizedFrames {
62    /// Total number of frames (`start` + chunks + `end`).
63    pub fn frame_count(&self) -> usize {
64        self.chunks.len() + 2
65    }
66
67    /// Consume into a flat vector in canonical send order: `start`, chunks…, `end`.
68    pub fn into_ordered(self) -> Vec<JsonRpcNotification> {
69        let mut frames = Vec::with_capacity(self.frame_count());
70        frames.push(self.start);
71        frames.extend(self.chunks);
72        frames.push(self.end);
73        frames
74    }
75}
76
77/// UTF-8 byte length of a string (Rust strings are already UTF-8).
78pub fn utf8_byte_len(value: &str) -> usize {
79    value.len()
80}
81
82/// Compute the CEP-22 digest of `value`: `"sha256:"` + lowercase hex of the
83/// SHA-256 of its UTF-8 bytes.
84pub fn sha256_digest(value: &str) -> String {
85    let mut hasher = Sha256::new();
86    hasher.update(value.as_bytes());
87    format!("{DIGEST_PREFIX}{}", hex::encode(hasher.finalize()))
88}
89
90/// Split `value` into fragments each at most `max_bytes` UTF-8 bytes, never
91/// breaking a multibyte codepoint across a boundary.
92///
93/// Concatenating the fragments in order reproduces `value` exactly. Returns an
94/// error if `max_bytes` is zero or a single character is larger than `max_bytes`.
95pub fn split_string_by_byte_size(
96    value: &str,
97    max_bytes: usize,
98) -> Result<Vec<String>, OversizedTransferError> {
99    if max_bytes == 0 {
100        return Err(OversizedTransferError::Policy(
101            "chunk_size_bytes must be greater than zero".to_string(),
102        ));
103    }
104
105    let mut chunks: Vec<String> = Vec::new();
106    let mut current = String::new();
107    let mut current_bytes = 0usize;
108
109    for ch in value.chars() {
110        let char_bytes = ch.len_utf8();
111        if char_bytes > max_bytes {
112            return Err(OversizedTransferError::Policy(format!(
113                "Unable to split message: single character exceeds chunk size ({char_bytes} > {max_bytes})"
114            )));
115        }
116
117        if current_bytes > 0 && current_bytes + char_bytes > max_bytes {
118            chunks.push(std::mem::take(&mut current));
119            current_bytes = 0;
120        }
121
122        current.push(ch);
123        current_bytes += char_bytes;
124    }
125
126    if !current.is_empty() {
127        chunks.push(current);
128    }
129
130    Ok(chunks)
131}
132
133/// Build the ordered `start` / `chunk…` / `end` frames for `serialized`.
134///
135/// Serializes once at the call site (the caller passes the exact JSON-RPC
136/// string), then: computes the digest over that string, char-aware splits it,
137/// and lays the chunks out on `progress` slots. When
138/// [`needs_accept_handshake`](OversizedSenderOptions::needs_accept_handshake) is
139/// set, progress 2 is reserved for the receiver's `accept` and chunks begin at 3.
140pub fn build_oversized_frames(
141    serialized: &str,
142    options: &OversizedSenderOptions,
143) -> crate::core::error::Result<BuiltOversizedFrames> {
144    let total_bytes = utf8_byte_len(serialized) as u64;
145    let digest = sha256_digest(serialized);
146
147    let text_chunks = split_string_by_byte_size(serialized, options.chunk_size_bytes)?;
148    let total_chunks = text_chunks.len() as u64;
149
150    // No-handshake: chunks reuse the reserved accept slot (progress 2).
151    // Handshake: accept occupies slot 2, so chunks begin at slot 3.
152    let chunk_base_progress = if options.needs_accept_handshake {
153        ACCEPT_PROGRESS + 1
154    } else {
155        ACCEPT_PROGRESS
156    };
157
158    let token = options.progress_token.as_str();
159
160    let start = OversizedFrame::Start {
161        completion_mode: CompletionMode::Render,
162        digest,
163        total_bytes,
164        total_chunks,
165    }
166    .into_progress_notification(token, START_PROGRESS, Some("starting oversized transfer"))?;
167
168    let mut chunks = Vec::with_capacity(text_chunks.len());
169    for (i, data) in text_chunks.into_iter().enumerate() {
170        let progress = chunk_base_progress + i as u64;
171        let frame = OversizedFrame::Chunk { data };
172        chunks.push(frame.into_progress_notification(token, progress, None)?);
173    }
174
175    let end = OversizedFrame::End.into_progress_notification(
176        token,
177        chunk_base_progress + total_chunks,
178        Some("oversized transfer complete"),
179    )?;
180
181    Ok(BuiltOversizedFrames { start, chunks, end })
182}
183
184#[cfg(test)]
185mod tests {
186    use super::*;
187    use serde_json::Value;
188
189    fn frame_type(notification: &JsonRpcNotification) -> String {
190        notification.params.as_ref().unwrap()["cvm"]["frameType"]
191            .as_str()
192            .unwrap()
193            .to_string()
194    }
195
196    fn progress(notification: &JsonRpcNotification) -> u64 {
197        notification.params.as_ref().unwrap()["progress"]
198            .as_u64()
199            .unwrap()
200    }
201
202    // ── digest ──────────────────────────────────────────────────────
203
204    #[test]
205    fn sha256_digest_known_vectors() {
206        // Reference SHA-256 vectors with the CEP-22 "sha256:" prefix.
207        assert_eq!(
208            sha256_digest(""),
209            "sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"
210        );
211        assert_eq!(
212            sha256_digest("abc"),
213            "sha256:ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad"
214        );
215    }
216
217    #[test]
218    fn sha256_digest_has_prefix_and_lowercase_hex() {
219        let digest = sha256_digest("hello world");
220        assert!(digest.starts_with("sha256:"));
221        let hex_part = &digest["sha256:".len()..];
222        assert_eq!(hex_part.len(), 64);
223        assert!(hex_part
224            .chars()
225            .all(|c| c.is_ascii_hexdigit() && !c.is_ascii_uppercase()));
226    }
227
228    #[test]
229    fn utf8_byte_len_counts_bytes_not_chars() {
230        assert_eq!(utf8_byte_len("abc"), 3);
231        assert_eq!(utf8_byte_len("é"), 2);
232        assert_eq!(utf8_byte_len("🦀"), 4);
233        assert_eq!(utf8_byte_len("a🦀"), 5);
234    }
235
236    // ── split_string_by_byte_size ───────────────────────────────────
237
238    #[test]
239    fn split_ascii_exact_and_remainder() {
240        let chunks = split_string_by_byte_size("abcdefghij", 4).unwrap();
241        assert_eq!(chunks, vec!["abcd", "efgh", "ij"]);
242        assert_eq!(chunks.concat(), "abcdefghij");
243    }
244
245    #[test]
246    fn split_empty_yields_no_chunks() {
247        assert!(split_string_by_byte_size("", 4).unwrap().is_empty());
248    }
249
250    #[test]
251    fn split_rejects_zero_max() {
252        let err = split_string_by_byte_size("abc", 0).unwrap_err();
253        assert!(matches!(err, OversizedTransferError::Policy(_)));
254    }
255
256    #[test]
257    fn split_multibyte_never_breaks_codepoints() {
258        // Mixed 1/2/3/4-byte codepoints; chunk size 5 forces boundaries mid-string.
259        let original = "aé🦀b日x☃yz🦀";
260        let chunks = split_string_by_byte_size(original, 5).unwrap();
261
262        // Concatenation in order reproduces the exact bytes.
263        assert_eq!(chunks.concat(), original);
264        // No chunk exceeds the byte budget, and every chunk is valid UTF-8
265        // (guaranteed by `String`, i.e. no codepoint was split).
266        for chunk in &chunks {
267            assert!(utf8_byte_len(chunk) <= 5, "chunk {chunk:?} exceeds 5 bytes");
268            assert!(!chunk.is_empty());
269        }
270    }
271
272    #[test]
273    fn split_packs_multibyte_up_to_budget() {
274        // Each 🦀 is 4 bytes; with a budget of 8 two fit per chunk.
275        assert_eq!(
276            split_string_by_byte_size("🦀🦀🦀", 8).unwrap(),
277            vec!["🦀🦀", "🦀"]
278        );
279    }
280
281    #[test]
282    fn split_rejects_single_char_larger_than_budget() {
283        // 🦀 is 4 bytes; a 3-byte budget cannot hold it.
284        let err = split_string_by_byte_size("🦀", 3).unwrap_err();
285        assert!(matches!(err, OversizedTransferError::Policy(_)));
286    }
287
288    // ── build_oversized_frames ──────────────────────────────────────
289
290    #[test]
291    fn build_lays_out_progress_slots_without_handshake() {
292        let opts = OversizedSenderOptions::new("tok").with_chunk_size(4);
293        let frames = build_oversized_frames("hello world", &opts).unwrap();
294
295        // "hello world" (11 bytes) / 4 → 3 chunks.
296        assert_eq!(frames.chunks.len(), 3);
297        assert_eq!(frames.frame_count(), 5);
298
299        assert_eq!(progress(&frames.start), 1);
300        assert_eq!(frame_type(&frames.start), "start");
301        // No handshake: chunks reuse the reserved accept slot (start at 2).
302        assert_eq!(progress(&frames.chunks[0]), 2);
303        assert_eq!(progress(&frames.chunks[1]), 3);
304        assert_eq!(progress(&frames.chunks[2]), 4);
305        assert_eq!(frame_type(&frames.chunks[0]), "chunk");
306        assert_eq!(progress(&frames.end), 5);
307        assert_eq!(frame_type(&frames.end), "end");
308    }
309
310    #[test]
311    fn build_reserves_accept_slot_with_handshake() {
312        let opts = OversizedSenderOptions::new("tok")
313            .with_chunk_size(4)
314            .with_accept_handshake(true);
315        let frames = build_oversized_frames("hello world", &opts).unwrap();
316
317        assert_eq!(progress(&frames.start), 1);
318        // Handshake: slot 2 reserved for accept, chunks begin at 3.
319        assert_eq!(progress(&frames.chunks[0]), 3);
320        assert_eq!(progress(&frames.chunks[2]), 5);
321        assert_eq!(progress(&frames.end), 6);
322    }
323
324    #[test]
325    fn build_start_frame_declares_digest_bytes_chunks() {
326        let opts = OversizedSenderOptions::new("tok").with_chunk_size(4);
327        let frames = build_oversized_frames("hello world", &opts).unwrap();
328        let cvm = &frames.start.params.as_ref().unwrap()["cvm"];
329
330        assert_eq!(cvm["type"].as_str(), Some("oversized-transfer"));
331        assert_eq!(cvm["completionMode"], Value::String("render".to_string()));
332        assert_eq!(cvm["digest"], Value::String(sha256_digest("hello world")));
333        assert_eq!(cvm["totalBytes"].as_u64(), Some(11));
334        assert_eq!(cvm["totalChunks"].as_u64(), Some(3));
335    }
336
337    #[test]
338    fn build_propagates_single_char_over_budget_error() {
339        let opts = OversizedSenderOptions::new("tok").with_chunk_size(3);
340        let err = build_oversized_frames("🦀", &opts).unwrap_err();
341        // Surfaced through the crate error as an oversized-transfer policy error.
342        assert!(matches!(
343            err,
344            crate::core::error::Error::OversizedTransfer(OversizedTransferError::Policy(_))
345        ));
346    }
347
348    #[test]
349    fn build_into_ordered_is_start_chunks_end() {
350        let opts = OversizedSenderOptions::new("tok").with_chunk_size(4);
351        let frames = build_oversized_frames("hello world", &opts).unwrap();
352        let ordered = frames.into_ordered();
353        assert_eq!(ordered.len(), 5);
354        assert_eq!(frame_type(&ordered[0]), "start");
355        assert_eq!(frame_type(&ordered[4]), "end");
356        // Progress is strictly increasing across the whole sequence.
357        let progresses: Vec<u64> = ordered.iter().map(progress).collect();
358        assert_eq!(progresses, vec![1, 2, 3, 4, 5]);
359    }
360}