| 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 | /// Hard ceiling for one pending SSE line, matching the MCP frame limit. |
| 99 | /// This does not bound the size of a multi-line event or the whole stream. |
| 100 | const MAX_SSE_LINE_BYTES: usize = 8 * 1024 * 1024; |
| 101 | /// Hard ceiling for one assembled SSE event (its `data:` fields joined), so a |
| 102 | /// peer that never sends the blank line ending an event cannot grow memory |
| 103 | /// without bound. Same size as the single-line ceiling. |
| 104 | const MAX_SSE_EVENT_BYTES: usize = MAX_SSE_LINE_BYTES; |
| 105 | |
| 106 | /// Invalid or oversized SSE line (or unterminated flush). |
| 107 | /// |
| 108 | /// HTTP/2 DATA and other transports may split a multi-byte character across |
| 109 | /// chunks. That is not this error: callers must buffer raw bytes until a |
| 110 | /// complete line (or stream end) before decoding. We never substitute |
| 111 | /// U+FFFD — fail closed so garbled CJK cannot enter the transcript. |
| 112 | #[derive(Debug, Clone, Copy, PartialEq, Eq)] |
| 113 | pub(super) enum SseLineError { |
| 114 | InvalidUtf8 { valid_up_to: usize }, |
| 115 | TooLong, |
| 116 | EventTooLong, |
| 117 | } |
| 118 | |
| 119 | impl std::fmt::Display for SseLineError { |
| 120 | fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { |
| 121 | match self { |
| 122 | Self::InvalidUtf8 { valid_up_to } => { |
| 123 | write!(f, "invalid UTF-8 in SSE stream at byte {valid_up_to}") |
| 124 | } |
| 125 | Self::TooLong => write!( |
| 126 | f, |
| 127 | "SSE line exceeded {MAX_SSE_LINE_BYTES} bytes (8 MiB) — aborting stream" |
| 128 | ), |
| 129 | Self::EventTooLong => write!( |
| 130 | f, |
| 131 | "SSE event exceeded {MAX_SSE_EVENT_BYTES} bytes (8 MiB) — aborting stream" |
| 132 | ), |
| 133 | } |
| 134 | } |
| 135 | } |
| 136 | |
| 137 | impl std::error::Error for SseLineError {} |
| 138 | |
| 139 | /// Decode one assembled SSE line (or stream-end tail) with `str::from_utf8`. |
| 140 | /// Does not substitute U+FFFD. |
| 141 | fn decode_sse_line_bytes(bytes: &[u8]) -> Result<&str, SseLineError> { |
| 142 | std::str::from_utf8(bytes).map_err(|err| SseLineError::InvalidUtf8 { |
| 143 | valid_up_to: err.valid_up_to(), |
| 144 | }) |
| 145 | } |
| 146 | |
| 147 | /// Take the next COMPLETE line (up to the first `\n`) off a raw byte buffer, |
| 148 | /// draining it, and return it trimmed. Returns `Ok(None)` when no full line is |
| 149 | /// buffered yet. Decoding only complete lines (never an arbitrary network-read |
| 150 | /// boundary) means a multi-byte UTF-8 char — CJK, emoji, accented letter — |
| 151 | /// split across two reads is never corrupted to U+FFFD, since the `\n` |
| 152 | /// delimiter is ASCII and can never fall inside a multi-byte sequence. |
| 153 | /// |
| 154 | /// Genuine invalid bytes fail closed (`Err(SseLineError)`); we do not |
| 155 | /// substitute U+FFFD. |
| 156 | pub(super) fn take_sse_line(buffer: &mut Vec<u8>) -> Result<Option<String>, SseLineError> { |
| 157 | let line_end = buffer.iter().position(|&b| b == b'\n'); |
| 158 | if line_end.unwrap_or(buffer.len()) > MAX_SSE_LINE_BYTES { |
| 159 | buffer.clear(); |
| 160 | return Err(SseLineError::TooLong); |
| 161 | } |
| 162 | let Some(line_end) = line_end else { |
| 163 | return Ok(None); |
| 164 | }; |
| 165 | // Strip a preceding `\r` so CRLF-delimited SSE frames do not leave CR. |
| 166 | let mut end = line_end; |
| 167 | if end > 0 && buffer[end - 1] == b'\r' { |
| 168 | end -= 1; |
| 169 | } |
| 170 | let decoded = decode_sse_line_bytes(&buffer[..end]).map(|text| text.trim().to_string()); |
| 171 | buffer.drain(..=line_end); |
| 172 | decoded.map(Some) |
| 173 | } |
| 174 | |
| 175 | /// Decode the unterminated tail left in `buffer` at stream end. |
| 176 | /// |
| 177 | /// Same fail-closed UTF-8 contract as [`take_sse_line`]. Empty / whitespace-only |
| 178 | /// tails yield `Ok(None)`. |
| 179 | pub(super) fn flush_sse_line(buffer: &mut Vec<u8>) -> Result<Option<String>, SseLineError> { |
| 180 | if buffer.len() > MAX_SSE_LINE_BYTES { |
| 181 | buffer.clear(); |
| 182 | return Err(SseLineError::TooLong); |
| 183 | } |
| 184 | if buffer.is_empty() { |
| 185 | return Ok(None); |
| 186 | } |
| 187 | let mut end = buffer.len(); |
| 188 | if buffer[end - 1] == b'\r' { |
| 189 | end -= 1; |
| 190 | } |
| 191 | let decoded = decode_sse_line_bytes(&buffer[..end]).map(|text| text.trim().to_string()); |
| 192 | buffer.clear(); |
| 193 | decoded.map(|line| (!line.is_empty()).then_some(line)) |
| 194 | } |
| 195 | |
| 196 | /// Append one `data:` field to the event being assembled. The SSE spec joins |
| 197 | /// the `data:` fields of one event with '\n' and dispatches the event at the |
| 198 | /// blank line that ends it; a provider may split one JSON payload across |
| 199 | /// several fields, so parsing each field alone would drop the event. Fails |
| 200 | /// closed (clearing the event) past [`MAX_SSE_EVENT_BYTES`]. |
| 201 | pub(super) fn push_sse_event_data(event: &mut String, data: &str) -> Result<(), SseLineError> { |
| 202 | let needed = event |
| 203 | .len() |
| 204 | .saturating_add(usize::from(!event.is_empty())) |
| 205 | .saturating_add(data.len()); |
| 206 | if needed > MAX_SSE_EVENT_BYTES { |
| 207 | event.clear(); |
| 208 | return Err(SseLineError::EventTooLong); |
| 209 | } |
| 210 | if !event.is_empty() { |
| 211 | event.push('\n'); |
| 212 | } |
| 213 | event.push_str(data); |
| 214 | Ok(()) |
| 215 | } |
| 216 | |
| 217 | /// Next decoded SSE line. When `at_end` is false, wait for `\n`. When `at_end` |
| 218 | /// is true, also flush an unterminated tail (stream closed). |
| 219 | pub(super) fn next_sse_line( |
| 220 | buffer: &mut Vec<u8>, |
| 221 | at_end: bool, |
| 222 | ) -> Result<Option<String>, SseLineError> { |
| 223 | match take_sse_line(buffer)? { |
| 224 | Some(line) => Ok(Some(line)), |
| 225 | None if at_end => flush_sse_line(buffer), |
| 226 | None => Ok(None), |
| 227 | } |
| 228 | } |
| 229 | |
| 230 | /// Incremental raw-byte SSE line assembler for tests and the Chat Completions |
| 231 | /// decoder. HTTP/2 DATA may split a multi-byte UTF-8 character across chunks; |
| 232 | /// we never decode until a complete line or [`SseLineDecoder::finish`]. |
| 233 | #[cfg(test)] |
| 234 | pub(super) struct SseLineDecoder { |
| 235 | buffer: Vec<u8>, |
| 236 | } |
| 237 | |
| 238 | #[cfg(test)] |
| 239 | impl SseLineDecoder { |
| 240 | pub(super) fn new() -> Self { |
| 241 | Self { buffer: Vec::new() } |
| 242 | } |
| 243 | |
| 244 | pub(super) fn push(&mut self, chunk: &[u8]) -> Result<Vec<String>, SseLineError> { |
| 245 | self.buffer.extend_from_slice(chunk); |
| 246 | let mut lines = Vec::new(); |
| 247 | while let Some(line) = take_sse_line(&mut self.buffer)? { |
| 248 | lines.push(line); |
| 249 | } |
| 250 | Ok(lines) |
| 251 | } |
| 252 | |
| 253 | pub(super) fn finish(mut self) -> Result<Option<String>, SseLineError> { |
| 254 | flush_sse_line(&mut self.buffer) |
| 255 | } |
| 256 | } |
| 257 | |
| 258 | #[cfg(test)] |
| 259 | mod tests { |
| 260 | use super::*; |
| 261 | use serde_json::json; |
| 262 | |
| 263 | #[test] |
| 264 | fn sse_line_limit_rejects_oversized_pending_complete_and_final_lines() { |
| 265 | let mut decoder = SseLineDecoder::new(); |
| 266 | let at_limit = vec![b'x'; MAX_SSE_LINE_BYTES]; |
| 267 | assert!(decoder.push(&at_limit).unwrap().is_empty()); |
| 268 | let error = decoder |
| 269 | .push(b"x") |
| 270 | .expect_err("pending line exceeds the limit"); |
| 271 | assert_eq!(error, SseLineError::TooLong); |
| 272 | assert!(error.to_string().contains("8 MiB")); |
| 273 | assert!(decoder.buffer.is_empty()); |
| 274 | |
| 275 | let mut complete = at_limit.clone(); |
| 276 | complete.push(b'\n'); |
| 277 | assert_eq!( |
| 278 | take_sse_line(&mut complete).unwrap().unwrap().len(), |
| 279 | MAX_SSE_LINE_BYTES |
| 280 | ); |
| 281 | assert!(complete.is_empty()); |
| 282 | |
| 283 | let mut oversized = vec![b'x'; MAX_SSE_LINE_BYTES + 1]; |
| 284 | assert_eq!(flush_sse_line(&mut oversized), Err(SseLineError::TooLong)); |
| 285 | assert!(oversized.is_empty()); |
| 286 | oversized.resize(MAX_SSE_LINE_BYTES + 1, b'x'); |
| 287 | oversized.push(b'\n'); |
| 288 | assert_eq!( |
| 289 | next_sse_line(&mut oversized, false), |
| 290 | Err(SseLineError::TooLong) |
| 291 | ); |
| 292 | assert!(oversized.is_empty()); |
| 293 | |
| 294 | // The cap applies per line, not to a chunk containing several lines. |
| 295 | let mut multiple = at_limit; |
| 296 | multiple.extend_from_slice(b"\nnext\n"); |
| 297 | assert!(next_sse_line(&mut multiple, false).unwrap().is_some()); |
| 298 | assert_eq!( |
| 299 | next_sse_line(&mut multiple, false).unwrap().as_deref(), |
| 300 | Some("next") |
| 301 | ); |
| 302 | } |
| 303 | |
| 304 | #[test] |
| 305 | fn parse_usage_scenario() { |
| 306 | // 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 |
| 307 | // from parse_usage_reads_deepseek_cache_and_reasoning_tokens |
| 308 | { |
| 309 | let usage = parse_usage(Some(&json!({ |
| 310 | "prompt_tokens": 100, |
| 311 | "completion_tokens": 20, |
| 312 | "prompt_cache_hit_tokens": 70, |
| 313 | "prompt_cache_miss_tokens": 30, |
| 314 | "completion_tokens_details": { |
| 315 | "reasoning_tokens": 12 |
| 316 | } |
| 317 | }))); |
| 318 | |
| 319 | assert_eq!(usage.input_tokens, 100); |
| 320 | assert_eq!(usage.output_tokens, 20); |
| 321 | assert_eq!(usage.prompt_cache_hit_tokens, Some(70)); |
| 322 | assert_eq!(usage.prompt_cache_miss_tokens, Some(30)); |
| 323 | assert_eq!(usage.reasoning_tokens, Some(12)); |
| 324 | } |
| 325 | // from parse_usage_saturates_every_u64_token_field |
| 326 | { |
| 327 | let usage = parse_usage(Some(&json!({ |
| 328 | "input_tokens": u64::MAX, |
| 329 | "output_tokens": u64::MAX, |
| 330 | "prompt_cache_hit_tokens": u64::MAX, |
| 331 | "prompt_cache_miss_tokens": u64::MAX, |
| 332 | "completion_tokens_details": { "reasoning_tokens": u64::MAX }, |
| 333 | "server_tool_use": { |
| 334 | "code_execution_requests": u64::MAX, |
| 335 | "tool_search_requests": u64::MAX |
| 336 | } |
| 337 | }))); |
| 338 | assert_eq!(usage.input_tokens, u32::MAX); |
| 339 | assert_eq!(usage.output_tokens, u32::MAX); |
| 340 | assert_eq!(usage.prompt_cache_hit_tokens, Some(u32::MAX)); |
| 341 | assert_eq!(usage.prompt_cache_miss_tokens, Some(u32::MAX)); |
| 342 | assert_eq!(usage.reasoning_tokens, Some(u32::MAX)); |
| 343 | let server = usage.server_tool_use.expect("server usage"); |
| 344 | assert_eq!(server.code_execution_requests, Some(u32::MAX)); |
| 345 | assert_eq!(server.tool_search_requests, Some(u32::MAX)); |
| 346 | } |
| 347 | // from parse_usage_counts_reasoning_tokens_when_completion_tokens_are_zero |
| 348 | { |
| 349 | let usage = parse_usage(Some(&json!({ |
| 350 | "prompt_tokens": 100, |
| 351 | "completion_tokens": 0, |
| 352 | "completion_tokens_details": { |
| 353 | "reasoning_tokens": 12 |
| 354 | } |
| 355 | }))); |
| 356 | |
| 357 | assert_eq!(usage.input_tokens, 100); |
| 358 | assert_eq!(usage.output_tokens, 12); |
| 359 | assert_eq!(usage.reasoning_tokens, Some(12)); |
| 360 | assert!( |
| 361 | crate::pricing::calculate_turn_cost_from_usage("deepseek-v4-pro", &usage) |
| 362 | .expect("DeepSeek V4 Pro pricing should apply") |
| 363 | > 0.0 |
| 364 | ); |
| 365 | } |
| 366 | // from parse_usage_derives_completion_tokens_from_total_tokens_when_needed |
| 367 | { |
| 368 | let usage = parse_usage(Some(&json!({ |
| 369 | "prompt_tokens": 100, |
| 370 | "total_tokens": 125, |
| 371 | "prompt_cache_hit_tokens": 70, |
| 372 | "prompt_cache_miss_tokens": 30 |
| 373 | }))); |
| 374 | |
| 375 | assert_eq!(usage.input_tokens, 100); |
| 376 | assert_eq!(usage.output_tokens, 25); |
| 377 | assert_eq!(usage.prompt_cache_hit_tokens, Some(70)); |
| 378 | assert_eq!(usage.prompt_cache_miss_tokens, Some(30)); |
| 379 | } |
| 380 | // from parse_usage_reads_v4_prompt_tokens_details_cached_tokens |
| 381 | { |
| 382 | let usage = parse_usage(Some(&json!({ |
| 383 | "prompt_tokens": 4000, |
| 384 | "completion_tokens": 20, |
| 385 | "prompt_tokens_details": { |
| 386 | "cached_tokens": 3000 |
| 387 | } |
| 388 | }))); |
| 389 | |
| 390 | assert_eq!(usage.input_tokens, 4000); |
| 391 | assert_eq!(usage.output_tokens, 20); |
| 392 | assert_eq!(usage.prompt_cache_hit_tokens, Some(3000)); |
| 393 | assert_eq!(usage.prompt_cache_miss_tokens, Some(1000)); |
| 394 | } |
| 395 | // from parse_usage_infers_cache_miss_from_selected_hit_source |
| 396 | { |
| 397 | let usage = parse_usage(Some(&json!({ |
| 398 | "prompt_tokens": 4000, |
| 399 | "completion_tokens": 20, |
| 400 | "prompt_cache_hit_tokens": 3000, |
| 401 | "prompt_tokens_details": { |
| 402 | "cached_tokens": 1000 |
| 403 | } |
| 404 | }))); |
| 405 | |
| 406 | assert_eq!(usage.input_tokens, 4000); |
| 407 | assert_eq!(usage.prompt_cache_hit_tokens, Some(3000)); |
| 408 | assert_eq!(usage.prompt_cache_miss_tokens, Some(1000)); |
| 409 | } |
| 410 | } |
| 411 | |
| 412 | /// Real-shaped Chat-Completions usage payloads from the three providers most |
| 413 | /// likely to report reasoning tokens, carried end-to-end into pricing. |
| 414 | /// |
| 415 | /// Two invariants hold for every fixture: `reasoning_tokens <= output_tokens`, |
| 416 | /// and pricing never adds reasoning on top of output — dropping the reasoning |
| 417 | /// field entirely must not change the cost by a single cent. |
| 418 | #[test] |
| 419 | fn reasoning_parser_fixtures_never_exceed_or_add_to_billable_output() { |
| 420 | use crate::config::ProviderKind; |
| 421 | use crate::pricing::{calculate_turn_cost_estimate_for_provider, token_usage_for_pricing}; |
| 422 | |
| 423 | // (label, provider, model, payload) |
| 424 | let fixtures: [(&str, ProviderKind, &str, serde_json::Value); 3] = [ |
| 425 | ( |
| 426 | "moonshot", |
| 427 | ProviderKind::Moonshot, |
| 428 | "kimi-k2.7-code", |
| 429 | json!({ |
| 430 | "prompt_tokens": 30_000, |
| 431 | "completion_tokens": 2_400, |
| 432 | "total_tokens": 32_400, |
| 433 | "prompt_tokens_details": { "cached_tokens": 24_000 }, |
| 434 | "completion_tokens_details": { "reasoning_tokens": 1_900 } |
| 435 | }), |
| 436 | ), |
| 437 | ( |
| 438 | "minimax", |
| 439 | ProviderKind::Minimax, |
| 440 | "minimax-m3", |
| 441 | json!({ |
| 442 | "prompt_tokens": 12_000, |
| 443 | "completion_tokens": 3_000, |
| 444 | "total_tokens": 15_000, |
| 445 | "prompt_tokens_details": { "cached_tokens": 4_000 }, |
| 446 | "completion_tokens_details": { "reasoning_tokens": 2_950 } |
| 447 | }), |
| 448 | ), |
| 449 | ( |
| 450 | "openrouter", |
| 451 | ProviderKind::Openrouter, |
| 452 | "qwen/qwen3.7-plus", |
| 453 | json!({ |
| 454 | "prompt_tokens": 8_000, |
| 455 | "completion_tokens": 1_500, |
| 456 | "total_tokens": 9_500, |
| 457 | "prompt_tokens_details": { "cached_tokens": 2_000 }, |
| 458 | "completion_tokens_details": { "reasoning_tokens": 1_500 } |
| 459 | }), |
| 460 | ), |
| 461 | ]; |
| 462 | |
| 463 | for (label, provider, model, payload) in fixtures { |
| 464 | let usage = parse_usage(Some(&payload)); |
| 465 | let reasoning = usage.reasoning_tokens.expect("fixture reports reasoning"); |
| 466 | |
| 467 | // Invariant 1: reasoning is a subset of the billed completion count. |
| 468 | assert!( |
| 469 | reasoning <= usage.output_tokens, |
| 470 | "{label}: reasoning {reasoning} exceeds output {}", |
| 471 | usage.output_tokens |
| 472 | ); |
| 473 | // Billable output is exactly the reported completion count. |
| 474 | let classes = token_usage_for_pricing(&usage); |
| 475 | assert_eq!( |
| 476 | classes.output, |
| 477 | u64::from(usage.output_tokens), |
| 478 | "{label}: reasoning leaked into billable output" |
| 479 | ); |
| 480 | |
| 481 | // Invariant 2: pricing does not add reasoning a second time. The same |
| 482 | // usage with the reasoning field removed must cost the same. |
| 483 | let without = codewhale_models::Usage { |
| 484 | reasoning_tokens: None, |
| 485 | ..usage.clone() |
| 486 | }; |
| 487 | assert_eq!( |
| 488 | calculate_turn_cost_estimate_for_provider(provider, model, &usage), |
| 489 | calculate_turn_cost_estimate_for_provider(provider, model, &without), |
| 490 | "{label}: reasoning changed the price" |
| 491 | ); |
| 492 | } |
| 493 | } |
| 494 | |
| 495 | /// A payload claiming more reasoning than output contradicts the subset |
| 496 | /// invariant. That is broken telemetry, so the field is discarded — and it |
| 497 | /// must never become extra billable output. |
| 498 | #[test] |
| 499 | fn pathological_reasoning_above_output_is_rejected_not_billed() { |
| 500 | let usage = parse_usage(Some(&json!({ |
| 501 | "prompt_tokens": 1_000, |
| 502 | "completion_tokens": 100, |
| 503 | "completion_tokens_details": { "reasoning_tokens": 5_000 } |
| 504 | }))); |
| 505 | |
| 506 | assert_eq!(usage.output_tokens, 100, "output stays as reported"); |
| 507 | assert_eq!( |
| 508 | usage.reasoning_tokens, None, |
| 509 | "impossible reasoning telemetry is dropped rather than trusted" |
| 510 | ); |
| 511 | let classes = crate::pricing::token_usage_for_pricing(&usage); |
| 512 | assert_eq!(classes.output, 100); |
| 513 | |
| 514 | // `completion_tokens: 0` with reasoning present is the *legitimate* |
| 515 | // shape this filter must not break: providers that report only reasoning |
| 516 | // set output from it, keeping reasoning == output. |
| 517 | let zero_output = parse_usage(Some(&json!({ |
| 518 | "prompt_tokens": 1_000, |
| 519 | "completion_tokens": 0, |
| 520 | "completion_tokens_details": { "reasoning_tokens": 12 } |
| 521 | }))); |
| 522 | assert_eq!(zero_output.output_tokens, 12); |
| 523 | assert_eq!(zero_output.reasoning_tokens, Some(12)); |
| 524 | } |
| 525 | |
| 526 | fn mid_char_split(text: &str, ch: char) -> usize { |
| 527 | let needle = ch.to_string(); |
| 528 | let start = text |
| 529 | .as_bytes() |
| 530 | .windows(needle.len()) |
| 531 | .position(|window| window == needle.as_bytes()) |
| 532 | .unwrap_or_else(|| panic!("{ch:?} present in {text:?}")); |
| 533 | start + 1 |
| 534 | } |
| 535 | |
| 536 | #[test] |
| 537 | fn take_sse_scenario() { |
| 538 | // 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 |
| 539 | // from take_sse_line_preserves_multibyte_split_across_reads |
| 540 | { |
| 541 | // "你好" streamed so the 3-byte '好' straddles a read boundary. |
| 542 | let full = "data: 你好\n"; |
| 543 | let bytes = full.as_bytes(); |
| 544 | let split = mid_char_split(full, '好'); |
| 545 | let mut buffer: Vec<u8> = Vec::new(); |
| 546 | // First read: no complete line yet. |
| 547 | buffer.extend_from_slice(&bytes[..split]); |
| 548 | assert_eq!(take_sse_line(&mut buffer).expect("valid prefix"), None); |
| 549 | // Second read completes the line; '好' must be intact, not U+FFFD. |
| 550 | buffer.extend_from_slice(&bytes[split..]); |
| 551 | let line = take_sse_line(&mut buffer) |
| 552 | .expect("valid utf-8") |
| 553 | .expect("a complete line"); |
| 554 | assert_eq!(line, "data: 你好"); |
| 555 | assert!(!line.contains('\u{FFFD}'), "multibyte char was corrupted"); |
| 556 | assert_eq!(extract_sse_data_value(&line), Some("你好")); |
| 557 | // Buffer fully drained. |
| 558 | assert!(buffer.is_empty()); |
| 559 | } |
| 560 | // from take_sse_line_returns_none_without_newline |
| 561 | { |
| 562 | let mut buffer = b"data: partial".to_vec(); |
| 563 | assert_eq!(take_sse_line(&mut buffer).expect("valid utf-8"), None); |
| 564 | assert_eq!(buffer, b"data: partial"); |
| 565 | } |
| 566 | // from take_sse_line_reassembles_cjk_and_rejects_invalid_bytes |
| 567 | { |
| 568 | let full = "data: 测试中文\n"; |
| 569 | let split = mid_char_split(full, '试'); |
| 570 | let mut buffer = full.as_bytes()[..split].to_vec(); |
| 571 | assert_eq!(take_sse_line(&mut buffer).expect("valid prefix"), None); |
| 572 | buffer.extend_from_slice(&full.as_bytes()[split..]); |
| 573 | let line = take_sse_line(&mut buffer) |
| 574 | .expect("valid utf-8") |
| 575 | .expect("complete line"); |
| 576 | assert_eq!(line, "data: 测试中文"); |
| 577 | assert!(!line.contains('\u{FFFD}')); |
| 578 | |
| 579 | let mut invalid = b"data: ok".to_vec(); |
| 580 | invalid.push(0xFF); |
| 581 | invalid.push(b'\n'); |
| 582 | let err = take_sse_line(&mut invalid).expect_err("invalid bytes must fail closed"); |
| 583 | assert!(!err.to_string().contains('\u{FFFD}')); |
| 584 | assert_eq!(err, SseLineError::InvalidUtf8 { valid_up_to: 8 }); |
| 585 | assert!( |
| 586 | invalid.is_empty(), |
| 587 | "invalid line is consumed so retries cannot loop" |
| 588 | ); |
| 589 | } |
| 590 | // from take_sse_line_rejects_invalid_bytes_without_replacement |
| 591 | { |
| 592 | let mut buffer = b"data: ok".to_vec(); |
| 593 | buffer.push(0xFF); |
| 594 | buffer.extend_from_slice(b"\n"); |
| 595 | let err = take_sse_line(&mut buffer).expect_err("0xFF is not UTF-8"); |
| 596 | assert_eq!(err, SseLineError::InvalidUtf8 { valid_up_to: 8 }); |
| 597 | assert!(!err.to_string().contains('\u{FFFD}')); |
| 598 | assert!(buffer.is_empty(), "invalid line must be drained"); |
| 599 | } |
| 600 | } |
| 601 | |
| 602 | #[test] |
| 603 | fn flush_sse_scenario() { |
| 604 | // 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 |
| 605 | // from flush_sse_line_reassembles_cjk_and_rejects_invalid_bytes |
| 606 | { |
| 607 | let text = "data: 你好世界"; |
| 608 | let split = mid_char_split(text, '好'); |
| 609 | let mut buffer = text.as_bytes()[..split].to_vec(); |
| 610 | assert_eq!(take_sse_line(&mut buffer).expect("no newline yet"), None); |
| 611 | buffer.extend_from_slice(&text.as_bytes()[split..]); |
| 612 | let line = flush_sse_line(&mut buffer) |
| 613 | .expect("valid utf-8") |
| 614 | .expect("unterminated tail"); |
| 615 | assert_eq!(line, "data: 你好世界"); |
| 616 | assert!(!line.contains('\u{FFFD}')); |
| 617 | assert!(buffer.is_empty()); |
| 618 | assert_eq!(flush_sse_line(&mut buffer).expect("empty"), None); |
| 619 | |
| 620 | let mut invalid = vec![0x80, 0xBF]; |
| 621 | let err = flush_sse_line(&mut invalid).expect_err("invalid flush must fail closed"); |
| 622 | assert!(!err.to_string().contains('\u{FFFD}')); |
| 623 | assert_eq!(err, SseLineError::InvalidUtf8 { valid_up_to: 0 }); |
| 624 | assert!(invalid.is_empty()); |
| 625 | } |
| 626 | // from flush_sse_line_preserves_unterminated_cjk |
| 627 | { |
| 628 | let mut buffer = "data: 你好".as_bytes().to_vec(); |
| 629 | let line = flush_sse_line(&mut buffer) |
| 630 | .expect("valid utf-8") |
| 631 | .expect("residual line"); |
| 632 | assert_eq!(line, "data: 你好"); |
| 633 | assert!(!line.contains('\u{FFFD}')); |
| 634 | assert!(buffer.is_empty()); |
| 635 | } |
| 636 | // from flush_sse_line_rejects_truncated_multibyte_sequence |
| 637 | { |
| 638 | let mut buffer = "data: ".as_bytes().to_vec(); |
| 639 | buffer.extend_from_slice(&"好".as_bytes()[..2]); |
| 640 | let err = flush_sse_line(&mut buffer).expect_err("truncated UTF-8"); |
| 641 | assert_eq!(err, SseLineError::InvalidUtf8 { valid_up_to: 6 }); |
| 642 | assert!(!err.to_string().contains('\u{FFFD}')); |
| 643 | assert!(buffer.is_empty()); |
| 644 | } |
| 645 | } |
| 646 | |
| 647 | #[test] |
| 648 | fn decode_sse_line_bytes_rejects_invalid_without_replacement() { |
| 649 | let ok = decode_sse_line_bytes("data: 你好".as_bytes()).expect("valid"); |
| 650 | assert_eq!(ok, "data: 你好"); |
| 651 | assert!(!ok.contains('\u{FFFD}')); |
| 652 | |
| 653 | let err = decode_sse_line_bytes(&[0xFF]).expect_err("bare 0xFF is invalid"); |
| 654 | assert!(!err.to_string().contains('\u{FFFD}')); |
| 655 | assert_eq!(err, SseLineError::InvalidUtf8 { valid_up_to: 0 }); |
| 656 | } |
| 657 | |
| 658 | #[test] |
| 659 | fn extract_sse_scenario() { |
| 660 | // 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 |
| 661 | // from extract_sse_data_value_accepts_optional_space |
| 662 | { |
| 663 | assert_eq!( |
| 664 | extract_sse_data_value("data: {\"ok\":true}"), |
| 665 | Some("{\"ok\":true}") |
| 666 | ); |
| 667 | assert_eq!( |
| 668 | extract_sse_data_value("data:{\"ok\":true}"), |
| 669 | Some("{\"ok\":true}") |
| 670 | ); |
| 671 | } |
| 672 | // from extract_sse_data_value_handles_done_marker |
| 673 | { |
| 674 | assert_eq!(extract_sse_data_value("data: [DONE]"), Some("[DONE]")); |
| 675 | assert_eq!(extract_sse_data_value("data:[DONE]"), Some("[DONE]")); |
| 676 | } |
| 677 | // from extract_sse_data_value_rejects_non_data_lines |
| 678 | { |
| 679 | assert_eq!(extract_sse_data_value("event: message"), None); |
| 680 | assert_eq!(extract_sse_data_value(": heartbeat"), None); |
| 681 | } |
| 682 | } |
| 683 | } |
| 684 |