fix(token_usage): recover from SSE overflow instead of silently dropping usage - #782
Conversation
|
Missing Signed-off-by: 10f7d98. All commits require sign-off (via |
10f7d98 to
aa19ca7
Compare
…ing 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 praxis-proxy#674 Signed-off-by: mkoushni <mkoushni@redhat.com>
aa19ca7 to
fe8437c
Compare
praxis-bot
left a comment
There was a problem hiding this comment.
PR Review
Solid fix. The SSE scanner recovery via SkipPhase is well-designed — it bounds memory without losing later events, which is exactly what #674 needs. The explicit overflow status signal for billing consumers is the right call, and the configurable limits are a natural extension. Test coverage is thorough: the regression tests match the issue's reproduction steps, and the cross-chunk boundary test catches a real edge case.
Two medium findings below.
Findings
[Medium] Duplicate SkipPhase encode/decode logic — the A2A load_sse_scan_state/save_sse_scan_state and the token_usage count.rs equivalents both hand-roll the same match for serializing and deserializing SkipPhase to/from metadata strings. If the encoding ever changes (new variant, renamed string), both sites must be updated in lockstep. Consider adding fn as_str(&self) -> &'static str and fn from_metadata_str(s: Option<&str>) -> Self directly on SkipPhase in sse.rs so both callers share one source of truth.
[Medium] max_body_bytes and max_scratch_bytes accept any usize without a minimum bound — a value of 0 would make every response immediately overflow, rendering the filter non-functional. The project conventions call for #[serde(try_from)] on constrained numerics; a newtype enforcing a sensible minimum (e.g. 1 KiB) would prevent misconfiguration and keep the pattern consistent with other filters.
Bring the token-accounting overflow branch current with main so Praxis 0.5.3 API adaptations and other CI fixes are included.
…mits 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 <mkoushni@redhat.com>
f68a60c to
836b9c7
Compare
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 <mkoushni@redhat.com>
Signed-off-by: mkoushni <mkoushni@redhat.com>
Merging main into the prompt-cache branch recombined the overflow-status guard (praxis-proxy#782) with the cache breakdown recording, pushing handle_json_body one line over the too_many_lines threshold. Extract the end-of-stream JSON usage recording into a dedicated helper to keep both concerns and stay under the limit. Signed-off-by: Noy Itzikowitz <nitzikow@redhat.com>
Summary
Fixes #674. Token accounting stopped permanently after a single oversized SSE event or JSON response, silently clearing all working state — including any terminal usage event still to come — with no signal that data was lost.
filters/src/agentic/a2a/sse.rs, used by bothtoken_usageand the A2A task-route capture filter) now discards only the oversized event and resumes scanning at the next event boundary, instead of aborting the whole stream. A terminal usage event arriving after an oversized one is now captured correctly.token.status = overflowmetadata key andPraxis-Token-Statusresponse header are now set explicitly, so billing/monitoring consumers cannot mistake missing counts for a genuine zero-usage response.max_body_bytes/max_scratch_bytesare now configurable pertoken_countfilter instance instead of fixed at 1 MiB / 64 KiB.Changes
filters/src/agentic/a2a/sse.rs: shared scanner recovers at the next blank-line event boundary instead of returning an abort signal; addsSkipPhasestate tracking (discarded bytes are never buffered).filters/src/agentic/a2a/mod.rs: task-route SSE capture adapted to the new scanner contract — capture no longer disables itself after one oversized event.filters/src/token_usage/count.rs:handle_sse_bodyno longer aborts on overflow;handle_json_bodyandfinalize_streaming_countsemit an explicit overflow status instead of silence; capture limits are now filter config fields.filters/src/token_usage/headers.rs: injectsPraxis-Token-Statuswhen set.Test plan
cargo test --workspace— all tests pass (1080 filters/apis lib tests, plus proxy/server/main suites)cargo clippy --workspace --all-targets— cleancargo fmt --check— clean for changed files