| 1 | //! Shared token-usage and SSE byte-line decoding extracted from `client.rs`. |
| 2 | //! |
| 3 | //! Protocol adapters retain event interpretation; this module owns the existing |
| 4 | //! accounting and fail-closed UTF-8 line helpers used across those adapters. |
| 5 | |
| 6 | use anyhow::Result; |
| 7 | use serde_json::Value; |
| 8 | |
| 9 | use codewhale_models::{ServerToolUsage, Usage}; |
| 10 | |
| 11 | pub(crate) fn saturating_u32(value: u64) -> u32 { |
| 12 | u32::try_from(value).unwrap_or(u32::MAX) |
| 13 | } |
| 14 | |
| 15 | pub(crate) fn parse_usage(usage: Option<&Value>) -> Usage { |
| 16 | let input_tokens = usage |
| 17 | .and_then(|u| u.get("input_tokens").or_else(|| u.get("prompt_tokens"))) |
| 18 | .and_then(Value::as_u64) |
| 19 | .unwrap_or(0); |
| 20 | let mut output_tokens = usage |
| 21 | .and_then(|u| { |
| 22 | u.get("output_tokens") |
| 23 | .or_else(|| u.get("completion_tokens")) |
| 24 | }) |
| 25 | .and_then(Value::as_u64) |
| 26 | .unwrap_or(0); |
| 27 | let total_tokens = usage |
| 28 | .and_then(|u| u.get("total_tokens")) |
| 29 | .and_then(Value::as_u64); |
| 30 | let reasoning_tokens_raw = usage |
| 31 | .and_then(|u| u.get("completion_tokens_details")) |
| 32 | .and_then(|details| details.get("reasoning_tokens")) |
| 33 | .and_then(Value::as_u64); |
| 34 | if output_tokens == 0 |
| 35 | && let Some(reasoning_tokens) = reasoning_tokens_raw |
| 36 | { |
| 37 | output_tokens = reasoning_tokens; |
| 38 | } else if output_tokens == 0 |
| 39 | && let Some(total_tokens) = total_tokens |
| 40 | { |
| 41 | output_tokens = total_tokens.saturating_sub(input_tokens); |
| 42 | } |
| 43 | let cached_tokens = usage |
| 44 | .and_then(|u| u.get("prompt_tokens_details")) |
| 45 | .and_then(|details| details.get("cached_tokens")) |
| 46 | .and_then(Value::as_u64); |
| 47 | let prompt_cache_hit_tokens = usage |
| 48 | .and_then(|u| u.get("prompt_cache_hit_tokens")) |
| 49 | .and_then(Value::as_u64) |
| 50 | .or(cached_tokens) |
| 51 | .map(saturating_u32); |
| 52 | let prompt_cache_miss_tokens = usage |
| 53 | .and_then(|u| u.get("prompt_cache_miss_tokens")) |
| 54 | .and_then(Value::as_u64) |
| 55 | .or_else(|| prompt_cache_hit_tokens.map(|hit| input_tokens.saturating_sub(u64::from(hit)))) |
| 56 | .map(saturating_u32); |
| 57 | // Reasoning tokens are a *subset* of the completion count every provider |
| 58 | // bills, so they are never added to output. A payload claiming more |
| 59 | // reasoning than output contradicts that invariant, which makes the figure |
| 60 | // invalid telemetry rather than extra billable output: drop it instead of |
| 61 | // letting a bad number reach the cost surfaces (#4318). |
| 62 | let reasoning_tokens = reasoning_tokens_raw |
| 63 | .filter(|reasoning| *reasoning <= output_tokens) |
| 64 | .map(saturating_u32); |
| 65 | |
| 66 | let server_tool_use = usage.and_then(|u| u.get("server_tool_use")).map(|server| { |
| 67 | let code_execution_requests = server |
| 68 | .get("code_execution_requests") |
| 69 | .and_then(Value::as_u64) |
| 70 | .map(saturating_u32); |
| 71 | let tool_search_requests = server |
| 72 | .get("tool_search_requests") |
| 73 | .and_then(Value::as_u64) |
| 74 | .map(saturating_u32); |
| 75 | ServerToolUsage { |
| 76 | code_execution_requests, |
| 77 | tool_search_requests, |
| 78 | } |
| 79 | }); |
| 80 | |
| 81 | Usage { |
| 82 | input_tokens: saturating_u32(input_tokens), |
| 83 | output_tokens: saturating_u32(output_tokens), |
| 84 | prompt_cache_hit_tokens, |
| 85 | prompt_cache_miss_tokens, |
| 86 | prompt_cache_write_tokens: None, |
| 87 | reasoning_tokens, |
| 88 | reasoning_replay_tokens: None, |
| 89 | server_tool_use, |
| 90 | } |
| 91 | } |
| 92 | |
| 93 | pub(super) fn extract_sse_data_value(line: &str) -> Option<&str> { |
| 94 | line.strip_prefix("data:") |
| 95 | .map(|value| value.strip_prefix(' ').unwrap_or(value)) |
| 96 | } |
| 97 | |
| 98 | /// Genuine invalid UTF-8 in an SSE line (or an unterminated flush). |
| 99 | /// |
| 100 | /// HTTP/2 DATA and other transports may split a multi-byte character across |
| 101 | /// chunks. That is not this error: callers must buffer raw bytes until a |
| 102 | /// complete line (or stream end) before decoding. This type is only returned |
| 103 | /// when `str::from_utf8` rejects the assembled bytes. We never substitute |
| 104 | /// U+FFFD — fail closed so garbled CJK cannot enter the transcript. |
| 105 | #[derive(Debug, Clone, Copy, PartialEq, Eq)] |
| 106 | pub(super) struct InvalidSseUtf8 { |
| 107 | valid_up_to: usize, |
| 108 | } |
| 109 | |
| 110 | impl std::fmt::Display for InvalidSseUtf8 { |
| 111 | fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { |
| 112 | write!( |
| 113 | f, |
| 114 | "invalid UTF-8 in SSE stream at byte {}", |
| 115 | self.valid_up_to |
| 116 | ) |
| 117 | } |
| 118 | } |
| 119 | |
| 120 | impl std::error::Error for InvalidSseUtf8 {} |
| 121 | |
| 122 | /// Decode one assembled SSE line (or stream-end tail) with `str::from_utf8`. |
| 123 | /// Does not substitute U+FFFD. |
| 124 | fn decode_sse_line_bytes(bytes: &[u8]) -> Result<&str, InvalidSseUtf8> { |
| 125 | std::str::from_utf8(bytes).map_err(|err| InvalidSseUtf8 { |
| 126 | valid_up_to: err.valid_up_to(), |
| 127 | }) |
| 128 | } |
| 129 | |
| 130 | /// Take the next COMPLETE line (up to the first `\n`) off a raw byte buffer, |
| 131 | /// draining it, and return it trimmed. Returns `Ok(None)` when no full line is |
| 132 | /// buffered yet. Decoding only complete lines (never an arbitrary network-read |
| 133 | /// boundary) means a multi-byte UTF-8 char — CJK, emoji, accented letter — |
| 134 | /// split across two reads is never corrupted to U+FFFD, since the `\n` |
| 135 | /// delimiter is ASCII and can never fall inside a multi-byte sequence. |
| 136 | /// |
| 137 | /// Genuine invalid bytes fail closed (`Err(InvalidSseUtf8)`); we do not |
| 138 | /// substitute U+FFFD. |
| 139 | pub(super) fn take_sse_line(buffer: &mut Vec<u8>) -> Result<Option<String>, InvalidSseUtf8> { |
| 140 | let Some(line_end) = buffer.iter().position(|&b| b == b'\n') else { |
| 141 | return Ok(None); |
| 142 | }; |
| 143 | // Strip a preceding `\r` so CRLF-delimited SSE frames do not leave CR. |
| 144 | let mut end = line_end; |
| 145 | if end > 0 && buffer[end - 1] == b'\r' { |
| 146 | end -= 1; |
| 147 | } |
| 148 | let decoded = decode_sse_line_bytes(&buffer[..end]).map(|text| text.trim().to_string()); |
| 149 | buffer.drain(..=line_end); |
| 150 | decoded.map(Some) |
| 151 | } |
| 152 | |
| 153 | /// Decode the unterminated tail left in `buffer` at stream end. |
| 154 | /// |
| 155 | /// Same fail-closed UTF-8 contract as [`take_sse_line`]. Empty / whitespace-only |
| 156 | /// tails yield `Ok(None)`. |
| 157 | pub(super) fn flush_sse_line(buffer: &mut Vec<u8>) -> Result<Option<String>, InvalidSseUtf8> { |
| 158 | if buffer.is_empty() { |
| 159 | return Ok(None); |
| 160 | } |
| 161 | let mut end = buffer.len(); |
| 162 | if buffer[end - 1] == b'\r' { |
| 163 | end -= 1; |
| 164 | } |
| 165 | let decoded = decode_sse_line_bytes(&buffer[..end]).map(|text| text.trim().to_string()); |
| 166 | buffer.clear(); |
| 167 | decoded.map(|line| (!line.is_empty()).then_some(line)) |
| 168 | } |
| 169 | |
| 170 | /// Next decoded SSE line. When `at_end` is false, wait for `\n`. When `at_end` |
| 171 | /// is true, also flush an unterminated tail (stream closed). |
| 172 | pub(super) fn next_sse_line( |
| 173 | buffer: &mut Vec<u8>, |
| 174 | at_end: bool, |
| 175 | ) -> Result<Option<String>, InvalidSseUtf8> { |
| 176 | match take_sse_line(buffer)? { |
| 177 | Some(line) => Ok(Some(line)), |
| 178 | None if at_end => flush_sse_line(buffer), |
| 179 | None => Ok(None), |
| 180 | } |
| 181 | } |
| 182 | |
| 183 | /// Incremental raw-byte SSE line assembler for tests and the Chat Completions |
| 184 | /// decoder. HTTP/2 DATA may split a multi-byte UTF-8 character across chunks; |
| 185 | /// we never decode until a complete line or [`SseLineDecoder::finish`]. |
| 186 | #[cfg(test)] |
| 187 | pub(super) struct SseLineDecoder { |
| 188 | buffer: Vec<u8>, |
| 189 | } |
| 190 | |
| 191 | #[cfg(test)] |
| 192 | impl SseLineDecoder { |
| 193 | pub(super) fn new() -> Self { |
| 194 | Self { buffer: Vec::new() } |
| 195 | } |
| 196 | |
| 197 | pub(super) fn push(&mut self, chunk: &[u8]) -> Result<Vec<String>, InvalidSseUtf8> { |
| 198 | self.buffer.extend_from_slice(chunk); |
| 199 | let mut lines = Vec::new(); |
| 200 | while let Some(line) = take_sse_line(&mut self.buffer)? { |
| 201 | lines.push(line); |
| 202 | } |
| 203 | Ok(lines) |
| 204 | } |
| 205 | |
| 206 | pub(super) fn finish(mut self) -> Result<Option<String>, InvalidSseUtf8> { |
| 207 | flush_sse_line(&mut self.buffer) |
| 208 | } |
| 209 | } |
| 210 | |
| 211 | #[cfg(test)] |
| 212 | mod tests { |
| 213 | use super::*; |
| 214 | use serde_json::json; |
| 215 | |
| 216 | #[test] |
| 217 | fn parse_usage_scenario() { |
| 218 | // Scenario consolidation of: parse_usage_reads_deepseek_cache_and_reasoning_tokens, parse_usage_saturates_every_u64_token_field, parse_usage_counts_reasoning_tokens_when_completion_tokens_are_zero, parse_usage_derives_completion_tokens_from_total_tokens_when_needed, parse_usage_reads_v4_prompt_tokens_details_cached_tokens, parse_usage_infers_cache_miss_from_selected_hit_source |
| 219 | // from parse_usage_reads_deepseek_cache_and_reasoning_tokens |
| 220 | { |
| 221 | let usage = parse_usage(Some(&json!({ |
| 222 | "prompt_tokens": 100, |
| 223 | "completion_tokens": 20, |
| 224 | "prompt_cache_hit_tokens": 70, |
| 225 | "prompt_cache_miss_tokens": 30, |
| 226 | "completion_tokens_details": { |
| 227 | "reasoning_tokens": 12 |
| 228 | } |
| 229 | }))); |
| 230 | |
| 231 | assert_eq!(usage.input_tokens, 100); |
| 232 | assert_eq!(usage.output_tokens, 20); |
| 233 | assert_eq!(usage.prompt_cache_hit_tokens, Some(70)); |
| 234 | assert_eq!(usage.prompt_cache_miss_tokens, Some(30)); |
| 235 | assert_eq!(usage.reasoning_tokens, Some(12)); |
| 236 | } |
| 237 | // from parse_usage_saturates_every_u64_token_field |
| 238 | { |
| 239 | let usage = parse_usage(Some(&json!({ |
| 240 | "input_tokens": u64::MAX, |
| 241 | "output_tokens": u64::MAX, |
| 242 | "prompt_cache_hit_tokens": u64::MAX, |
| 243 | "prompt_cache_miss_tokens": u64::MAX, |
| 244 | "completion_tokens_details": { "reasoning_tokens": u64::MAX }, |
| 245 | "server_tool_use": { |
| 246 | "code_execution_requests": u64::MAX, |
| 247 | "tool_search_requests": u64::MAX |
| 248 | } |
| 249 | }))); |
| 250 | assert_eq!(usage.input_tokens, u32::MAX); |
| 251 | assert_eq!(usage.output_tokens, u32::MAX); |
| 252 | assert_eq!(usage.prompt_cache_hit_tokens, Some(u32::MAX)); |
| 253 | assert_eq!(usage.prompt_cache_miss_tokens, Some(u32::MAX)); |
| 254 | assert_eq!(usage.reasoning_tokens, Some(u32::MAX)); |
| 255 | let server = usage.server_tool_use.expect("server usage"); |
| 256 | assert_eq!(server.code_execution_requests, Some(u32::MAX)); |
| 257 | assert_eq!(server.tool_search_requests, Some(u32::MAX)); |
| 258 | } |
| 259 | // from parse_usage_counts_reasoning_tokens_when_completion_tokens_are_zero |
| 260 | { |
| 261 | let usage = parse_usage(Some(&json!({ |
| 262 | "prompt_tokens": 100, |
| 263 | "completion_tokens": 0, |
| 264 | "completion_tokens_details": { |
| 265 | "reasoning_tokens": 12 |
| 266 | } |
| 267 | }))); |
| 268 | |
| 269 | assert_eq!(usage.input_tokens, 100); |
| 270 | assert_eq!(usage.output_tokens, 12); |
| 271 | assert_eq!(usage.reasoning_tokens, Some(12)); |
| 272 | assert!( |
| 273 | crate::pricing::calculate_turn_cost_from_usage("deepseek-v4-pro", &usage) |
| 274 | .expect("DeepSeek V4 Pro pricing should apply") |
| 275 | > 0.0 |
| 276 | ); |
| 277 | } |
| 278 | // from parse_usage_derives_completion_tokens_from_total_tokens_when_needed |
| 279 | { |
| 280 | let usage = parse_usage(Some(&json!({ |
| 281 | "prompt_tokens": 100, |
| 282 | "total_tokens": 125, |
| 283 | "prompt_cache_hit_tokens": 70, |
| 284 | "prompt_cache_miss_tokens": 30 |
| 285 | }))); |
| 286 | |
| 287 | assert_eq!(usage.input_tokens, 100); |
| 288 | assert_eq!(usage.output_tokens, 25); |
| 289 | assert_eq!(usage.prompt_cache_hit_tokens, Some(70)); |
| 290 | assert_eq!(usage.prompt_cache_miss_tokens, Some(30)); |
| 291 | } |
| 292 | // from parse_usage_reads_v4_prompt_tokens_details_cached_tokens |
| 293 | { |
| 294 | let usage = parse_usage(Some(&json!({ |
| 295 | "prompt_tokens": 4000, |
| 296 | "completion_tokens": 20, |
| 297 | "prompt_tokens_details": { |
| 298 | "cached_tokens": 3000 |
| 299 | } |
| 300 | }))); |
| 301 | |
| 302 | assert_eq!(usage.input_tokens, 4000); |
| 303 | assert_eq!(usage.output_tokens, 20); |
| 304 | assert_eq!(usage.prompt_cache_hit_tokens, Some(3000)); |
| 305 | assert_eq!(usage.prompt_cache_miss_tokens, Some(1000)); |
| 306 | } |
| 307 | // from parse_usage_infers_cache_miss_from_selected_hit_source |
| 308 | { |
| 309 | let usage = parse_usage(Some(&json!({ |
| 310 | "prompt_tokens": 4000, |
| 311 | "completion_tokens": 20, |
| 312 | "prompt_cache_hit_tokens": 3000, |
| 313 | "prompt_tokens_details": { |
| 314 | "cached_tokens": 1000 |
| 315 | } |
| 316 | }))); |
| 317 | |
| 318 | assert_eq!(usage.input_tokens, 4000); |
| 319 | assert_eq!(usage.prompt_cache_hit_tokens, Some(3000)); |
| 320 | assert_eq!(usage.prompt_cache_miss_tokens, Some(1000)); |
| 321 | } |
| 322 | } |
| 323 | |
| 324 | /// Real-shaped Chat-Completions usage payloads from the three providers most |
| 325 | /// likely to report reasoning tokens, carried end-to-end into pricing. |
| 326 | /// |
| 327 | /// Two invariants hold for every fixture: `reasoning_tokens <= output_tokens`, |
| 328 | /// and pricing never adds reasoning on top of output — dropping the reasoning |
| 329 | /// field entirely must not change the cost by a single cent. |
| 330 | #[test] |
| 331 | fn reasoning_parser_fixtures_never_exceed_or_add_to_billable_output() { |
| 332 | use crate::config::ApiProvider; |
| 333 | use crate::pricing::{calculate_turn_cost_estimate_for_provider, token_usage_for_pricing}; |
| 334 | |
| 335 | // (label, provider, model, payload) |
| 336 | let fixtures: [(&str, ApiProvider, &str, serde_json::Value); 3] = [ |
| 337 | ( |
| 338 | "moonshot", |
| 339 | ApiProvider::Moonshot, |
| 340 | "kimi-k2.7-code", |
| 341 | json!({ |
| 342 | "prompt_tokens": 30_000, |
| 343 | "completion_tokens": 2_400, |
| 344 | "total_tokens": 32_400, |
| 345 | "prompt_tokens_details": { "cached_tokens": 24_000 }, |
| 346 | "completion_tokens_details": { "reasoning_tokens": 1_900 } |
| 347 | }), |
| 348 | ), |
| 349 | ( |
| 350 | "minimax", |
| 351 | ApiProvider::Minimax, |
| 352 | "minimax-m3", |
| 353 | json!({ |
| 354 | "prompt_tokens": 12_000, |
| 355 | "completion_tokens": 3_000, |
| 356 | "total_tokens": 15_000, |
| 357 | "prompt_tokens_details": { "cached_tokens": 4_000 }, |
| 358 | "completion_tokens_details": { "reasoning_tokens": 2_950 } |
| 359 | }), |
| 360 | ), |
| 361 | ( |
| 362 | "openrouter", |
| 363 | ApiProvider::Openrouter, |
| 364 | "qwen/qwen3.7-plus", |
| 365 | json!({ |
| 366 | "prompt_tokens": 8_000, |
| 367 | "completion_tokens": 1_500, |
| 368 | "total_tokens": 9_500, |
| 369 | "prompt_tokens_details": { "cached_tokens": 2_000 }, |
| 370 | "completion_tokens_details": { "reasoning_tokens": 1_500 } |
| 371 | }), |
| 372 | ), |
| 373 | ]; |
| 374 | |
| 375 | for (label, provider, model, payload) in fixtures { |
| 376 | let usage = parse_usage(Some(&payload)); |
| 377 | let reasoning = usage.reasoning_tokens.expect("fixture reports reasoning"); |
| 378 | |
| 379 | // Invariant 1: reasoning is a subset of the billed completion count. |
| 380 | assert!( |
| 381 | reasoning <= usage.output_tokens, |
| 382 | "{label}: reasoning {reasoning} exceeds output {}", |
| 383 | usage.output_tokens |
| 384 | ); |
| 385 | // Billable output is exactly the reported completion count. |
| 386 | let classes = token_usage_for_pricing(&usage); |
| 387 | assert_eq!( |
| 388 | classes.output, |
| 389 | u64::from(usage.output_tokens), |
| 390 | "{label}: reasoning leaked into billable output" |
| 391 | ); |
| 392 | |
| 393 | // Invariant 2: pricing does not add reasoning a second time. The same |
| 394 | // usage with the reasoning field removed must cost the same. |
| 395 | let without = codewhale_models::Usage { |
| 396 | reasoning_tokens: None, |
| 397 | ..usage.clone() |
| 398 | }; |
| 399 | assert_eq!( |
| 400 | calculate_turn_cost_estimate_for_provider(provider, model, &usage), |
| 401 | calculate_turn_cost_estimate_for_provider(provider, model, &without), |
| 402 | "{label}: reasoning changed the price" |
| 403 | ); |
| 404 | } |
| 405 | } |
| 406 | |
| 407 | /// A payload claiming more reasoning than output contradicts the subset |
| 408 | /// invariant. That is broken telemetry, so the field is discarded — and it |
| 409 | /// must never become extra billable output. |
| 410 | #[test] |
| 411 | fn pathological_reasoning_above_output_is_rejected_not_billed() { |
| 412 | let usage = parse_usage(Some(&json!({ |
| 413 | "prompt_tokens": 1_000, |
| 414 | "completion_tokens": 100, |
| 415 | "completion_tokens_details": { "reasoning_tokens": 5_000 } |
| 416 | }))); |
| 417 | |
| 418 | assert_eq!(usage.output_tokens, 100, "output stays as reported"); |
| 419 | assert_eq!( |
| 420 | usage.reasoning_tokens, None, |
| 421 | "impossible reasoning telemetry is dropped rather than trusted" |
| 422 | ); |
| 423 | let classes = crate::pricing::token_usage_for_pricing(&usage); |
| 424 | assert_eq!(classes.output, 100); |
| 425 | |
| 426 | // `completion_tokens: 0` with reasoning present is the *legitimate* |
| 427 | // shape this filter must not break: providers that report only reasoning |
| 428 | // set output from it, keeping reasoning == output. |
| 429 | let zero_output = parse_usage(Some(&json!({ |
| 430 | "prompt_tokens": 1_000, |
| 431 | "completion_tokens": 0, |
| 432 | "completion_tokens_details": { "reasoning_tokens": 12 } |
| 433 | }))); |
| 434 | assert_eq!(zero_output.output_tokens, 12); |
| 435 | assert_eq!(zero_output.reasoning_tokens, Some(12)); |
| 436 | } |
| 437 | |
| 438 | fn mid_char_split(text: &str, ch: char) -> usize { |
| 439 | let needle = ch.to_string(); |
| 440 | let start = text |
| 441 | .as_bytes() |
| 442 | .windows(needle.len()) |
| 443 | .position(|window| window == needle.as_bytes()) |
| 444 | .unwrap_or_else(|| panic!("{ch:?} present in {text:?}")); |
| 445 | start + 1 |
| 446 | } |
| 447 | |
| 448 | #[test] |
| 449 | fn take_sse_scenario() { |
| 450 | // Scenario consolidation of: take_sse_line_preserves_multibyte_split_across_reads, take_sse_line_returns_none_without_newline, take_sse_line_reassembles_cjk_and_rejects_invalid_bytes, take_sse_line_rejects_invalid_bytes_without_replacement |
| 451 | // from take_sse_line_preserves_multibyte_split_across_reads |
| 452 | { |
| 453 | // "你好" streamed so the 3-byte '好' straddles a read boundary. |
| 454 | let full = "data: 你好\n"; |
| 455 | let bytes = full.as_bytes(); |
| 456 | let split = mid_char_split(full, '好'); |
| 457 | let mut buffer: Vec<u8> = Vec::new(); |
| 458 | // First read: no complete line yet. |
| 459 | buffer.extend_from_slice(&bytes[..split]); |
| 460 | assert_eq!(take_sse_line(&mut buffer).expect("valid prefix"), None); |
| 461 | // Second read completes the line; '好' must be intact, not U+FFFD. |
| 462 | buffer.extend_from_slice(&bytes[split..]); |
| 463 | let line = take_sse_line(&mut buffer) |
| 464 | .expect("valid utf-8") |
| 465 | .expect("a complete line"); |
| 466 | assert_eq!(line, "data: 你好"); |
| 467 | assert!(!line.contains('\u{FFFD}'), "multibyte char was corrupted"); |
| 468 | assert_eq!(extract_sse_data_value(&line), Some("你好")); |
| 469 | // Buffer fully drained. |
| 470 | assert!(buffer.is_empty()); |
| 471 | } |
| 472 | // from take_sse_line_returns_none_without_newline |
| 473 | { |
| 474 | let mut buffer = b"data: partial".to_vec(); |
| 475 | assert_eq!(take_sse_line(&mut buffer).expect("valid utf-8"), None); |
| 476 | assert_eq!(buffer, b"data: partial"); |
| 477 | } |
| 478 | // from take_sse_line_reassembles_cjk_and_rejects_invalid_bytes |
| 479 | { |
| 480 | let full = "data: 测试中文\n"; |
| 481 | let split = mid_char_split(full, '试'); |
| 482 | let mut buffer = full.as_bytes()[..split].to_vec(); |
| 483 | assert_eq!(take_sse_line(&mut buffer).expect("valid prefix"), None); |
| 484 | buffer.extend_from_slice(&full.as_bytes()[split..]); |
| 485 | let line = take_sse_line(&mut buffer) |
| 486 | .expect("valid utf-8") |
| 487 | .expect("complete line"); |
| 488 | assert_eq!(line, "data: 测试中文"); |
| 489 | assert!(!line.contains('\u{FFFD}')); |
| 490 | |
| 491 | let mut invalid = b"data: ok".to_vec(); |
| 492 | invalid.push(0xFF); |
| 493 | invalid.push(b'\n'); |
| 494 | let err = take_sse_line(&mut invalid).expect_err("invalid bytes must fail closed"); |
| 495 | assert!(!err.to_string().contains('\u{FFFD}')); |
| 496 | assert_eq!(err.valid_up_to, 8); |
| 497 | assert!( |
| 498 | invalid.is_empty(), |
| 499 | "invalid line is consumed so retries cannot loop" |
| 500 | ); |
| 501 | } |
| 502 | // from take_sse_line_rejects_invalid_bytes_without_replacement |
| 503 | { |
| 504 | let mut buffer = b"data: ok".to_vec(); |
| 505 | buffer.push(0xFF); |
| 506 | buffer.extend_from_slice(b"\n"); |
| 507 | let err = take_sse_line(&mut buffer).expect_err("0xFF is not UTF-8"); |
| 508 | assert_eq!(err.valid_up_to, 8); |
| 509 | assert!(!err.to_string().contains('\u{FFFD}')); |
| 510 | assert!(buffer.is_empty(), "invalid line must be drained"); |
| 511 | } |
| 512 | } |
| 513 | |
| 514 | #[test] |
| 515 | fn flush_sse_scenario() { |
| 516 | // Scenario consolidation of: flush_sse_line_reassembles_cjk_and_rejects_invalid_bytes, flush_sse_line_preserves_unterminated_cjk, flush_sse_line_rejects_truncated_multibyte_sequence |
| 517 | // from flush_sse_line_reassembles_cjk_and_rejects_invalid_bytes |
| 518 | { |
| 519 | let text = "data: 你好世界"; |
| 520 | let split = mid_char_split(text, '好'); |
| 521 | let mut buffer = text.as_bytes()[..split].to_vec(); |
| 522 | assert_eq!(take_sse_line(&mut buffer).expect("no newline yet"), None); |
| 523 | buffer.extend_from_slice(&text.as_bytes()[split..]); |
| 524 | let line = flush_sse_line(&mut buffer) |
| 525 | .expect("valid utf-8") |
| 526 | .expect("unterminated tail"); |
| 527 | assert_eq!(line, "data: 你好世界"); |
| 528 | assert!(!line.contains('\u{FFFD}')); |
| 529 | assert!(buffer.is_empty()); |
| 530 | assert_eq!(flush_sse_line(&mut buffer).expect("empty"), None); |
| 531 | |
| 532 | let mut invalid = vec![0x80, 0xBF]; |
| 533 | let err = flush_sse_line(&mut invalid).expect_err("invalid flush must fail closed"); |
| 534 | assert!(!err.to_string().contains('\u{FFFD}')); |
| 535 | assert_eq!(err.valid_up_to, 0); |
| 536 | assert!(invalid.is_empty()); |
| 537 | } |
| 538 | // from flush_sse_line_preserves_unterminated_cjk |
| 539 | { |
| 540 | let mut buffer = "data: 你好".as_bytes().to_vec(); |
| 541 | let line = flush_sse_line(&mut buffer) |
| 542 | .expect("valid utf-8") |
| 543 | .expect("residual line"); |
| 544 | assert_eq!(line, "data: 你好"); |
| 545 | assert!(!line.contains('\u{FFFD}')); |
| 546 | assert!(buffer.is_empty()); |
| 547 | } |
| 548 | // from flush_sse_line_rejects_truncated_multibyte_sequence |
| 549 | { |
| 550 | let mut buffer = "data: ".as_bytes().to_vec(); |
| 551 | buffer.extend_from_slice(&"好".as_bytes()[..2]); |
| 552 | let err = flush_sse_line(&mut buffer).expect_err("truncated UTF-8"); |
| 553 | assert_eq!(err.valid_up_to, 6); |
| 554 | assert!(!err.to_string().contains('\u{FFFD}')); |
| 555 | assert!(buffer.is_empty()); |
| 556 | } |
| 557 | } |
| 558 | |
| 559 | #[test] |
| 560 | fn decode_sse_line_bytes_rejects_invalid_without_replacement() { |
| 561 | let ok = decode_sse_line_bytes("data: 你好".as_bytes()).expect("valid"); |
| 562 | assert_eq!(ok, "data: 你好"); |
| 563 | assert!(!ok.contains('\u{FFFD}')); |
| 564 | |
| 565 | let err = decode_sse_line_bytes(&[0xFF]).expect_err("bare 0xFF is invalid"); |
| 566 | assert!(!err.to_string().contains('\u{FFFD}')); |
| 567 | assert_eq!(err.valid_up_to, 0); |
| 568 | } |
| 569 | |
| 570 | #[test] |
| 571 | fn extract_sse_scenario() { |
| 572 | // Scenario consolidation of: extract_sse_data_value_accepts_optional_space, extract_sse_data_value_handles_done_marker, extract_sse_data_value_rejects_non_data_lines |
| 573 | // from extract_sse_data_value_accepts_optional_space |
| 574 | { |
| 575 | assert_eq!( |
| 576 | extract_sse_data_value("data: {\"ok\":true}"), |
| 577 | Some("{\"ok\":true}") |
| 578 | ); |
| 579 | assert_eq!( |
| 580 | extract_sse_data_value("data:{\"ok\":true}"), |
| 581 | Some("{\"ok\":true}") |
| 582 | ); |
| 583 | } |
| 584 | // from extract_sse_data_value_handles_done_marker |
| 585 | { |
| 586 | assert_eq!(extract_sse_data_value("data: [DONE]"), Some("[DONE]")); |
| 587 | assert_eq!(extract_sse_data_value("data:[DONE]"), Some("[DONE]")); |
| 588 | } |
| 589 | // from extract_sse_data_value_rejects_non_data_lines |
| 590 | { |
| 591 | assert_eq!(extract_sse_data_value("event: message"), None); |
| 592 | assert_eq!(extract_sse_data_value(": heartbeat"), None); |
| 593 | } |
| 594 | } |
| 595 | } |
| 596 |