Skip to content

fix(token_usage): recover from SSE overflow instead of silently dropping usage - #782

Merged
aslakknutsen merged 5 commits into
praxis-proxy:mainfrom
mkoushni:fix/674-token-accounting-overflow
Aug 25, 2026
Merged

fix(token_usage): recover from SSE overflow instead of silently dropping usage#782
aslakknutsen merged 5 commits into
praxis-proxy:mainfrom
mkoushni:fix/674-token-accounting-overflow

Conversation

@mkoushni

Copy link
Copy Markdown
Contributor

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.

  • The shared SSE scanner (filters/src/agentic/a2a/sse.rs, used by both token_usage and 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.
  • For JSON responses, and for the residual SSE case where the usage event itself is the one that overflows, a token.status = overflow metadata key and Praxis-Token-Status response header are now set explicitly, so billing/monitoring consumers cannot mistake missing counts for a genuine zero-usage response.
  • max_body_bytes / max_scratch_bytes are now configurable per token_count filter 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; adds SkipPhase state 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_body no longer aborts on overflow; handle_json_body and finalize_streaming_counts emit an explicit overflow status instead of silence; capture limits are now filter config fields.
  • filters/src/token_usage/headers.rs: injects Praxis-Token-Status when set.
  • Regression tests added matching the issue's reproduction and suggested tests (oversized non-usage SSE event followed by a valid terminal usage event; oversized JSON with usage near the end).

Test plan

  • cargo test --workspace — all tests pass (1080 filters/apis lib tests, plus proxy/server/main suites)
  • cargo clippy --workspace --all-targets — clean
  • cargo fmt --check — clean for changed files
  • New regression tests cover: oversized SSE event followed by terminal usage event; oversized JSON with usage near the end; explicit overflow status distinguishable from zero; configurable limits

@mkoushni
mkoushni requested review from a team and aslakknutsen August 19, 2026 12:08
@praxis-bot-app

Copy link
Copy Markdown

Missing Signed-off-by: 10f7d98. All commits require sign-off (via git commit --signoff).

@mkoushni
mkoushni force-pushed the fix/674-token-accounting-overflow branch from 10f7d98 to aa19ca7 Compare August 19, 2026 12:10
…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>
@mkoushni
mkoushni force-pushed the fix/674-token-accounting-overflow branch from aa19ca7 to fe8437c Compare August 19, 2026 12:16

@praxis-bot praxis-bot left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread filters/src/agentic/a2a/sse.rs
Comment thread filters/src/token_usage/count.rs
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>
@mkoushni
mkoushni force-pushed the fix/674-token-accounting-overflow branch from f68a60c to 836b9c7 Compare August 23, 2026 12:44
Comment thread filters/src/token_usage/count.rs Outdated
Comment thread filters/src/token_usage/count.rs Outdated
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>
@aslakknutsen
aslakknutsen added this pull request to the merge queue Aug 24, 2026
Merged via the queue into praxis-proxy:main with commit db2881c Aug 25, 2026
20 checks passed
noyitz added a commit to noyitz/ai that referenced this pull request Aug 26, 2026
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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Token accounting silently stops after fixed JSON or SSE capture limits are exceeded

3 participants