contextvm_sdk/transport/oversized_transfer/
codec.rs1use 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#[derive(Debug, Clone)]
17pub struct OversizedSenderOptions {
18 pub progress_token: String,
20 pub chunk_size_bytes: usize,
22 pub needs_accept_handshake: bool,
25}
26
27impl OversizedSenderOptions {
28 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 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 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#[derive(Debug, Clone)]
52pub struct BuiltOversizedFrames {
53 pub start: JsonRpcNotification,
55 pub chunks: Vec<JsonRpcNotification>,
57 pub end: JsonRpcNotification,
59}
60
61impl BuiltOversizedFrames {
62 pub fn frame_count(&self) -> usize {
64 self.chunks.len() + 2
65 }
66
67 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
77pub fn utf8_byte_len(value: &str) -> usize {
79 value.len()
80}
81
82pub 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
90pub 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
133pub 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 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 #[test]
205 fn sha256_digest_known_vectors() {
206 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 #[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 let original = "aé🦀b日x☃yz🦀";
260 let chunks = split_string_by_byte_size(original, 5).unwrap();
261
262 assert_eq!(chunks.concat(), original);
264 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 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 let err = split_string_by_byte_size("🦀", 3).unwrap_err();
285 assert!(matches!(err, OversizedTransferError::Policy(_)));
286 }
287
288 #[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 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 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 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 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 let progresses: Vec<u64> = ordered.iter().map(progress).collect();
358 assert_eq!(progresses, vec![1, 2, 3, 4, 5]);
359 }
360}