From fe8437ccf84c314eb9e93c2847e5d58a900fb588 Mon Sep 17 00:00:00 2001 From: mkoushni Date: Wed, 19 Aug 2026 15:07:47 +0300 Subject: [PATCH 1/4] fix(token_usage): recover from SSE overflow instead of silently dropping usage Token accounting stopped permanently after a single oversized SSE event or JSON response, clearing all working state (including the terminal usage event) with no signal that data was lost. - The shared SSE scanner (used by token_usage and a2a) now discards only the oversized event and resumes at the next event boundary, instead of aborting the whole stream. A terminal usage event arriving after an oversized one is now captured. - JSON and any residual SSE overflow (e.g. the usage event itself being oversized) now set an explicit token.status=overflow metadata key and Praxis-Token-Status response header, so billing consumers cannot mistake missing counts for zero usage. - max_body_bytes/max_scratch_bytes are now configurable per token_count filter instance instead of fixed at 1 MiB / 64 KiB. Fixes #674 Signed-off-by: mkoushni --- filters/src/agentic/a2a/mod.rs | 29 +++-- filters/src/agentic/a2a/sse.rs | 147 ++++++++++++++++++------- filters/src/agentic/a2a/tests.rs | 31 ++++-- filters/src/token_usage/count.rs | 113 ++++++++++++++++--- filters/src/token_usage/count/tests.rs | 125 +++++++++++++++++++-- filters/src/token_usage/headers.rs | 79 +++++++++++-- filters/src/token_usage/mod.rs | 16 +++ 7 files changed, 455 insertions(+), 85 deletions(-) diff --git a/filters/src/agentic/a2a/mod.rs b/filters/src/agentic/a2a/mod.rs index 409838023b..e15dd37e59 100644 --- a/filters/src/agentic/a2a/mod.rs +++ b/filters/src/agentic/a2a/mod.rs @@ -845,16 +845,15 @@ fn process_sse_response_chunk( try_extract_task_from_sse_payload(payload, ctx, store, config); } - if result.overflowed { + if result.dropped_events > 0 { debug!( - scratch_bytes = state.scratch_bytes, + dropped_events = result.dropped_events, max_bytes = config.max_response_body_bytes, - "SSE scratch exceeds capture limit, disabling streaming capture" + "SSE event exceeded capture limit, discarding and resuming at next event boundary" ); - clear_sse_capture_metadata(ctx); - } else { - save_sse_scan_state(ctx, &state); } + + save_sse_scan_state(ctx, &state); } /// Drains any incomplete SSE event buffered in `filter_metadata` at @@ -896,6 +895,7 @@ fn try_extract_task_from_sse_payload( /// Reconstructs scanner state from hex-encoded `filter_metadata` keys. /// Metadata bypasses the 256-byte dynamic-value helper because the /// scanner buffers raw SSE line/data bytes that can exceed that limit. +#[expect(clippy::too_many_lines, reason = "36 lines; sequential per-field metadata reads")] fn load_sse_scan_state(ctx: &HttpFilterContext<'_>) -> sse::SseScanState { let line_buf = ctx .filter_metadata @@ -925,12 +925,19 @@ fn load_sse_scan_state(ctx: &HttpFilterContext<'_>) -> sse::SseScanState { .and_then(|v| v.parse().ok()) .unwrap_or(0); + let skip = match ctx.filter_metadata.get("a2a.response.sse_skip").map(String::as_str) { + Some("line_empty") => sse::SkipPhase::LineEmptySoFar, + Some("line_has_content") => sse::SkipPhase::LineHasContent, + _ => sse::SkipPhase::NotSkipping, + }; + sse::SseScanState { line_buf, data_buf, has_data, prev_cr, scratch_bytes, + skip, } } @@ -951,6 +958,14 @@ fn save_sse_scan_state(ctx: &mut HttpFilterContext<'_>, state: &sse::SseScanStat "a2a.response.sse_scratch_bytes".to_owned(), state.scratch_bytes.to_string(), ); + + let skip = match state.skip { + sse::SkipPhase::NotSkipping => "not_skipping", + sse::SkipPhase::LineEmptySoFar => "line_empty", + sse::SkipPhase::LineHasContent => "line_has_content", + }; + ctx.filter_metadata + .insert("a2a.response.sse_skip".to_owned(), skip.to_owned()); } /// Hex-encodes raw bytes into a metadata value, or removes the key if empty. @@ -975,10 +990,10 @@ fn clear_sse_capture_metadata(ctx: &mut HttpFilterContext<'_>) { ctx.filter_metadata.remove("a2a.response.sse_has_data"); ctx.filter_metadata.remove("a2a.response.sse_prev_cr"); ctx.filter_metadata.remove("a2a.response.sse_scratch_bytes"); + ctx.filter_metadata.remove("a2a.response.sse_skip"); ctx.filter_metadata.remove("a2a.response.cluster"); } - /// Build a `JsonRpcConfig` for the shared parser with A2A-appropriate defaults. fn build_json_rpc_config(max_body_bytes: usize) -> JsonRpcConfig { use praxis_filter::builtins::http::payload_processing::json_rpc::config::{BatchPolicy, JsonRpcHeaders}; diff --git a/filters/src/agentic/a2a/sse.rs b/filters/src/agentic/a2a/sse.rs index 1721053e2b..9e3bbc35db 100644 --- a/filters/src/agentic/a2a/sse.rs +++ b/filters/src/agentic/a2a/sse.rs @@ -40,6 +40,29 @@ pub(crate) struct SseScanState { /// Total scratch bytes consumed (`line_buf` + `data_buf`). pub scratch_bytes: usize, + + /// Progress discarding an event that exceeded `max_scratch_bytes`, + /// if any. See [`SkipPhase`]. + pub skip: SkipPhase, +} + +/// Progress discarding an oversized event's bytes without buffering them. +/// +/// An event exceeding `max_scratch_bytes` is dropped rather than +/// retained: its bytes are discarded as they arrive, and scanning +/// resumes normally once a blank line (the SSE event boundary) is +/// found — tracked here without ever buffering the discarded content. +#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)] +pub(crate) enum SkipPhase { + /// Not discarding; bytes are buffered normally. + #[default] + NotSkipping, + /// Discarding; the line since the last newline has been empty so + /// far. A newline seen now means a blank line — the boundary. + LineEmptySoFar, + /// Discarding; the line since the last newline has had at least + /// one byte. A newline seen now just ends that (non-blank) line. + LineHasContent, } // ----------------------------------------------------------------------------- @@ -47,30 +70,34 @@ pub(crate) struct SseScanState { // ----------------------------------------------------------------------------- /// Outcome of [`scan_sse_chunk`]. -/// -/// Always returns completed payloads, even when the scratch limit is -/// exceeded partway through a chunk. The caller should process the -/// payloads first, then check `overflowed` to decide whether to -/// continue or disable capture. pub(crate) struct SseScanResult { /// Completed `data:` payloads dispatched during this chunk. pub payloads: Vec>, - /// Whether the scratch limit was exceeded. When `true`, the - /// caller should store routes from `payloads`, then clear - /// capture state and stop scanning further chunks. - pub overflowed: bool, + /// Number of events discarded during this chunk for exceeding + /// `max_scratch_bytes`. Informational only — scanning already + /// recovered at the next event boundary, so the caller does not + /// need to take any corrective action. + pub dropped_events: usize, } // ----------------------------------------------------------------------------- // Scanning // ----------------------------------------------------------------------------- -/// Process one SSE chunk, returning completed `data:` payloads and -/// an overflow flag. +/// Process one SSE chunk, returning completed `data:` payloads and the +/// number of oversized events discarded. +/// +/// An event that would exceed `max_scratch_bytes` is dropped rather than +/// buffered: its bytes are discarded without ever being retained, and +/// scanning resumes normally at the next blank-line event boundary. This +/// bounds memory use per event without losing later events on the same +/// stream — in particular, a terminal usage/summary event arriving after +/// an oversized one is still captured. #[expect(clippy::too_many_lines, reason = "linear byte-processing loop")] pub(crate) fn scan_sse_chunk(state: &mut SseScanState, chunk: &[u8], max_scratch_bytes: usize) -> SseScanResult { let mut payloads = Vec::new(); + let mut dropped_events = 0_usize; let mut i = 0; // If previous chunk ended with CR and this starts with LF, consume it @@ -81,39 +108,67 @@ pub(crate) fn scan_sse_chunk(state: &mut SseScanState, chunk: &[u8], max_scratch state.prev_cr = false; while let Some(&b) = chunk.get(i) { - if b == b'\n' || b == b'\r' { - process_line(&state.line_buf, &mut state.data_buf, &mut state.has_data, &mut payloads); - state.line_buf.clear(); - - // CRLF within the same chunk: skip the LF. - if b == b'\r' { - if let Some(&next) = chunk.get(i + 1) { - if next == b'\n' { - i += 1; - } + let is_newline = b == b'\n' || b == b'\r'; + + match state.skip { + SkipPhase::NotSkipping => { + if is_newline { + process_line(&state.line_buf, &mut state.data_buf, &mut state.has_data, &mut payloads); + state.line_buf.clear(); + } else { + state.line_buf.push(b); + } + }, + SkipPhase::LineEmptySoFar => { + state.skip = if is_newline { + SkipPhase::NotSkipping // blank line: event boundary found } else { - state.prev_cr = true; + SkipPhase::LineHasContent + }; + }, + SkipPhase::LineHasContent => { + if is_newline { + state.skip = SkipPhase::LineEmptySoFar; + } + }, + } + + // CRLF within the same chunk: skip the LF. Applies regardless of + // skip state, since it tracks raw byte position, not buffering. + if b == b'\r' { + if let Some(&next) = chunk.get(i + 1) { + if next == b'\n' { + i += 1; } + } else { + state.prev_cr = true; } - } else { - state.line_buf.push(b); } - state.scratch_bytes = state.line_buf.len() + state.data_buf.len(); - if state.scratch_bytes > max_scratch_bytes { - return SseScanResult { - payloads, - overflowed: true, - }; + if state.skip == SkipPhase::NotSkipping { + state.scratch_bytes = state.line_buf.len() + state.data_buf.len(); + if state.scratch_bytes > max_scratch_bytes { + state.line_buf.clear(); + state.data_buf.clear(); + state.has_data = false; + state.scratch_bytes = 0; + // The line since the last newline already has content + // unless byte `b` itself was the newline starting it. + state.skip = if is_newline { + SkipPhase::LineEmptySoFar + } else { + SkipPhase::LineHasContent + }; + dropped_events += 1; + } } i += 1; } - state.scratch_bytes = state.line_buf.len() + state.data_buf.len(); SseScanResult { payloads, - overflowed: false, + dropped_events, } } @@ -429,14 +484,18 @@ mod tests { // ------------------------------------------------------------------------- #[test] - fn scratch_overflow_sets_overflowed_flag() { + fn scratch_overflow_recovers_at_next_event_boundary() { let mut state = SseScanState::default(); - let chunk = b"data: a]very long line that exceeds the limit\n"; + let chunk = b"data: this-line-is-too-long-to-fit\n\ndata: ok\n\n"; let result = scan_sse_chunk(&mut state, chunk, 10); - assert!(result.overflowed, "exceeding scratch limit should set overflowed"); - assert!(result.payloads.is_empty(), "no completed events before overflow"); + assert_eq!( + result.dropped_events, 1, + "oversized event should be dropped, not abort scanning" + ); + assert_eq!(result.payloads.len(), 1, "scanning should resume for the next event"); + assert_eq!(result.payloads[0], b"ok"); } #[test] @@ -447,7 +506,7 @@ mod tests { let result = scan_sse_chunk(&mut state, chunk, 15); - assert!(result.overflowed, "should overflow on the second event"); + assert_eq!(result.dropped_events, 1, "should drop the second, oversized event"); assert_eq!( result.payloads.len(), 1, @@ -456,6 +515,20 @@ mod tests { assert_eq!(result.payloads[0], b"ok"); } + #[test] + fn scratch_overflow_recovers_across_chunk_boundary() { + let mut state = SseScanState::default(); + + let r1 = scan_sse_chunk(&mut state, b"data: aaaaaaaaaaaaaaaa", 10); + assert_eq!(r1.dropped_events, 1, "overflow should be detected mid-line"); + assert!(r1.payloads.is_empty()); + + let r2 = scan_sse_chunk(&mut state, b"aaaaaaaa\n\ndata: ok\n\n", 10); + assert_eq!(r2.dropped_events, 0, "no further drops once the boundary is found"); + assert_eq!(r2.payloads.len(), 1, "next event should be captured normally"); + assert_eq!(r2.payloads[0], b"ok"); + } + #[test] fn scratch_resets_after_dispatch() { let mut state = SseScanState::default(); diff --git a/filters/src/agentic/a2a/tests.rs b/filters/src/agentic/a2a/tests.rs index 5b6d193857..d4714c624c 100644 --- a/filters/src/agentic/a2a/tests.rs +++ b/filters/src/agentic/a2a/tests.rs @@ -2810,7 +2810,7 @@ async fn sse_streaming_capture_invalid_json_does_not_fail() { } #[tokio::test] -async fn sse_streaming_capture_oversized_scratch_clears_state() { +async fn sse_streaming_capture_oversized_scratch_drops_event_and_continues() { let filter = make_task_routing_filter_with_config( r#"{"on_invalid": "continue", "task_routing": {"enabled": true, "max_response_body_bytes": 32}}"#, ); @@ -2826,9 +2826,13 @@ async fn sse_streaming_capture_oversized_scratch_clears_state() { assert!( store.get_by_task_id("task-big").is_none(), - "oversized SSE scratch should skip capture" + "oversized SSE event should be discarded, not parsed" + ); + assert_eq!( + ctx.get_metadata("a2a.response.sse_capture_enabled"), + Some("true"), + "capture should remain enabled after discarding one oversized event" ); - assert_sse_capture_cleared(&ctx); } #[tokio::test] @@ -3178,7 +3182,7 @@ async fn sse_streaming_capture_artifact_update_event() { // ----------------------------------------------------------------------------- #[tokio::test] -async fn valid_task_before_oversized_sse_frame_stores_route_and_clears_capture() { +async fn valid_task_before_and_after_oversized_sse_frame_are_both_stored() { let filter = make_task_routing_filter_with_config( r#"{"on_invalid": "continue", "task_routing": {"enabled": true, "max_response_body_bytes": 128}}"#, ); @@ -3189,11 +3193,13 @@ async fn valid_task_before_oversized_sse_frame_stores_route_and_clears_capture() seed_sse_capture(&mut ctx); // First event (~110 bytes) fits within 128-byte limit. - // Second event (~200+ bytes) overflows. + // Second event (~200+ bytes) overflows and is dropped. + // Third event fits again, verifying capture recovered. let padding = "x".repeat(100); let sse = format!( "data: {{\"jsonrpc\":\"2.0\",\"id\":1,\"result\":{{\"task\":{{\"id\":\"task-pre-overflow\",\"status\":{{\"state\":\"TASK_STATE_WORKING\"}}}}}}}}\n\n\ - data: {{\"jsonrpc\":\"2.0\",\"id\":2,\"result\":{{\"task\":{{\"id\":\"task-big\",\"status\":{{\"state\":\"TASK_STATE_WORKING\"}},\"padding\":\"{padding}\"}}}}}}\n\n" + data: {{\"jsonrpc\":\"2.0\",\"id\":2,\"result\":{{\"task\":{{\"id\":\"task-big\",\"status\":{{\"state\":\"TASK_STATE_WORKING\"}},\"padding\":\"{padding}\"}}}}}}\n\n\ + data: {{\"jsonrpc\":\"2.0\",\"id\":3,\"result\":{{\"task\":{{\"id\":\"task-post-overflow\",\"status\":{{\"state\":\"TASK_STATE_WORKING\"}}}}}}}}\n\n" ); let sse_data = sse.as_bytes(); let mut body = Some(Bytes::from(sse_data.to_vec())); @@ -3202,9 +3208,17 @@ async fn valid_task_before_oversized_sse_frame_stores_route_and_clears_capture() assert_eq!( store.get_by_task_id("task-pre-overflow").as_deref(), Some("agent-a"), - "completed event before overflow should still be captured" + "event before overflow should still be captured" + ); + assert!( + store.get_by_task_id("task-big").is_none(), + "the oversized event itself should be discarded" + ); + assert_eq!( + store.get_by_task_id("task-post-overflow").as_deref(), + Some("agent-a"), + "capture should recover and store the event after the oversized one" ); - assert_sse_capture_cleared(&ctx); } // ----------------------------------------------------------------------------- @@ -3389,7 +3403,6 @@ fn scan_json_balance_is_a_no_op_once_invalid() { assert!(!state.is_complete(), "an invalid scan must never report complete"); } - #[test] fn scan_json_balance_persists_across_chunk_calls() { let full = br#"{"jsonrpc":"2.0","id":1,"result":{"task":{"id":"t","status":{"state":"TASK_STATE_WORKING"}}}}"#; diff --git a/filters/src/token_usage/count.rs b/filters/src/token_usage/count.rs index c3822e56cb..29bb938283 100644 --- a/filters/src/token_usage/count.rs +++ b/filters/src/token_usage/count.rs @@ -38,7 +38,7 @@ use tracing::{debug, trace}; use super::{ TokenUsage, providers::{parse_anthropic, parse_bedrock, parse_google, parse_openai}, - set_token_usage, streaming, + set_token_status_overflow, set_token_usage, streaming, }; use crate::agentic::a2a::sse; @@ -85,6 +85,15 @@ const META_SSE_PREV_CR: &str = "token_count.sse_prev_cr"; /// Metadata key for SSE scanner scratch byte count. const META_SSE_SCRATCH: &str = "token_count.sse_scratch_bytes"; +/// Metadata key for the SSE scanner's oversized-event skip phase. +const META_SSE_SKIP: &str = "token_count.sse_skip"; + +/// Metadata key recording that at least one oversized SSE event was +/// discarded during this stream. Used by [`finalize_streaming_counts`] +/// to distinguish "no usage events at all" from "usage may have been in +/// a discarded event." +const META_SSE_DROPPED: &str = "token_count.sse_dropped"; + /// Bedrock `InvokeModel` response header carrying the input token count. const HEADER_BEDROCK_INPUT: &str = "x-amzn-bedrock-input-token-count"; @@ -101,6 +110,26 @@ const HEADER_BEDROCK_OUTPUT: &str = "x-amzn-bedrock-output-token-count"; struct TokenCountConfig { /// AI provider whose response format to parse. provider: ProviderKind, + + /// Maximum bytes to buffer for a non-streaming JSON response before + /// giving up on locating its usage field. + #[serde(default = "default_max_body_bytes")] + max_body_bytes: usize, + + /// Maximum scratch bytes (buffered line + in-progress event data) for + /// the SSE scanner before an event is discarded as oversized. + #[serde(default = "default_max_scratch_bytes")] + max_scratch_bytes: usize, +} + +/// Default for [`TokenCountConfig::max_body_bytes`]. +fn default_max_body_bytes() -> usize { + DEFAULT_MAX_BODY_BYTES +} + +/// Default for [`TokenCountConfig::max_scratch_bytes`]. +fn default_max_scratch_bytes() -> usize { + DEFAULT_MAX_SCRATCH_BYTES } /// Provider and transport shape selected by the `token_count` configuration. @@ -160,12 +189,20 @@ impl ProviderKind { /// ```yaml /// filter: token_count /// provider: openai # openai | anthropic | google | bedrock | bedrock_invoke_model | azure +/// max_body_bytes: 1048576 # optional, JSON capture limit +/// max_scratch_bytes: 65536 # optional, SSE per-event capture limit /// ``` /// /// [`filter_metadata`]: HttpFilterContext::filter_metadata pub struct TokenCountFilter { /// Which provider's response format to parse. provider: ProviderKind, + + /// Maximum bytes to buffer for a non-streaming JSON response. + max_body_bytes: usize, + + /// Maximum scratch bytes per SSE event before it is discarded. + max_scratch_bytes: usize, } impl TokenCountFilter { @@ -177,7 +214,11 @@ impl TokenCountFilter { pub fn from_config(config: &serde_yaml::Value) -> Result, FilterError> { let cfg: TokenCountConfig = parse_filter_config("token_count", config)?; - Ok(Box::new(Self { provider: cfg.provider })) + Ok(Box::new(Self { + provider: cfg.provider, + max_body_bytes: cfg.max_body_bytes, + max_scratch_bytes: cfg.max_scratch_bytes, + })) } } @@ -255,8 +296,8 @@ impl HttpFilter for TokenCountFilter { let mode = ctx.get_metadata(META_MODE).map(str::to_owned); match mode.as_deref() { - Some("sse") => handle_sse_body(ctx, body, end_of_stream, self.provider), - Some("json") => handle_json_body(ctx, body, end_of_stream, self.provider), + Some("sse") => handle_sse_body(ctx, body, end_of_stream, self.provider, self.max_scratch_bytes), + Some("json") => handle_json_body(ctx, body, end_of_stream, self.provider, self.max_body_bytes), _ => {}, } @@ -303,15 +344,24 @@ fn extract_bedrock_headers(ctx: &mut HttpFilterContext<'_>) { // ----------------------------------------------------------------------------- /// Accumulate JSON body chunks and extract token usage on completion. +/// +/// A response larger than `max_body_bytes` cannot be parsed for usage +/// (it may appear anywhere, including near the end), so extraction is +/// abandoned. Rather than leaving no signal at all, an explicit overflow +/// status is recorded so consumers do not mistake this for a genuine +/// zero-usage response. fn handle_json_body( ctx: &mut HttpFilterContext<'_>, body: &Option, end_of_stream: bool, provider: ProviderKind, + max_body_bytes: usize, ) { if let Some(chunk) = body.as_ref() - && !accumulate_response_hex(ctx, chunk, DEFAULT_MAX_BODY_BYTES) + && !accumulate_response_hex(ctx, chunk, max_body_bytes) { + set_token_status_overflow(ctx); + debug!("JSON response exceeded token_count capture limit, usage is unavailable"); clear_all_metadata(ctx); return; } @@ -345,24 +395,33 @@ fn handle_json_body( // ----------------------------------------------------------------------------- /// Process SSE body chunks and extract token usage incrementally. -fn handle_sse_body(ctx: &mut HttpFilterContext<'_>, body: &Option, end_of_stream: bool, provider: ProviderKind) { +/// +/// An oversized event is discarded by the scanner without aborting the +/// stream, so a terminal usage event arriving after it is still captured. +/// Final counts are only ever written at `end_of_stream`. +fn handle_sse_body( + ctx: &mut HttpFilterContext<'_>, + body: &Option, + end_of_stream: bool, + provider: ProviderKind, + max_scratch_bytes: usize, +) { if let Some(chunk) = body.as_ref() { let mut state = load_sse_scan_state(ctx); - let result = sse::scan_sse_chunk(&mut state, chunk, DEFAULT_MAX_SCRATCH_BYTES); + let result = sse::scan_sse_chunk(&mut state, chunk, max_scratch_bytes); for payload in &result.payloads { process_sse_payload(ctx, payload, provider); } - if result.overflowed { + if result.dropped_events > 0 { debug!( - scratch_bytes = state.scratch_bytes, - "SSE scratch exceeds limit, finalizing token counts" + dropped_events = result.dropped_events, + "SSE event exceeded scratch limit, discarding and resuming at next event boundary" ); - finalize_streaming_counts(ctx); - clear_all_metadata(ctx); - return; + ctx.filter_metadata + .insert(META_SSE_DROPPED.to_owned(), "true".to_owned()); } save_sse_scan_state(ctx, &state); @@ -437,9 +496,18 @@ fn merge_accumulated_count(ctx: &mut HttpFilterContext<'_>, key: &str, value: u6 } /// Write accumulated streaming counts to the well-known metadata keys. +/// +/// If no counts were ever accumulated but an oversized event was +/// discarded during the stream, the terminal usage event itself may +/// have been the one dropped — an explicit overflow status is recorded +/// so this is distinguishable from a genuine zero-usage response. fn finalize_streaming_counts(ctx: &mut HttpFilterContext<'_>) { let has_accumulated = ctx.filter_metadata.contains_key(META_INPUT) || ctx.filter_metadata.contains_key(META_OUTPUT); if !has_accumulated { + if ctx.filter_metadata.contains_key(META_SSE_DROPPED) { + set_token_status_overflow(ctx); + debug!("stream had oversized events discarded and produced no usage; marking status overflow"); + } return; } @@ -486,12 +554,24 @@ fn load_sse_scan_state(ctx: &HttpFilterContext<'_>) -> sse::SseScanState { .and_then(|v| v.parse().ok()) .unwrap_or(0); + let skip = decode_skip_phase(ctx.filter_metadata.get(META_SSE_SKIP).map(String::as_str)); + sse::SseScanState { line_buf, data_buf, has_data, prev_cr, scratch_bytes, + skip, + } +} + +/// Inverse of the encoding written by [`save_sse_scan_state`]. +fn decode_skip_phase(value: Option<&str>) -> sse::SkipPhase { + match value { + Some("line_empty") => sse::SkipPhase::LineEmptySoFar, + Some("line_has_content") => sse::SkipPhase::LineHasContent, + _ => sse::SkipPhase::NotSkipping, } } @@ -510,6 +590,13 @@ fn save_sse_scan_state(ctx: &mut HttpFilterContext<'_>, state: &sse::SseScanStat ); ctx.filter_metadata .insert(META_SSE_SCRATCH.to_owned(), state.scratch_bytes.to_string()); + + let skip = match state.skip { + sse::SkipPhase::NotSkipping => "not_skipping", + sse::SkipPhase::LineEmptySoFar => "line_empty", + sse::SkipPhase::LineHasContent => "line_has_content", + }; + ctx.filter_metadata.insert(META_SSE_SKIP.to_owned(), skip.to_owned()); } // ----------------------------------------------------------------------------- diff --git a/filters/src/token_usage/count/tests.rs b/filters/src/token_usage/count/tests.rs index 01d757dada..d3caf06a3c 100644 --- a/filters/src/token_usage/count/tests.rs +++ b/filters/src/token_usage/count/tests.rs @@ -538,8 +538,32 @@ async fn sse_final_event_without_trailing_blank_line() { // SSE: Scratch Overflow // ----------------------------------------------------------------------------- +/// Regression test for #674: an oversized event must not disable +/// extraction for the rest of the stream. The scanner recovers at the +/// next event boundary, so a terminal usage event arriving afterward is +/// still captured. #[tokio::test] -async fn sse_overflow_finalizes_partial_counts() { +async fn sse_oversized_event_recovers_and_captures_terminal_usage() { + let mut events = b"data: ".to_vec(); + events.extend(std::iter::repeat_n(b'x', DEFAULT_MAX_SCRATCH_BYTES + 1)); + events.extend_from_slice(b"\n\n"); + events.extend_from_slice( + b"data: {\"usage\":{\"prompt_tokens\":10,\"completion_tokens\":20,\"total_tokens\":30}}\n\n", + ); + + let (input, output, total) = run_sse_extraction(ProviderKind::OpenAi, &events).await; + + assert_eq!( + input.as_deref(), + Some("10"), + "terminal usage should be captured after the oversized event is discarded" + ); + assert_eq!(output.as_deref(), Some("20")); + assert_eq!(total.as_deref(), Some("30")); +} + +#[tokio::test] +async fn sse_overflow_does_not_finalize_mid_stream() { let filter = make_filter(ProviderKind::OpenAi); let req = crate::test_utils::make_request(http::Method::POST, "/v1/chat/completions"); let mut ctx = crate::test_utils::make_filter_context(&req); @@ -558,12 +582,46 @@ async fn sse_overflow_finalizes_partial_counts() { let mut body2 = Some(Bytes::from(overflow_chunk)); drop(filter.on_response_body(&mut ctx, &mut body2, false).unwrap()); + assert!( + ctx.get_metadata("token.input").is_none(), + "overflow mid-stream should not finalize early; already-accumulated counts wait for end_of_stream" + ); + + let mut body3 = Some(Bytes::from_static(b"\n\n")); + drop(filter.on_response_body(&mut ctx, &mut body3, true).unwrap()); + assert_eq!(ctx.get_metadata("token.input"), Some("10")); assert_eq!(ctx.get_metadata("token.output"), Some("20")); assert_eq!(ctx.get_metadata("token.total"), Some("30")); assert_no_working_metadata(&ctx); } +#[tokio::test] +async fn sse_overflow_with_no_usage_sets_overflow_status() { + let filter = make_filter(ProviderKind::OpenAi); + let req = crate::test_utils::make_request(http::Method::POST, "/v1/chat/completions"); + let mut ctx = crate::test_utils::make_filter_context(&req); + + let mut resp = make_response_with_content_type("text/event-stream"); + ctx.response_header = Some(&mut resp); + drop(filter.on_response(&mut ctx).await.unwrap()); + ctx.response_header = None; + + let overflow_chunk = vec![b'x'; DEFAULT_MAX_SCRATCH_BYTES + 1]; + let mut body = Some(Bytes::from(overflow_chunk)); + drop(filter.on_response_body(&mut ctx, &mut body, true).unwrap()); + + assert!( + ctx.get_metadata("token.input").is_none(), + "no usage was ever captured, so no counts should be written" + ); + assert_eq!( + ctx.get_metadata("token.status"), + Some("overflow"), + "billing consumers must not mistake this for zero usage" + ); +} + // ----------------------------------------------------------------------------- // Body Mode Without on_response // ----------------------------------------------------------------------------- @@ -591,8 +649,10 @@ fn on_response_body_noop_without_mode() { // Buffer Overflow // ----------------------------------------------------------------------------- +/// Regression test for #674: overflow must not be indistinguishable +/// from a genuine zero-usage response. #[tokio::test] -async fn json_overflow_sets_nothing() { +async fn json_overflow_sets_explicit_status_not_zero() { let filter = make_filter(ProviderKind::OpenAi); let req = crate::test_utils::make_request(http::Method::POST, "/v1/chat/completions"); let mut ctx = crate::test_utils::make_filter_context(&req); @@ -608,14 +668,55 @@ async fn json_overflow_sets_nothing() { assert!( ctx.get_metadata("token.input").is_none(), - "overflow should not set tokens" + "overflow should not set counts" ); - let working_keys: Vec<_> = ctx - .filter_metadata - .keys() - .filter(|k| k.starts_with(META_PREFIX)) - .collect(); - assert!(working_keys.is_empty(), "overflow should clear all working metadata"); + assert_eq!( + ctx.get_metadata("token.status"), + Some("overflow"), + "billing consumers must not mistake missing counts for zero usage" + ); + assert_no_working_metadata(&ctx); +} + +/// A valid response with usage near the end but larger than the +/// configured limit exhibits the same overflow signal, not silence. +#[tokio::test] +async fn json_valid_usage_beyond_configured_limit_sets_overflow_status() { + let mut filter = make_filter(ProviderKind::OpenAi); + filter.max_body_bytes = 64; + let req = crate::test_utils::make_request(http::Method::POST, "/v1/chat/completions"); + let mut ctx = crate::test_utils::make_filter_context(&req); + + let mut resp = make_response_with_content_type("application/json"); + ctx.response_header = Some(&mut resp); + drop(filter.on_response(&mut ctx).await.unwrap()); + ctx.response_header = None; + + let padding = "x".repeat(200); + let json = format!( + r#"{{"id":"resp-1","padding":"{padding}","usage":{{"prompt_tokens":10,"completion_tokens":20,"total_tokens":30}}}}"# + ); + let mut body = Some(Bytes::from(json.into_bytes())); + drop(filter.on_response_body(&mut ctx, &mut body, true).unwrap()); + + assert!(ctx.get_metadata("token.input").is_none()); + assert_eq!(ctx.get_metadata("token.status"), Some("overflow")); +} + +#[test] +fn max_body_bytes_configurable_via_yaml() { + let config: serde_yaml::Value = serde_yaml::from_str("provider: openai\nmax_body_bytes: 4096").unwrap(); + let filter = TokenCountFilter::from_config(&config).unwrap(); + + assert_eq!(filter.name(), "token_count"); +} + +#[test] +fn max_scratch_bytes_configurable_via_yaml() { + let config: serde_yaml::Value = serde_yaml::from_str("provider: openai\nmax_scratch_bytes: 512").unwrap(); + let filter = TokenCountFilter::from_config(&config).unwrap(); + + assert_eq!(filter.name(), "token_count"); } // ----------------------------------------------------------------------------- @@ -842,7 +943,11 @@ fn on_response_body_noop_for_bedrock_invoke_model() { use std::fmt::Write as _; fn make_filter(provider: ProviderKind) -> TokenCountFilter { - TokenCountFilter { provider } + TokenCountFilter { + provider, + max_body_bytes: DEFAULT_MAX_BODY_BYTES, + max_scratch_bytes: DEFAULT_MAX_SCRATCH_BYTES, + } } fn make_response_with_content_type(ct: &str) -> Response { diff --git a/filters/src/token_usage/headers.rs b/filters/src/token_usage/headers.rs index 400bec3c06..335fe139ae 100644 --- a/filters/src/token_usage/headers.rs +++ b/filters/src/token_usage/headers.rs @@ -8,7 +8,7 @@ use http::header::{HeaderName, HeaderValue}; use praxis_filter::{EmptyFilterConfig, FilterAction, FilterError, HttpFilter, HttpFilterContext, parse_filter_config}; use tracing::trace; -use super::{META_TOKEN_INPUT, META_TOKEN_OUTPUT, META_TOKEN_TOTAL}; +use super::{META_TOKEN_INPUT, META_TOKEN_OUTPUT, META_TOKEN_STATUS, META_TOKEN_TOTAL}; // ----------------------------------------------------------------------------- // Constants @@ -23,13 +23,21 @@ const HEADER_TOKEN_OUTPUT: HeaderName = HeaderName::from_static("praxis-token-ou /// Response header carrying the total token count. const HEADER_TOKEN_TOTAL: HeaderName = HeaderName::from_static("praxis-token-total"); +/// Response header signaling that usage could not be captured (e.g. +/// "overflow"), present instead of or alongside the count headers so +/// consumers cannot mistake missing counts for zero usage. +const HEADER_TOKEN_STATUS: HeaderName = HeaderName::from_static("praxis-token-status"); + // ----------------------------------------------------------------------------- // TokenUsageHeadersFilter // ----------------------------------------------------------------------------- /// Injects `Praxis-Token-Input`, `Praxis-Token-Output`, and /// `Praxis-Token-Total` headers into downstream responses when -/// token usage data is present in [`filter_metadata`]. +/// token usage data is present in [`filter_metadata`]. Also injects +/// `Praxis-Token-Status` when usage capture failed (e.g. overflow), +/// so an unavailable count is never silently indistinguishable from +/// a genuine zero. /// /// Reads token counts written by upstream filters and exposes them /// as HTTP response headers for infrastructure-level consumption @@ -92,25 +100,36 @@ impl HttpFilter for TokenUsageHeadersFilter { /// response headers. No-op when response headers are absent or /// no token data has been written by upstream filters. fn inject_usage_headers(ctx: &mut HttpFilterContext<'_>) { - let (Some(input), Some(output), Some(total)) = ( + let status = ctx.get_metadata(META_TOKEN_STATUS).map(str::to_owned); + let counts = match ( ctx.get_metadata(META_TOKEN_INPUT).map(str::to_owned), ctx.get_metadata(META_TOKEN_OUTPUT).map(str::to_owned), ctx.get_metadata(META_TOKEN_TOTAL).map(str::to_owned), - ) else { + ) { + (Some(input), Some(output), Some(total)) => Some((input, output, total)), + _ => None, + }; + + if status.is_none() && counts.is_none() { trace!("no token usage metadata found, skipping header injection"); return; - }; + } let Some(resp) = ctx.response_header.as_mut() else { return; }; ctx.response_headers_modified = true; - let headers = &mut resp.headers; - insert_token_header(headers, &HEADER_TOKEN_INPUT, &input); - insert_token_header(headers, &HEADER_TOKEN_OUTPUT, &output); - insert_token_header(headers, &HEADER_TOKEN_TOTAL, &total); + + if let Some(status) = &status { + insert_token_header(headers, &HEADER_TOKEN_STATUS, status); + } + if let Some((input, output, total)) = &counts { + insert_token_header(headers, &HEADER_TOKEN_INPUT, input); + insert_token_header(headers, &HEADER_TOKEN_OUTPUT, output); + insert_token_header(headers, &HEADER_TOKEN_TOTAL, total); + } trace!("injected token usage response headers"); } @@ -283,6 +302,48 @@ mod tests { ); } + #[tokio::test] + async fn injects_status_header_on_overflow_without_counts() { + let filter = TokenUsageHeadersFilter; + let req = crate::test_utils::make_request(http::Method::POST, "/v1/chat/completions"); + let mut ctx = crate::test_utils::make_filter_context(&req); + ctx.set_metadata(META_TOKEN_STATUS, "overflow".to_owned()); + + let mut resp = crate::test_utils::make_response(); + ctx.response_header = Some(&mut resp); + drop(filter.on_response(&mut ctx).await.unwrap()); + ctx.response_header = None; + + assert!( + ctx.response_headers_modified, + "status alone should still modify headers" + ); + assert_eq!(resp.headers["praxis-token-status"], "overflow"); + assert!( + resp.headers.get("praxis-token-input").is_none(), + "no counts available on overflow" + ); + } + + #[tokio::test] + async fn injects_counts_without_status_header_on_success() { + let filter = TokenUsageHeadersFilter; + let req = crate::test_utils::make_request(http::Method::POST, "/v1/chat/completions"); + let mut ctx = crate::test_utils::make_filter_context(&req); + set_token_metadata(&mut ctx, "10", "20", "30"); + + let mut resp = crate::test_utils::make_response(); + ctx.response_header = Some(&mut resp); + drop(filter.on_response(&mut ctx).await.unwrap()); + ctx.response_header = None; + + assert!( + resp.headers.get("praxis-token-status").is_none(), + "no status header on a normal, successful extraction" + ); + assert_eq!(resp.headers["praxis-token-input"], "10"); + } + #[test] fn from_config_succeeds_with_empty_config() { let config = serde_yaml::Value::Mapping(serde_yaml::Mapping::new()); diff --git a/filters/src/token_usage/mod.rs b/filters/src/token_usage/mod.rs index 51858f7e1c..37fa27861d 100644 --- a/filters/src/token_usage/mod.rs +++ b/filters/src/token_usage/mod.rs @@ -25,6 +25,16 @@ const META_TOKEN_OUTPUT: &str = "token.output"; /// Metadata key for the total token count. const META_TOKEN_TOTAL: &str = "token.total"; +/// Metadata key signaling that usage could not be captured because the +/// response exceeded the configured capture limit. Absent on success, +/// including when the provider genuinely reported no usage — consumers +/// must not treat "no counts" the same as "counts unavailable." +const META_TOKEN_STATUS: &str = "token.status"; + +/// Value of [`META_TOKEN_STATUS`] when capture was abandoned due to +/// exceeding the configured size limit. +const TOKEN_STATUS_OVERFLOW: &str = "overflow"; + /// Unified token usage extracted from an AI provider response. #[derive(Debug, Clone, Copy, PartialEq, Eq)] struct TokenUsage { @@ -72,3 +82,9 @@ fn set_token_usage(ctx: &mut HttpFilterContext<'_>, input: u64, output: u64, tot ctx.set_metadata(META_TOKEN_OUTPUT, output.to_string()); ctx.set_metadata(META_TOKEN_TOTAL, total.to_string()); } + +/// Marks token usage as unavailable due to exceeding the capture limit, +/// distinguishable from a genuine zero-usage response. +fn set_token_status_overflow(ctx: &mut HttpFilterContext<'_>) { + ctx.set_metadata(META_TOKEN_STATUS, TOKEN_STATUS_OVERFLOW.to_owned()); +} From 836b9c7c8c5e912863f2d094ca10c2f380fa615f Mon Sep 17 00:00:00 2001 From: mkoushni Date: Sun, 23 Aug 2026 15:42:12 +0300 Subject: [PATCH 2/4] fix(token_usage): share SkipPhase encoding and reject zero capture limits Give A2A and token_count a single SkipPhase metadata codec, reject zero max_body_bytes/max_scratch_bytes, and regenerate filter docs so lint CI matches the new config fields. Signed-off-by: mkoushni --- docs/filters/reference.md | 2 +- docs/filters/token_count.md | 4 +++ docs/filters/token_usage_headers.md | 2 +- filters/src/agentic/a2a/mod.rs | 13 ++----- filters/src/agentic/a2a/sse.rs | 47 ++++++++++++++++++++++++++ filters/src/token_usage/count.rs | 31 +++++++---------- filters/src/token_usage/count/tests.rs | 24 +++++++++++++ 7 files changed, 91 insertions(+), 32 deletions(-) diff --git a/docs/filters/reference.md b/docs/filters/reference.md index 32452bc2f3..a4637806d2 100644 --- a/docs/filters/reference.md +++ b/docs/filters/reference.md @@ -91,4 +91,4 @@ see the [Praxis core filter reference][core-ref]. | Filter | Description | |--------|-------------| | [`token_count`](token_count.md) | Extracts token usage from AI inference responses and writes unified counts to [`filter_metadata`]. | -| [`token_usage_headers`](token_usage_headers.md) | Injects `Praxis-Token-Input`, `Praxis-Token-Output`, and `Praxis-Token-Total` headers into downstream responses when token usage data is present in [`filter_metadata`]. | +| [`token_usage_headers`](token_usage_headers.md) | Injects `Praxis-Token-Input`, `Praxis-Token-Output`, and `Praxis-Token-Total` headers into downstream responses when token usage data is present in [`filter_metadata`]. Also injects `Praxis-Token-Status` when usage capture failed (e.g. overflow), so an unavailable count is never silently indistinguishable from a genuine zero. | diff --git a/docs/filters/token_count.md b/docs/filters/token_count.md index fe2b0228bd..49952a7a69 100644 --- a/docs/filters/token_count.md +++ b/docs/filters/token_count.md @@ -14,10 +14,14 @@ Supports both streaming (SSE) and non-streaming (JSON) responses across five pro | Field | Type | Required | Description | |-------|------|---------|-------------| | `provider` | `openai` \| `anthropic` \| `google` \| `bedrock` \| `bedrock_invoke_model` \| `azure` | yes | AI provider whose response format to parse. | +| `max_body_bytes` | integer | no | Maximum bytes to buffer for a non-streaming JSON response before giving up on locating its usage field. Must be greater than 0. | +| `max_scratch_bytes` | integer | no | Maximum scratch bytes (buffered line + in-progress event data) for the SSE scanner before an event is discarded as oversized. Must be greater than 0. | ## Example ```yaml filter: token_count provider: openai # openai | anthropic | google | bedrock | bedrock_invoke_model | azure +max_body_bytes: 1048576 # optional, JSON capture limit +max_scratch_bytes: 65536 # optional, SSE per-event capture limit ``` diff --git a/docs/filters/token_usage_headers.md b/docs/filters/token_usage_headers.md index 6ce6bd7a21..5456771bba 100644 --- a/docs/filters/token_usage_headers.md +++ b/docs/filters/token_usage_headers.md @@ -3,7 +3,7 @@ # `token_usage_headers` -Injects `Praxis-Token-Input`, `Praxis-Token-Output`, and `Praxis-Token-Total` headers into downstream responses when token usage data is present in [`filter_metadata`]. +Injects `Praxis-Token-Input`, `Praxis-Token-Output`, and `Praxis-Token-Total` headers into downstream responses when token usage data is present in [`filter_metadata`]. Also injects `Praxis-Token-Status` when usage capture failed (e.g. overflow), so an unavailable count is never silently indistinguishable from a genuine zero. ## Configuration Notes diff --git a/filters/src/agentic/a2a/mod.rs b/filters/src/agentic/a2a/mod.rs index bc40c18634..0b13cc08b6 100644 --- a/filters/src/agentic/a2a/mod.rs +++ b/filters/src/agentic/a2a/mod.rs @@ -925,11 +925,7 @@ fn load_sse_scan_state(ctx: &HttpFilterContext<'_>) -> sse::SseScanState { .and_then(|v| v.parse().ok()) .unwrap_or(0); - let skip = match ctx.filter_metadata.get("a2a.response.sse_skip").map(String::as_str) { - Some("line_empty") => sse::SkipPhase::LineEmptySoFar, - Some("line_has_content") => sse::SkipPhase::LineHasContent, - _ => sse::SkipPhase::NotSkipping, - }; + let skip = sse::SkipPhase::from_metadata_str(ctx.filter_metadata.get("a2a.response.sse_skip").map(String::as_str)); sse::SseScanState { line_buf, @@ -959,13 +955,8 @@ fn save_sse_scan_state(ctx: &mut HttpFilterContext<'_>, state: &sse::SseScanStat state.scratch_bytes.to_string(), ); - let skip = match state.skip { - sse::SkipPhase::NotSkipping => "not_skipping", - sse::SkipPhase::LineEmptySoFar => "line_empty", - sse::SkipPhase::LineHasContent => "line_has_content", - }; ctx.filter_metadata - .insert("a2a.response.sse_skip".to_owned(), skip.to_owned()); + .insert("a2a.response.sse_skip".to_owned(), state.skip.as_str().to_owned()); } /// Hex-encodes raw bytes into a metadata value, or removes the key if empty. diff --git a/filters/src/agentic/a2a/sse.rs b/filters/src/agentic/a2a/sse.rs index 9e3bbc35db..68b7601746 100644 --- a/filters/src/agentic/a2a/sse.rs +++ b/filters/src/agentic/a2a/sse.rs @@ -65,6 +65,27 @@ pub(crate) enum SkipPhase { LineHasContent, } +impl SkipPhase { + /// Encode this phase for `filter_metadata` persistence. + pub(crate) const fn as_str(self) -> &'static str { + match self { + Self::NotSkipping => "not_skipping", + Self::LineEmptySoFar => "line_empty", + Self::LineHasContent => "line_has_content", + } + } + + /// Decode a persisted metadata value. Missing or unknown strings + /// map to [`Self::NotSkipping`] so a restart resumes buffering. + pub(crate) fn from_metadata_str(s: Option<&str>) -> Self { + match s { + Some("line_empty") => Self::LineEmptySoFar, + Some("line_has_content") => Self::LineHasContent, + _ => Self::NotSkipping, + } + } +} + // ----------------------------------------------------------------------------- // SseScanResult // ----------------------------------------------------------------------------- @@ -529,6 +550,32 @@ mod tests { assert_eq!(r2.payloads[0], b"ok"); } + #[test] + fn skip_phase_round_trips_through_metadata_strings() { + assert_eq!(SkipPhase::NotSkipping.as_str(), "not_skipping"); + assert_eq!(SkipPhase::LineEmptySoFar.as_str(), "line_empty"); + assert_eq!(SkipPhase::LineHasContent.as_str(), "line_has_content"); + + assert_eq!( + SkipPhase::from_metadata_str(Some("not_skipping")), + SkipPhase::NotSkipping + ); + assert_eq!( + SkipPhase::from_metadata_str(Some("line_empty")), + SkipPhase::LineEmptySoFar + ); + assert_eq!( + SkipPhase::from_metadata_str(Some("line_has_content")), + SkipPhase::LineHasContent + ); + assert_eq!(SkipPhase::from_metadata_str(None), SkipPhase::NotSkipping); + assert_eq!( + SkipPhase::from_metadata_str(Some("unknown")), + SkipPhase::NotSkipping, + "unknown values should resume normal buffering" + ); + } + #[test] fn scratch_resets_after_dispatch() { let mut state = SseScanState::default(); diff --git a/filters/src/token_usage/count.rs b/filters/src/token_usage/count.rs index 29bb938283..1a78a0a049 100644 --- a/filters/src/token_usage/count.rs +++ b/filters/src/token_usage/count.rs @@ -30,7 +30,8 @@ use std::fmt::Write as _; use async_trait::async_trait; use bytes::Bytes; use praxis_filter::{ - BodyAccess, BodyMode, FilterAction, FilterError, HttpFilter, HttpFilterContext, parse_filter_config, + BodyAccess, BodyMode, FilterAction, FilterError, HttpFilter, HttpFilterContext, + builtins::http::payload_processing::config_validation::validate_max_body_bytes, parse_filter_config, }; use serde::Deserialize; use tracing::{debug, trace}; @@ -112,12 +113,13 @@ struct TokenCountConfig { provider: ProviderKind, /// Maximum bytes to buffer for a non-streaming JSON response before - /// giving up on locating its usage field. + /// giving up on locating its usage field. Must be greater than 0. #[serde(default = "default_max_body_bytes")] max_body_bytes: usize, /// Maximum scratch bytes (buffered line + in-progress event data) for - /// the SSE scanner before an event is discarded as oversized. + /// the SSE scanner before an event is discarded as oversized. Must be + /// greater than 0. #[serde(default = "default_max_scratch_bytes")] max_scratch_bytes: usize, } @@ -213,6 +215,10 @@ impl TokenCountFilter { /// Returns [`FilterError`] if the YAML config is invalid. pub fn from_config(config: &serde_yaml::Value) -> Result, FilterError> { let cfg: TokenCountConfig = parse_filter_config("token_count", config)?; + validate_max_body_bytes("token_count", cfg.max_body_bytes)?; + if cfg.max_scratch_bytes == 0 { + return Err("token_count: 'max_scratch_bytes' must be greater than 0".into()); + } Ok(Box::new(Self { provider: cfg.provider, @@ -554,7 +560,7 @@ fn load_sse_scan_state(ctx: &HttpFilterContext<'_>) -> sse::SseScanState { .and_then(|v| v.parse().ok()) .unwrap_or(0); - let skip = decode_skip_phase(ctx.filter_metadata.get(META_SSE_SKIP).map(String::as_str)); + let skip = sse::SkipPhase::from_metadata_str(ctx.filter_metadata.get(META_SSE_SKIP).map(String::as_str)); sse::SseScanState { line_buf, @@ -566,15 +572,6 @@ fn load_sse_scan_state(ctx: &HttpFilterContext<'_>) -> sse::SseScanState { } } -/// Inverse of the encoding written by [`save_sse_scan_state`]. -fn decode_skip_phase(value: Option<&str>) -> sse::SkipPhase { - match value { - Some("line_empty") => sse::SkipPhase::LineEmptySoFar, - Some("line_has_content") => sse::SkipPhase::LineHasContent, - _ => sse::SkipPhase::NotSkipping, - } -} - /// Persists scanner state back to `filter_metadata` for the next chunk. fn save_sse_scan_state(ctx: &mut HttpFilterContext<'_>, state: &sse::SseScanState) { set_hex_metadata(ctx, META_SSE_LINE_BUF, &state.line_buf); @@ -591,12 +588,8 @@ fn save_sse_scan_state(ctx: &mut HttpFilterContext<'_>, state: &sse::SseScanStat ctx.filter_metadata .insert(META_SSE_SCRATCH.to_owned(), state.scratch_bytes.to_string()); - let skip = match state.skip { - sse::SkipPhase::NotSkipping => "not_skipping", - sse::SkipPhase::LineEmptySoFar => "line_empty", - sse::SkipPhase::LineHasContent => "line_has_content", - }; - ctx.filter_metadata.insert(META_SSE_SKIP.to_owned(), skip.to_owned()); + ctx.filter_metadata + .insert(META_SSE_SKIP.to_owned(), state.skip.as_str().to_owned()); } // ----------------------------------------------------------------------------- diff --git a/filters/src/token_usage/count/tests.rs b/filters/src/token_usage/count/tests.rs index d3caf06a3c..b603c46674 100644 --- a/filters/src/token_usage/count/tests.rs +++ b/filters/src/token_usage/count/tests.rs @@ -719,6 +719,30 @@ fn max_scratch_bytes_configurable_via_yaml() { assert_eq!(filter.name(), "token_count"); } +#[test] +fn from_config_rejects_zero_max_body_bytes() { + let config: serde_yaml::Value = serde_yaml::from_str("provider: openai\nmax_body_bytes: 0").unwrap(); + match TokenCountFilter::from_config(&config) { + Err(err) => assert!( + err.to_string().contains("max_body_bytes"), + "should reject zero max_body_bytes: {err}" + ), + Ok(_) => panic!("should reject zero max_body_bytes"), + } +} + +#[test] +fn from_config_rejects_zero_max_scratch_bytes() { + let config: serde_yaml::Value = serde_yaml::from_str("provider: openai\nmax_scratch_bytes: 0").unwrap(); + match TokenCountFilter::from_config(&config) { + Err(err) => assert!( + err.to_string().contains("max_scratch_bytes"), + "should reject zero max_scratch_bytes: {err}" + ), + Ok(_) => panic!("should reject zero max_scratch_bytes"), + } +} + // ----------------------------------------------------------------------------- // Content-Type Helpers // ----------------------------------------------------------------------------- From e6077a1efae7259df9d62cd611cf489c490c4bf5 Mon Sep 17 00:00:00 2001 From: mkoushni Date: Mon, 24 Aug 2026 16:23:06 +0300 Subject: [PATCH 3/4] fix(token_usage): mark overflow when a dropped SSE event is the tail Partial Anthropic/Bedrock counts followed by an oversized terminal usage event now keep the captured maxima and set token.status=overflow, while a recovered usage event after a drop stays authoritative. Signed-off-by: mkoushni --- filters/src/agentic/a2a/mod.rs | 12 ++++- filters/src/agentic/a2a/sse.rs | 59 ++++++++++++++++++++++++ filters/src/token_usage/count.rs | 42 ++++++++++------- filters/src/token_usage/count/tests.rs | 64 ++++++++++++++++++++++++-- 4 files changed, 154 insertions(+), 23 deletions(-) diff --git a/filters/src/agentic/a2a/mod.rs b/filters/src/agentic/a2a/mod.rs index 0b13cc08b6..7cd8c1ef96 100644 --- a/filters/src/agentic/a2a/mod.rs +++ b/filters/src/agentic/a2a/mod.rs @@ -895,7 +895,7 @@ fn try_extract_task_from_sse_payload( /// Reconstructs scanner state from hex-encoded `filter_metadata` keys. /// Metadata bypasses the 256-byte dynamic-value helper because the /// scanner buffers raw SSE line/data bytes that can exceed that limit. -#[expect(clippy::too_many_lines, reason = "36 lines; sequential per-field metadata reads")] +#[expect(clippy::too_many_lines, reason = "sequential per-field metadata reads")] fn load_sse_scan_state(ctx: &HttpFilterContext<'_>) -> sse::SseScanState { let line_buf = ctx .filter_metadata @@ -926,6 +926,11 @@ fn load_sse_scan_state(ctx: &HttpFilterContext<'_>) -> sse::SseScanState { .unwrap_or(0); let skip = sse::SkipPhase::from_metadata_str(ctx.filter_metadata.get("a2a.response.sse_skip").map(String::as_str)); + let tail = sse::ScanTail::from_metadata_str( + ctx.filter_metadata + .get("a2a.response.sse_dropped_tail") + .map(String::as_str), + ); sse::SseScanState { line_buf, @@ -934,6 +939,7 @@ fn load_sse_scan_state(ctx: &HttpFilterContext<'_>) -> sse::SseScanState { prev_cr, scratch_bytes, skip, + tail, } } @@ -957,6 +963,10 @@ fn save_sse_scan_state(ctx: &mut HttpFilterContext<'_>, state: &sse::SseScanStat ctx.filter_metadata .insert("a2a.response.sse_skip".to_owned(), state.skip.as_str().to_owned()); + ctx.filter_metadata.insert( + "a2a.response.sse_dropped_tail".to_owned(), + state.tail.as_str().to_owned(), + ); } /// Hex-encodes raw bytes into a metadata value, or removes the key if empty. diff --git a/filters/src/agentic/a2a/sse.rs b/filters/src/agentic/a2a/sse.rs index 68b7601746..dc731f6b13 100644 --- a/filters/src/agentic/a2a/sse.rs +++ b/filters/src/agentic/a2a/sse.rs @@ -44,6 +44,10 @@ pub(crate) struct SseScanState { /// Progress discarding an event that exceeded `max_scratch_bytes`, /// if any. See [`SkipPhase`]. pub skip: SkipPhase, + + /// Most recently completed stream action. Distinguishes "drop then + /// recovered usage" from "usage then dropped terminal event." + pub tail: ScanTail, } /// Progress discarding an oversized event's bytes without buffering them. @@ -86,6 +90,37 @@ impl SkipPhase { } } +/// Whether the last completed SSE action dispatched a payload or dropped +/// an oversized event. Persisted across chunks so finalize can tell a +/// recovered terminal usage event from a dropped one. +#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)] +pub(crate) enum ScanTail { + /// Last completed action dispatched a `data:` payload. + #[default] + Payload, + /// Last completed action discarded an oversized event. + Dropped, +} + +impl ScanTail { + /// Encode this tail for `filter_metadata` persistence. + pub(crate) const fn as_str(self) -> &'static str { + match self { + Self::Payload => "payload", + Self::Dropped => "dropped", + } + } + + /// Decode a persisted metadata value. Missing or unknown strings + /// map to [`Self::Payload`]. + pub(crate) fn from_metadata_str(s: Option<&str>) -> Self { + match s { + Some("dropped") => Self::Dropped, + _ => Self::Payload, + } + } +} + // ----------------------------------------------------------------------------- // SseScanResult // ----------------------------------------------------------------------------- @@ -134,7 +169,11 @@ pub(crate) fn scan_sse_chunk(state: &mut SseScanState, chunk: &[u8], max_scratch match state.skip { SkipPhase::NotSkipping => { if is_newline { + let before = payloads.len(); process_line(&state.line_buf, &mut state.data_buf, &mut state.has_data, &mut payloads); + if payloads.len() > before { + state.tail = ScanTail::Payload; + } state.line_buf.clear(); } else { state.line_buf.push(b); @@ -181,6 +220,7 @@ pub(crate) fn scan_sse_chunk(state: &mut SseScanState, chunk: &[u8], max_scratch SkipPhase::LineHasContent }; dropped_events += 1; + state.tail = ScanTail::Dropped; } } @@ -199,6 +239,8 @@ pub(crate) fn scan_sse_chunk(state: &mut SseScanState, chunk: &[u8], max_scratch /// closing the connection, so the scanner must dispatch buffered state on /// `end_of_stream` rather than waiting for another `\n\n` boundary. pub(crate) fn flush_sse_state(state: &mut SseScanState, payloads: &mut Vec>) { + let before = payloads.len(); + if !state.line_buf.is_empty() { process_line(&state.line_buf, &mut state.data_buf, &mut state.has_data, payloads); state.line_buf.clear(); @@ -209,6 +251,10 @@ pub(crate) fn flush_sse_state(state: &mut SseScanState, payloads: &mut Vec before { + state.tail = ScanTail::Payload; + } + state.scratch_bytes = 0; } @@ -517,6 +563,10 @@ mod tests { ); assert_eq!(result.payloads.len(), 1, "scanning should resume for the next event"); assert_eq!(result.payloads[0], b"ok"); + assert!( + state.tail != ScanTail::Dropped, + "a payload after the drop means the tail is recovered usage, not a drop" + ); } #[test] @@ -534,6 +584,10 @@ mod tests { "first completed event should still be returned" ); assert_eq!(result.payloads[0], b"ok"); + assert!( + state.tail == ScanTail::Dropped, + "dropping the last event must leave dropped_tail set" + ); } #[test] @@ -543,11 +597,16 @@ mod tests { let r1 = scan_sse_chunk(&mut state, b"data: aaaaaaaaaaaaaaaa", 10); assert_eq!(r1.dropped_events, 1, "overflow should be detected mid-line"); assert!(r1.payloads.is_empty()); + assert_eq!(state.tail, ScanTail::Dropped, "mid-line overflow is a dropped tail"); let r2 = scan_sse_chunk(&mut state, b"aaaaaaaa\n\ndata: ok\n\n", 10); assert_eq!(r2.dropped_events, 0, "no further drops once the boundary is found"); assert_eq!(r2.payloads.len(), 1, "next event should be captured normally"); assert_eq!(r2.payloads[0], b"ok"); + assert!( + state.tail != ScanTail::Dropped, + "dispatching a later payload clears the dropped tail" + ); } #[test] diff --git a/filters/src/token_usage/count.rs b/filters/src/token_usage/count.rs index 1a78a0a049..cdc7367f5a 100644 --- a/filters/src/token_usage/count.rs +++ b/filters/src/token_usage/count.rs @@ -89,11 +89,10 @@ const META_SSE_SCRATCH: &str = "token_count.sse_scratch_bytes"; /// Metadata key for the SSE scanner's oversized-event skip phase. const META_SSE_SKIP: &str = "token_count.sse_skip"; -/// Metadata key recording that at least one oversized SSE event was -/// discarded during this stream. Used by [`finalize_streaming_counts`] -/// to distinguish "no usage events at all" from "usage may have been in -/// a discarded event." -const META_SSE_DROPPED: &str = "token_count.sse_dropped"; +/// Metadata key recording that the most recent SSE stream action was +/// discarding an oversized event. Used by [`finalize_streaming_counts`] +/// to mark overflow when the terminal usage event itself was dropped. +const META_SSE_DROPPED_TAIL: &str = "token_count.sse_dropped_tail"; /// Bedrock `InvokeModel` response header carrying the input token count. const HEADER_BEDROCK_INPUT: &str = "x-amzn-bedrock-input-token-count"; @@ -426,8 +425,6 @@ fn handle_sse_body( dropped_events = result.dropped_events, "SSE event exceeded scratch limit, discarding and resuming at next event boundary" ); - ctx.filter_metadata - .insert(META_SSE_DROPPED.to_owned(), "true".to_owned()); } save_sse_scan_state(ctx, &state); @@ -441,7 +438,7 @@ fn handle_sse_body( process_sse_payload(ctx, payload, provider); } - finalize_streaming_counts(ctx); + finalize_streaming_counts(ctx, state.tail == sse::ScanTail::Dropped); clear_all_metadata(ctx); } } @@ -503,17 +500,24 @@ fn merge_accumulated_count(ctx: &mut HttpFilterContext<'_>, key: &str, value: u6 /// Write accumulated streaming counts to the well-known metadata keys. /// -/// If no counts were ever accumulated but an oversized event was -/// discarded during the stream, the terminal usage event itself may -/// have been the one dropped — an explicit overflow status is recorded -/// so this is distinguishable from a genuine zero-usage response. -fn finalize_streaming_counts(ctx: &mut HttpFilterContext<'_>) { +/// If the stream ended after discarding an oversized event — including +/// the case where Anthropic/Bedrock partial counts were already stored +/// and the terminal usage event itself overflowed — an explicit overflow +/// status is recorded so this is distinguishable from a complete capture. +/// Recovered usage that arrived *after* an oversized event is treated as +/// authoritative and does not set overflow. +fn finalize_streaming_counts(ctx: &mut HttpFilterContext<'_>, dropped_tail: bool) { let has_accumulated = ctx.filter_metadata.contains_key(META_INPUT) || ctx.filter_metadata.contains_key(META_OUTPUT); + + if dropped_tail { + set_token_status_overflow(ctx); + debug!( + has_accumulated, + "stream ended after an oversized SSE event; marking status overflow" + ); + } + if !has_accumulated { - if ctx.filter_metadata.contains_key(META_SSE_DROPPED) { - set_token_status_overflow(ctx); - debug!("stream had oversized events discarded and produced no usage; marking status overflow"); - } return; } @@ -561,6 +565,7 @@ fn load_sse_scan_state(ctx: &HttpFilterContext<'_>) -> sse::SseScanState { .unwrap_or(0); let skip = sse::SkipPhase::from_metadata_str(ctx.filter_metadata.get(META_SSE_SKIP).map(String::as_str)); + let tail = sse::ScanTail::from_metadata_str(ctx.filter_metadata.get(META_SSE_DROPPED_TAIL).map(String::as_str)); sse::SseScanState { line_buf, @@ -569,6 +574,7 @@ fn load_sse_scan_state(ctx: &HttpFilterContext<'_>) -> sse::SseScanState { prev_cr, scratch_bytes, skip, + tail, } } @@ -590,6 +596,8 @@ fn save_sse_scan_state(ctx: &mut HttpFilterContext<'_>, state: &sse::SseScanStat ctx.filter_metadata .insert(META_SSE_SKIP.to_owned(), state.skip.as_str().to_owned()); + ctx.filter_metadata + .insert(META_SSE_DROPPED_TAIL.to_owned(), state.tail.as_str().to_owned()); } // ----------------------------------------------------------------------------- diff --git a/filters/src/token_usage/count/tests.rs b/filters/src/token_usage/count/tests.rs index b603c46674..5dcdec835a 100644 --- a/filters/src/token_usage/count/tests.rs +++ b/filters/src/token_usage/count/tests.rs @@ -541,7 +541,7 @@ async fn sse_final_event_without_trailing_blank_line() { /// Regression test for #674: an oversized event must not disable /// extraction for the rest of the stream. The scanner recovers at the /// next event boundary, so a terminal usage event arriving afterward is -/// still captured. +/// still captured and treated as authoritative (no overflow status). #[tokio::test] async fn sse_oversized_event_recovers_and_captures_terminal_usage() { let mut events = b"data: ".to_vec(); @@ -551,15 +551,29 @@ async fn sse_oversized_event_recovers_and_captures_terminal_usage() { b"data: {\"usage\":{\"prompt_tokens\":10,\"completion_tokens\":20,\"total_tokens\":30}}\n\n", ); - let (input, output, total) = run_sse_extraction(ProviderKind::OpenAi, &events).await; + let filter = make_filter(ProviderKind::OpenAi); + let req = crate::test_utils::make_request(http::Method::POST, "/v1/chat/completions"); + let mut ctx = crate::test_utils::make_filter_context(&req); + + let mut resp = make_response_with_content_type("text/event-stream"); + ctx.response_header = Some(&mut resp); + drop(filter.on_response(&mut ctx).await.unwrap()); + ctx.response_header = None; + + let mut body = Some(Bytes::from(events)); + drop(filter.on_response_body(&mut ctx, &mut body, true).unwrap()); assert_eq!( - input.as_deref(), + ctx.get_metadata("token.input"), Some("10"), "terminal usage should be captured after the oversized event is discarded" ); - assert_eq!(output.as_deref(), Some("20")); - assert_eq!(total.as_deref(), Some("30")); + assert_eq!(ctx.get_metadata("token.output"), Some("20")); + assert_eq!(ctx.get_metadata("token.total"), Some("30")); + assert!( + ctx.get_metadata("token.status").is_none(), + "recovered terminal usage after a dropped content event is authoritative" + ); } #[tokio::test] @@ -593,6 +607,11 @@ async fn sse_overflow_does_not_finalize_mid_stream() { assert_eq!(ctx.get_metadata("token.input"), Some("10")); assert_eq!(ctx.get_metadata("token.output"), Some("20")); assert_eq!(ctx.get_metadata("token.total"), Some("30")); + assert_eq!( + ctx.get_metadata("token.status"), + Some("overflow"), + "a drop after captured usage means later events (possibly a usage correction) were lost" + ); assert_no_working_metadata(&ctx); } @@ -622,6 +641,41 @@ async fn sse_overflow_with_no_usage_sets_overflow_status() { ); } +/// Anthropic/Bedrock split usage across events. If partial counts were +/// stored and the terminal usage event itself overflows, emit the +/// partials *and* `token.status = overflow` so they are not treated as +/// a complete capture. +#[tokio::test] +async fn sse_partial_then_oversized_terminal_event_sets_overflow_status() { + let mut filter = make_filter(ProviderKind::Anthropic); + filter.max_scratch_bytes = 256; + let req = crate::test_utils::make_request(http::Method::POST, "/v1/messages"); + let mut ctx = crate::test_utils::make_filter_context(&req); + let mut resp = make_response_with_content_type("text/event-stream"); + ctx.response_header = Some(&mut resp); + drop(filter.on_response(&mut ctx).await.unwrap()); + ctx.response_header = None; + + let mut events = b"data: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":25}}}\n\n".to_vec(); + events.extend_from_slice(b"data: "); + events.extend(std::iter::repeat_n(b'x', 257)); + events.extend_from_slice(b"\n\n"); + let mut body = Some(Bytes::from(events)); + drop(filter.on_response_body(&mut ctx, &mut body, true).unwrap()); + + assert_eq!(ctx.get_metadata("token.input"), Some("25"), "partial input kept"); + assert_eq!( + ctx.get_metadata("token.output"), + Some("0"), + "terminal output was dropped" + ); + assert_eq!( + ctx.get_metadata("token.status"), + Some("overflow"), + "partial maxima are not a complete capture when the terminal event overflowed" + ); +} + // ----------------------------------------------------------------------------- // Body Mode Without on_response // ----------------------------------------------------------------------------- From 7bf1e258af5f7b9e15432fac405cccc330631481 Mon Sep 17 00:00:00 2001 From: mkoushni Date: Mon, 24 Aug 2026 16:45:30 +0300 Subject: [PATCH 4/4] fix(token_usage): cap max_scratch_bytes at the shared 64 MiB ceiling Signed-off-by: mkoushni --- docs/filters/token_count.md | 4 ++-- filters/src/token_usage/count.rs | 23 +++++++++++++++++------ filters/src/token_usage/count/tests.rs | 24 ++++++++++++++++++++++++ 3 files changed, 43 insertions(+), 8 deletions(-) diff --git a/docs/filters/token_count.md b/docs/filters/token_count.md index 49952a7a69..76e562c10a 100644 --- a/docs/filters/token_count.md +++ b/docs/filters/token_count.md @@ -14,8 +14,8 @@ Supports both streaming (SSE) and non-streaming (JSON) responses across five pro | Field | Type | Required | Description | |-------|------|---------|-------------| | `provider` | `openai` \| `anthropic` \| `google` \| `bedrock` \| `bedrock_invoke_model` \| `azure` | yes | AI provider whose response format to parse. | -| `max_body_bytes` | integer | no | Maximum bytes to buffer for a non-streaming JSON response before giving up on locating its usage field. Must be greater than 0. | -| `max_scratch_bytes` | integer | no | Maximum scratch bytes (buffered line + in-progress event data) for the SSE scanner before an event is discarded as oversized. Must be greater than 0. | +| `max_body_bytes` | integer | no | Maximum bytes to buffer for a non-streaming JSON response before giving up on locating its usage field. Must be greater than 0 and at most 64 MiB. | +| `max_scratch_bytes` | integer | no | Maximum scratch bytes (buffered line + in-progress event data) for the SSE scanner before an event is discarded as oversized. Must be greater than 0 and at most 64 MiB. | ## Example diff --git a/filters/src/token_usage/count.rs b/filters/src/token_usage/count.rs index cdc7367f5a..d778a10796 100644 --- a/filters/src/token_usage/count.rs +++ b/filters/src/token_usage/count.rs @@ -30,7 +30,7 @@ use std::fmt::Write as _; use async_trait::async_trait; use bytes::Bytes; use praxis_filter::{ - BodyAccess, BodyMode, FilterAction, FilterError, HttpFilter, HttpFilterContext, + BodyAccess, BodyMode, FilterAction, FilterError, HttpFilter, HttpFilterContext, body::MAX_JSON_BODY_BYTES, builtins::http::payload_processing::config_validation::validate_max_body_bytes, parse_filter_config, }; use serde::Deserialize; @@ -112,13 +112,14 @@ struct TokenCountConfig { provider: ProviderKind, /// Maximum bytes to buffer for a non-streaming JSON response before - /// giving up on locating its usage field. Must be greater than 0. + /// giving up on locating its usage field. Must be greater than 0 and + /// at most 64 MiB. #[serde(default = "default_max_body_bytes")] max_body_bytes: usize, /// Maximum scratch bytes (buffered line + in-progress event data) for /// the SSE scanner before an event is discarded as oversized. Must be - /// greater than 0. + /// greater than 0 and at most 64 MiB. #[serde(default = "default_max_scratch_bytes")] max_scratch_bytes: usize, } @@ -215,9 +216,7 @@ impl TokenCountFilter { pub fn from_config(config: &serde_yaml::Value) -> Result, FilterError> { let cfg: TokenCountConfig = parse_filter_config("token_count", config)?; validate_max_body_bytes("token_count", cfg.max_body_bytes)?; - if cfg.max_scratch_bytes == 0 { - return Err("token_count: 'max_scratch_bytes' must be greater than 0".into()); - } + validate_max_scratch_bytes(cfg.max_scratch_bytes)?; Ok(Box::new(Self { provider: cfg.provider, @@ -227,6 +226,18 @@ impl TokenCountFilter { } } +/// Reject a `max_scratch_bytes` that is zero or above the shared 64 MiB +/// capture ceiling used by other payload-processing filters. +fn validate_max_scratch_bytes(value: usize) -> Result<(), FilterError> { + if value == 0 { + return Err("token_count: 'max_scratch_bytes' must be greater than 0".into()); + } + if value > MAX_JSON_BODY_BYTES { + return Err(format!("token_count: max_scratch_bytes ({value}) exceeds maximum ({MAX_JSON_BODY_BYTES})").into()); + } + Ok(()) +} + #[async_trait] impl HttpFilter for TokenCountFilter { fn name(&self) -> &'static str { diff --git a/filters/src/token_usage/count/tests.rs b/filters/src/token_usage/count/tests.rs index 5dcdec835a..080d77a3f7 100644 --- a/filters/src/token_usage/count/tests.rs +++ b/filters/src/token_usage/count/tests.rs @@ -797,6 +797,30 @@ fn from_config_rejects_zero_max_scratch_bytes() { } } +#[test] +fn from_config_rejects_max_body_bytes_above_ceiling() { + let config: serde_yaml::Value = serde_yaml::from_str("provider: openai\nmax_body_bytes: 67108865").unwrap(); + match TokenCountFilter::from_config(&config) { + Err(err) => assert!( + err.to_string().contains("exceeds maximum"), + "should reject max_body_bytes above 64 MiB: {err}" + ), + Ok(_) => panic!("should reject max_body_bytes above 64 MiB"), + } +} + +#[test] +fn from_config_rejects_max_scratch_bytes_above_ceiling() { + let config: serde_yaml::Value = serde_yaml::from_str("provider: openai\nmax_scratch_bytes: 67108865").unwrap(); + match TokenCountFilter::from_config(&config) { + Err(err) => assert!( + err.to_string().contains("max_scratch_bytes") && err.to_string().contains("exceeds maximum"), + "should reject max_scratch_bytes above 64 MiB: {err}" + ), + Ok(_) => panic!("should reject max_scratch_bytes above 64 MiB"), + } +} + // ----------------------------------------------------------------------------- // Content-Type Helpers // -----------------------------------------------------------------------------